BitLinear Backward (Transpose MatVec)

Computes dInput = scale * W^T @ dOutput for the O projection backward.

i2_s Element Indexing (from forward kernel)

Given a row's u32 array, u32 index u32Idx decodes to 16 elements:

group128 = u32Idx / 8
groupPos = (u32Idx % 8) * 4

For byte b in [0..3], shift s in [0..3]:
  elemIdx = group128 * 128 + groupPos + b + s * 32
  code = ((packed >> (b*8)) >> (6 - s*2)) & 3
  weight = code - 1

For transpose (column j access across rows):

group128 = j / 128
posInGroup = j % 128
s = posInGroup / 32        (shift group)
subPos = posInGroup % 32   (within 32-element sub-group)
b = subPos % 4             (byte index within u32)
u32InGroup = subPos / 4    (u32 within group)
u32Idx = group128 * 8 + u32InGroup

byte_shift = b * 8
code_shift = 6 - s * 2

def Hesper.Training.BitLinearBackward.bitLinearTransposeKernel (inDim outDim : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

Transpose matmul kernel: dInput[j] = scale * Σ_i W[i,j] * dOutput[i]

W is [outDim, inDim] in i2_s format. One workgroup per input element j, with threads cooperating over outDim. Uses shared memory reduction.

Equations

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

def Hesper.Training.BitLinearBackward.executeBitLinearTranspose (device : WebGPU.Device) (layer : Layers.BitLinear.BitLinear WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (dOutputBuf dInputBuf : WebGPU.Buffer) : IO Unit

Execute BitLinear transpose: dInput = scale * W^T @ dOutput

Equations

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