Attention Backward GPU Kernels

GPU kernels implementing the backward pass through attention for LoRA training. Each kernel corresponds to a verified CPU spec in VerifiedBackward.lean.

Gradient Flow (reverse order of forward)

dOutput [dim]
  ↓ O projection backward (BitLinear transpose)
dAttnOut [dim]
  ↓ RMSNorm backward (sub-norm)
dAttnWeighted [numHeads * headDim]
  ↓ Attention apply backward
dAttn [numHeads * cacheLen] + dV [kvDim] (not needed for LoRA Q)
  ↓ Softmax backward
dScores [numHeads * cacheLen]
  ↓ Score backward (Q @ K^T)
dQ [numHeads * headDim]
  ↓ RoPE backward (inverse rotation)
dQpre [numHeads * headDim]  ← This is ∂L/∂(BitLinear_Q output) = LoRA Q gradient signal

Softmax Backward

def Hesper.Training.AttentionBackward.softmaxBackwardKernel (numHeads cacheLen : Nat) : WGSL.Monad.ShaderM Unit

Softmax backward kernel: dScores[h, s] = attn[h, s] * (dAttn[h, s] - Σ_s' attn[h, s'] * dAttn[h, s'])

One thread per (head, seq_pos) pair. Uses shared memory for the dot product reduction per head.

Equations

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

def Hesper.Training.AttentionBackward.executeSoftmaxBackward (device : WebGPU.Device) (attnBuf dAttnBuf dScoresBuf : WebGPU.Buffer) (numHeads cacheLen : Nat) : IO Unit

Equations

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

Attention Score Backward (dQ from dScores)

def Hesper.Training.AttentionBackward.scoreBackwardQKernel (numHeads numKVHeads cacheLen headDim : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Score backward kernel for Q: dQ[h, d] = scale * Σ_s dScores[h, s] * K_cache[kvHead(h), s, d]

One thread per (head, dim) pair. GQA: multiple heads map to the same KV head.

Equations

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

def Hesper.Training.AttentionBackward.executeScoreBackwardQ (device : WebGPU.Device) (dScoresBuf kCacheBuf dQBuf : WebGPU.Buffer) (numHeads numKVHeads cacheLen headDim : Nat) (scale : Float) : IO Unit

Equations

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

Attention Apply Backward (dAttn from dOutput @ V^T)

def Hesper.Training.AttentionBackward.applyBackwardKernel (numHeads numKVHeads cacheLen headDim : Nat) : WGSL.Monad.ShaderM Unit

Attention apply backward kernel: dAttn[h, s] = Σ_d dOutput[h, d] * V_cache[kvHead(h), s, d]

One thread per (head, seq_pos) pair.

Equations

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

def Hesper.Training.AttentionBackward.executeApplyBackward (device : WebGPU.Device) (dOutputBuf vCacheBuf dAttnBuf : WebGPU.Buffer) (numHeads numKVHeads cacheLen headDim : Nat) : IO Unit

Equations

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

RoPE Backward (inverse rotation)

def Hesper.Training.AttentionBackward.ropeBackwardKernel (numHeads headDim : Nat) (ropeBase : Float) (pos : Nat) : WGSL.Monad.ShaderM Unit

RoPE backward kernel: apply inverse rotation R(-θ) to gradient. For NeoX split-half layout: dx[h, d] = dy[h, d] * cos(θ) + dy[h, d+half] * sin(θ) dx[h, d+half] = -dy[h, d] * sin(θ) + dy[h, d+half] * cos(θ)

where θ = pos * base^(-2d/headDim), same as forward.

Equations

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

def Hesper.Training.AttentionBackward.executeRopeBackward (device : WebGPU.Device) (dOutBuf dInBuf : WebGPU.Buffer) (numHeads headDim : Nat) (ropeBase : Float) (pos : Nat) : IO Unit

Equations

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

RMSNorm Backward

def Hesper.Training.AttentionBackward.rmsNormBackwardKernel (dim : Nat) (eps : Float) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

RMSNorm backward kernel (single workgroup, shared memory reduction): dx[i] = (1/rms) * (dy[i]γ[i] - x[i] * dot(dyγ, x) / (n * rms²))

Uses the same workgroup reduction pattern as forward RMSNorm.

Equations

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

def Hesper.Training.AttentionBackward.executeRmsNormBackward (device : WebGPU.Device) (xBuf gammaBuf dOutBuf dInBuf : WebGPU.Buffer) (dim : Nat) (eps : Float := 1e-6) : IO Unit

Equations

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