LoRA Forward Pass GPU Kernels

Implements the LoRA forward computation on GPU:

output = BitLinear(x) + (alpha / rank) * B @ (A @ x)

Decomposed into three GPU operations:

1. loraProjectA: h = A @ x ([rank] = [rank, inDim] @ [inDim])
2. loraProjectB: y = B @ h ([outDim] = [outDim, rank] @ [rank])
3. loraFusedAdd: output[i] += scale * y[i]

For single-token training (rank=8, dim=2560), these are very small matmuls.

GPU Kernels

def Hesper.LoRA.Forward.loraProjectAKernel (rank inDim : Nat) : WGSL.Monad.ShaderM Unit

Kernel: h = A @ x A is [rank, inDim] row-major, x is [inDim], h is [rank]. Each thread computes one element of h (one dot product over inDim).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.loraProjectBKernel (outDim rank : Nat) : WGSL.Monad.ShaderM Unit

Kernel: y = B @ h B is [outDim, rank] row-major, h is [rank], y is [outDim]. Each thread computes one element of y.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.loraAddScaledKernel (numElements : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Kernel: output[i] += scale * y[i] Adds the LoRA contribution to the base BitLinear output in-place. output is read-write (already contains base output).

Equations

• One or more equations did not get rendered due to their size.

Execution Functions

def Hesper.LoRA.Forward.executeProjectA (device : WebGPU.Device) (weight : Weight) (xBuf hBuf : WebGPU.Buffer) : IO Unit

Execute LoRA A projection: h = A @ x

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.executeProjectB (device : WebGPU.Device) (weight : Weight) (hBuf yBuf : WebGPU.Buffer) : IO Unit

Execute LoRA B projection: y = B @ h

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.executeAddScaled (device : WebGPU.Device) (yBuf outputBuf : WebGPU.Buffer) (numElements : Nat) (scale : Float) : IO Unit

Execute LoRA add: output += scale * y

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.executeLoRAForward (device : WebGPU.Device) (weight : Weight) (scale : Float) (inputBuf outputBuf hBuf yBuf : WebGPU.Buffer) : IO Unit

Full LoRA forward pass for a single projection. Computes: outputBuf += (alpha/rank) * B @ (A @ inputBuf)

@param device GPU device @param weight LoRA weight pair (A, B) @param scale The alpha/rank scaling factor @param inputBuf Input buffer [inDim] (shared with base BitLinear input) @param outputBuf Output buffer [outDim] (already contains base BitLinear output) @param hBuf Temporary buffer [rank] for intermediate h = A @ x @param yBuf Temporary buffer [outDim] for y = B @ h

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.saveActivation (device : WebGPU.Device) (srcBuf dstBuf : WebGPU.Buffer) (numElements : Nat) : IO Unit

Save input activation for backward pass (copy inputBuf to savedBuf)

Equations

• One or more equations did not get rendered due to their size.

Fused LoRA Kernels

def Hesper.LoRA.Forward.loraFusedBAddKernel (outDim rank : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Fused projectB + addScaled: output[i] += scale * Σ_r B[i,r] * h[r] Combines B@h matmul and scaled add into 1 dispatch (saves 1 dispatch per call).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.executeFusedBAdd (device : WebGPU.Device) (weight : Weight) (scale : Float) (hBuf outputBuf : WebGPU.Buffer) : IO Unit

Execute fused B@h + add: output += scale * B @ h (1 dispatch instead of 2)

Equations

• One or more equations did not get rendered due to their size.

def Hesper.LoRA.Forward.executeLoRAForwardFused (device : WebGPU.Device) (weight : Weight) (scale : Float) (inputBuf outputBuf hBuf : WebGPU.Buffer) : IO Unit

Execute full LoRA forward with fused B@h+add: 2 dispatches instead of 3. projectA (1 dispatch) → fusedBAdd (1 dispatch)

Equations

• One or more equations did not get rendered due to their size.