LoRA Backward Pass GPU Kernels

Given upstream gradient dOutput [outDim], computes:

1. dB = scale * outer(dOutput, h) where h = A @ x (saved from forward)
2. dA = scale * outer(B^T @ dOutput, x) where x is saved from forward
3. dInput += A^T @ (B^T @ dOutput) * scale (gradient to residual stream)

All operations are small due to low rank (4-16).

GPU Kernels

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

Kernel: dB[i, r] += scale * dOutput[i] * h[r] Outer product of dOutput [outDim] and h [rank]. Each thread computes one element of the [outDim, rank] gradient matrix.

Equations

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

def Hesper.LoRA.Backward.gradDhKernel (outDim rank : Nat) : WGSL.Monad.ShaderM Unit

Kernel: dh[r] = sum_i B[i, r] * dOutput[i] Computes B^T @ dOutput. Each thread computes one element of dh [rank].

Equations

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

def Hesper.LoRA.Backward.gradAKernel (rank inDim : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Kernel: dA[r, j] += scale * dh[r] * x[j] Outer product of dh [rank] and x [inDim]. Each thread computes one element of the [rank, inDim] gradient matrix.

Equations

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

def Hesper.LoRA.Backward.inputGradKernel (rank inDim : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Kernel: dInput[j] += scale * sum_r A[r, j] * dh[r] Propagates gradient back through LoRA to the residual stream. Each thread computes one element of dInput [inDim].

Equations

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

Execution Functions

def Hesper.LoRA.Backward.executeGradB (device : WebGPU.Device) (dOutputBuf hBuf dBBuf : WebGPU.Buffer) (outDim rank : Nat) (scale : Float) : IO Unit

Execute gradient computation for B: dB += scale * outer(dOutput, h)

Equations

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

def Hesper.LoRA.Backward.executeGradDh (device : WebGPU.Device) (bBuf dOutputBuf dhBuf : WebGPU.Buffer) (outDim rank : Nat) : IO Unit

Execute B^T @ dOutput to get dh [rank]

Equations

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

def Hesper.LoRA.Backward.executeGradA (device : WebGPU.Device) (dhBuf xBuf dABuf : WebGPU.Buffer) (rank inDim : Nat) (scale : Float) : IO Unit

Execute gradient computation for A: dA += scale * outer(dh, x)

Equations

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

def Hesper.LoRA.Backward.executeInputGrad (device : WebGPU.Device) (aBuf dhBuf dInputBuf : WebGPU.Buffer) (rank inDim : Nat) (scale : Float) : IO Unit

Execute input gradient propagation: dInput += scale * A^T @ dh

Equations

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

def Hesper.LoRA.Backward.executeLoRABackward (device : WebGPU.Device) (weight : Weight) (grad : WeightGrad) (scale : Float) (dOutputBuf savedX savedH dInputBuf dhBuf : WebGPU.Buffer) : IO Unit

Full LoRA backward pass for a single projection. Computes dA, dB gradients and propagates dInput.

@param device GPU device @param weight LoRA weight (A, B matrices) @param grad Gradient buffers to accumulate into @param scale alpha/rank scaling factor @param dOutputBuf Upstream gradient [outDim] @param savedX Saved input from forward pass [inDim] @param savedH Saved intermediate h = A @ x from forward [rank] @param dInputBuf Buffer to accumulate input gradient into [inDim] @param dhBuf Temporary buffer [rank] for dh = B^T @ dOutput

Equations

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