Kernel Fusion Framework

Compose multiple ShaderM operations into a single GPU dispatch.

Key Insight

ShaderM is a monad that generates WGSL code. When two ShaderM computations write to / read from the same buffer, fusing them eliminates the intermediate buffer and reduces dispatch count.

Fusion Types

1. Element-wise chain: op1 writes out[i], op2 reads out[i] → inline
2. Multi-copy: N independent copies → 1 kernel with N read/writes
3. Sequential with shared memory: reduction → element-wise

Multi-Buffer Copy (fused save activations)

def Hesper.WGSL.Fusion.fusedCopy4Kernel (numElements numPairs : Nat) : Monad.ShaderM Unit

Fused copy of up to 4 buffers in a single dispatch. Each (src, dst) pair is copied element-wise. All copies must have the same element count.

Equations

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

def Hesper.WGSL.Fusion.executeFusedCopy (device : WebGPU.Device) (pairs : Array (WebGPU.Buffer × WebGPU.Buffer)) (numElements : Nat) : IO Unit

Execute fused copy of up to 4 buffer pairs of the same size

Equations

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

def Hesper.WGSL.Fusion.fusedSaveAttentionActivations (device : WebGPU.Device) (normedBuf savedNormed attnOutBuf savedAttnOut : WebGPU.Buffer) (dim : Nat) (attnBuf savedAttn : WebGPU.Buffer) (attnSize : Nat) : IO Unit

Fused save of attention activations (normed + attnOut = 2 pairs of dim size)

• attention weights (1 pair of attnSize) in 2 dispatches instead of 3

Equations

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

def Hesper.WGSL.Fusion.fusedSaveFFNActivations (device : WebGPU.Device) (gateBuf savedGate upBuf savedUp hiddenBuf savedHidden : WebGPU.Buffer) (ffnDim : Nat) (residual1Buf savedResidual1 : WebGPU.Buffer) (dim : Nat) : IO Unit

Fused save of FFN activations (gate + up + hidden = 3 pairs of ffnDim)

• residual1 (1 pair of dim) in 2 dispatches instead of 4

Equations

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