Matrix Multiplication Kernels

Implements various matrix multiplication strategies for GPU execution.

Matrix Multiply Variants

1. Naive MatMul

• Simple triple-loop implementation
• Each thread computes one output element
• Good for small matrices, easy to understand

2. Tiled MatMul

• Block-based computation using workgroup shared memory
• Reduces global memory access by ~N/tile_size
• Optimal tile size: 16x16 or 32x32 (depends on GPU)

3. Attention-Specific MatMul

• Q @ K^T: [batch, heads, seq, d] @ [batch, heads, d, seq] → [batch, heads, seq, seq]
• attn @ V: [batch, heads, seq, seq] @ [batch, heads, seq, d] → [batch, heads, seq, d]
• Often includes scaling factor (1/sqrt(d_k))

Performance Considerations

Memory bandwidth (typically the bottleneck):

Naive: Each output element reads full row + column
  Operations: M×N×K
  Memory reads: M×N×K (A) + M×N×K (B) = 2×M×N×K
  Arithmetic intensity: 0.5 FLOP/byte (very low!)

Tiled: Reuse data in shared memory
  Memory reads: M×N×K / tile_size
  Arithmetic intensity: tile_size × 0.5 FLOP/byte (much better!)

For attention (seq_len=2048, d=80):

Q @ K^T: [2048, 80] @ [80, 2048]
  Naive: 2048² × 80 × 2 = 671 MB reads
  Tiled (16×16): 671 / 16 = 42 MB reads (16x reduction!)

References

• CUDA GEMM optimization guide: https://docs.nvidia.com/cuda/cublas/
• Attention implementation: llama.cpp/ggml-cuda/attention.cu

Configuration

structure Hesper.WGSL.MatMul.Config : Type

Matrix multiplication configuration

• M : Nat
• N : Nat
• K : Nat

instance Hesper.WGSL.MatMul.instReprConfig : Repr Config

Equations

• Hesper.WGSL.MatMul.instReprConfig = { reprPrec := Hesper.WGSL.MatMul.instReprConfig.repr }

def Hesper.WGSL.MatMul.instReprConfig.repr : Config → Nat → Std.Format

Equations

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

Naive Matrix Multiply

def Hesper.WGSL.MatMul.naiveMatMulKernel (config : Config) : Monad.ShaderM Unit

Naive matrix multiply: C = A @ B

A: [M, K] B: [K, N] C: [M, N]

Each thread computes one output element C[i,j]: C[i,j] = Σₖ A[i,k] × B[k,j]

@param config Matrix dimensions

Equations

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

Transposed Matrix Multiply

def Hesper.WGSL.MatMul.matMulTransposeBKernel (config : Config) : Monad.ShaderM Unit

Matrix multiply with B transposed: C = A @ B^T

A: [M, K] B: [N, K] (note: same K dimension, not transposed in memory) C: [M, N]

This is more efficient when B is already stored in row-major format. Common in attention: Q @ K^T where K is [seq, d]

@param config Matrix dimensions (K is shared dimension)

Equations

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

F16 Transposed Matrix Multiply (for LM Head)

def Hesper.WGSL.MatMul.matMulTransposeF16Kernel (config : Config) : Monad.ShaderM Unit

Matrix multiply with B transposed, B stored as packed F16: C = A @ B^T

A: [M, K] in F32 B: [N, K] stored as packed F16 (each u32 = 2 F16 values via pack2x16float) C: [M, N] in F32

Uses hardware unpack2x16float to convert F16→F32 during computation. Processes 2 K-elements per loop iteration for 2x bandwidth reduction on B.

Primary use: LM head projection where B is the F16 embedding table (656 MB vs 1.3 GB F32).

@param config Matrix dimensions (K must be even)

Equations

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

Block-Cooperative + SW-Pipelined F16 Transposed MatMul (M=1 decode)

def Hesper.WGSL.MatMul.matMulTransposeF16BlockCoopKernel (config : Config) (gridXWidth : Nat := 0) : Monad.ShaderM Unit

Block-cooperative F16 matrix-vector multiply with B transposed and stored as packed F16: C = A @ B^T, specialised to M = 1 (single-row A). Mirrors the Q4_K / Q6_K block-coop pattern:

• 1 workgroup per output element (N workgroups), 32 threads each
• The K dim is tiled into kTilesPerRow = K / 64 tiles of 64 F16s = 32 u32s. Each lane tid ∈ [0, 32) reads one u32 per tile (contiguous across lanes → one coalesced 128-byte load), unpacks it into two F16→F32 values, pairs them with the matching two A f32s, and accumulates two FMAs.
• Depth-1 software pipelining: each iteration, the u32 for tile N+1 is prefetched into a mutable nextU32 var BEFORE the dequant+FMA chain that consumes the current snapshot, so memory load latency overlaps with FMA latency (same trick that cut Q4_K ffnGate/Up to 0.059 ms at 335 GB/s).

Constraints: M = 1 (the inner loop keeps the single A row in registers), K % 64 == 0 (per-lane tile of 64 F16 values needs 32 u32s cleanly), subgroup support required.

Weight layout: row-major [N, K/2] packed F16 (same as matMulTransposeF16Kernel). Row n starts at u32 index n * (K/2).

Equations

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

def Hesper.WGSL.MatMul.executeMatMulTransposeF16BlockCoop {β : Type} [GPUBackend β] (ctx : β) (aBuf bF16Buf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute matMulTransposeF16BlockCoopKernel. See the kernel doc for preconditions (M = 1, K % 64 == 0, subgroup support).

Equations

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

Optimized F16 Transposed MatMul with Shared Memory

def Hesper.WGSL.MatMul.matMulTransposeF16SharedKernel (config : Config) : Monad.ShaderM Unit

Optimized matrix multiply with B transposed, B stored as packed F16: C = A @ B^T

Key optimizations over matMulTransposeF16Kernel:

1. A vector loaded into workgroup shared memory (eliminates redundant global reads)
2. Inner loop processes 4 u32 (= 8 F16 values) per iteration (4x fewer loop iterations)

For M=1 LM head (1×2560 @ 128256×2560):

• Old: each of 128K threads reads all 2560 A values from global memory = 1.3 GB total A reads
• New: each workgroup (256 threads) loads A once into shared memory = 5 MB total A reads

Requirements: K must be divisible by 8

@param config Matrix dimensions (K must be divisible by 8)

Equations

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

Cooperative-Matrix (WMMA) F16 Transposed MatMul

def Hesper.WGSL.MatMul.matMulTransposeF16WMMAKernel (config : Config) : Monad.ShaderM Unit

C = A @ B^T using NVIDIA / Chromium subgroup matrix operations.

• A: [M, K] f32 — cast to f16 on load into workgroup memory
• B: [N, K] f16 (packed as u32, two halves per u32) — the same layout used by matMulTransposeF16Kernel / -SharedKernel. The kernel reads B in its native f16 form and never materialises a f32 copy.
• C: [M, N] f32 — cooperative matrix accumulator, written back unscaled.

The kernel targets the NVIDIA Ada/Ampere WMMA config (in=f16, out=f32, M=N=K=16), which is exposed by Dawn when both ShaderF16 and ChromiumExperimentalSubgroupMatrix are enabled.

Workgroup layout: 1 subgroup (32 threads) per 16×16 output tile. Dispatch (N/16, M/16, 1) workgroups.

Constraints (enforced by the caller):

• config.M % 16 == 0
• config.N % 16 == 0
• config.K % 16 == 0

Note: this kernel is generic (unlike fusedBitLinearSubgroupMatrixKernel it does no ternary dequant) and is intended to back the Gemma 4 Q4_K/Q6_K linears and the LM head once those paths are wired up. The pre-dequant-to-f16 variant is what makes it reusable.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMA8x8Kernel (config : Config) : Monad.ShaderM Unit

8×8×8 variant of matMulTransposeF16WMMAKernel for hardware that only supports 8×8 subgroup matrices (Apple M-series / Metal: probe reports M=N=K=8 only; the 16×16 kernel crashes at pipeline creation). A: f32 [M,K]; B: u32-packed f16 [N, K/2] (i.e. B^T); C: f32 [M,N]. One subgroup (32 lanes) per 8×8 output tile. Requires M,N,K multiples of 8.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMA8x8RegKernel (config : Config) (TM TN : Nat) : Monad.ShaderM Unit

Register-blocked 8×8 WMMA: one workgroup computes a TM×TN grid of 8×8 output tiles (= 8·TM × 8·TN output), keeping TM·TN f32 result matrices in registers across the whole K-loop and reusing each loaded A/B 8×8 tile TN/TM ways. Amortizes the shared-memory loads and barriers TM·TN×. A: f32 [M,K]; B: u32-packed f16 [N,K/2]; C: f32 [M,N]. Requires M % (8·TM) = 0, N % (8·TN) = 0, K % 8 = 0. Grid: (N/(8·TN), M/(8·TM)) × 32.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMA8x8SgKernel (config : Config) (TM TN : Nat) : Monad.ShaderM Unit

SUBGROUP-blocked 8×8 WMMA: one workgroup runs TM·TN subgroups (32 lanes each), each computing ONE 8×8 output tile (1 result matrix → no register spill), but the workgroup loads TM A-tiles + TN B-tiles into shared memory ONCE per K-step and all subgroups reuse them. Cuts the A-row/B-col re-read (TN/TM×) and the total workgroup/barrier count (TM·TN×) without the matrix-register pressure that sinks register-blocking on Apple GPUs. workgroupSize = TM·TN·32 (S subgroups). Grid: (N/(8·TN), M/(8·TM)) × (TM·TN·32).

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMA8x8KbKernel (config : Config) (BK : Nat) : Monad.ShaderM Unit

K-batched 8×8 WMMA: like the naive 8×8 kernel (one 32-lane subgroup, ONE 8×8 output tile, ONE result matrix — no register/subgroup pressure) but loads BK K-tiles into shared per barrier and does BK back-to-back MMAs, cutting the barrier count BK×. A: f32 [M,K]; B: u32 f16 [N,K/2]; C: f32 [M,N]. Requires K % (8·BK) = 0. Grid: (N/8, M/8) × 32.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMARegKernel (config : Config) (weightRowOffset weightRows : Nat := 0) (grouped : Bool := false) (nExpert : Nat := 0) : Monad.ShaderM Unit

llama.cpp-style register-blocked 8×8 WMMA matmul. Mirrors ggml-metal kernel_mul_mm: threadgroup tile BM=64 (M) × BN=32 (N), BK=32 (K); 4 simdgroups (128 threads) in a 2×2 grid, each computing a 32×16 sub-tile = 4×2 = 8 result 8×8 matrices kept in registers (14 matrix registers total: 8 result + 4 left + 2 right), reusing each loaded 8×8 tile 2-4 ways. A: f32 [M,K]; B: u32-packed f16 [N,K/2] (B^T); C: f32 [M,N]. Requires M%64=N%32=K%32=0. Grid: (N/32, M/64) × 128.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMARegKernelGen (config : Config) (TMsg TNsg sgRows sgCols BK : Nat) (weightRowOffset : Nat := 0) (directB : Bool := false) : Monad.ShaderM Unit

DEVPLAN M2/M4 — GENERALIZED register-blocked WMMA matmul: the deployed 64×32 kernel (matMulTransposeF16WMMARegKernel) with its hardcoded tile/simdgroup constants replaced by parameters, as the autotune sweep substrate (runtime K loop — full-unroll generators explode Tint/Metal compile at K=2816, see DEVPLAN).

Parameters (Triton correspondence): TMsg, TNsg — 8×8 fragments per simdgroup in M / N (register-pressure knob; BLOCK_M/N) sgRows, sgCols — simdgroup grid inside the workgroup (num_warps = sgRows·sgCols) BK — K-tile depth staged in shared per iteration (BLOCK_K), BK % 8 = 0

Derived: wgSize = sgRows·sgCols·32 threads; workgroup tile = (sgRows·8·TMsg) × (sgCols·8·TNsg); shared = (Mtile·BK + BK·Ntile) f16. The deployed kernel = (TMsg=4, TNsg=2, sgRows=2, sgCols=2, BK=32) — the sweep must reproduce or beat it (M4 acceptance).

Caller contract (same as deployed): M padded to an Mtile multiple (WMMA tail rows write past M otherwise — the known heap-stomp class), N % Ntile = 0, K % BK = 0, (Mtile·BK) % wgSize = 0 and (Ntile·BK/2) % wgSize = 0 (exact cooperative-load iterations — filter variants violating these BEFORE generation). Non-grouped only. A: f32 [M,K]; B: u32-packed f16 [N,K/2] (Bᵀ); C: f32 [M,N]. Grid: (N/Ntile, M/Mtile) × wgSize.

Equations

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

def Hesper.WGSL.MatMul.matMulTransposeF16WMMARegSmallMKernel (config : Config) (weightRowOffset weightRows : Nat := 0) : Monad.ShaderM Unit

Small-M variant of the register-blocked WMMA matmul, tuned for MoE expert matmuls where M (tokens routed to one expert) is small (~16). Tile BM=16 (M) × BN=64 (N), BK=32. 4 simdgroups split N (each does 16×16 = 2×2 = 4 result 8×8 matrices, sharing the same 16 A-rows), so a 64-row tile's ~4× wasted compute at M≈16 is avoided. 8 matrix registers (4 result + 2 left + 2 right). A: f32 [M,K]; B: u32 f16 [N,K/2] (B^T); C: f32 [M,N]. Requires K%32=0; M,N rounded up by the grid. Grid: (N/64, M/16) × 128. weightRowOffset/weightRows select a column-block of a wider weight.

Equations

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

Scaled Matrix Multiply (for Attention)

def Hesper.WGSL.MatMul.scaledMatMulTransposeKernel (config : Config) (scale : Float) : Monad.ShaderM Unit

Scaled matrix multiply: C = (A @ B) / scale

Used in attention: scores = (Q @ K^T) / sqrt(d_k)

@param config Matrix dimensions @param scale Scaling factor (e.g., 1/sqrt(d_k))

Equations

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

Batched Matrix Multiply

def Hesper.WGSL.MatMul.batchedMatMulKernel (config : Config) (batchSize : Nat) : Monad.ShaderM Unit

Batched matrix multiply: C[b] = A[b] @ B[b] for each batch b

A: [batch, M, K] B: [batch, K, N] C: [batch, M, N]

Common in multi-head attention where batch includes both actual batch dimension and number of heads.

@param config Matrix dimensions (per batch) @param batchSize Number of batches

Equations

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

def Hesper.WGSL.MatMul.batchedScaledMatMulTransposeKernel (config : Config) (batchSize : Nat) (scale : Float) : Monad.ShaderM Unit

Batched scaled transposed matmul: C[b] = (A[b] @ B[b]^T) * scale

A: [batch, M, K] B: [batch, N, K] (transposed: B is stored row-major as [N, K]) C: [batch, M, N]

Used for attention scores: scores = (Q @ K^T) / sqrt(d_k)

Equations

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

High-Level API

def Hesper.WGSL.MatMul.executeMatMul {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute naive matrix multiply: C = A @ B

@param device WebGPU device @param aBuf Matrix A [M, K] @param bBuf Matrix B [K, N] @param cBuf Output matrix C [M, N] @param config Matrix dimensions

Equations

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

def Hesper.WGSL.MatMul.executeMatMulTranspose {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute transposed matrix multiply: C = A @ B^T

@param device WebGPU device @param aBuf Matrix A [M, K] @param bBuf Matrix B [N, K] (will be transposed) @param cBuf Output matrix C [M, N] @param config Matrix dimensions

Equations

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

def Hesper.WGSL.MatMul.executeMatMulTransposeF16 {β : Type} [GPUBackend β] (ctx : β) (aBuf bF16Buf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute transposed matrix multiply with B in packed F16: C = A @ B^T

B is stored as packed F16 (each u32 = 2 F16 values). Uses hardware unpack2x16float for 2x bandwidth reduction. Primary use: LM head with F16 embedding table.

@param device WebGPU device @param aBuf Matrix A [M, K] in F32 @param bF16Buf Matrix B [N, K] stored as packed F16 ([N, K/2] u32 values) @param cBuf Output matrix C [M, N] in F32 @param config Matrix dimensions (K must be even)

Equations

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

def Hesper.WGSL.MatMul.executeMatMulTransposeF16Shared {β : Type} [GPUBackend β] (ctx : β) (aBuf bF16Buf cBuf : GPUBackend.Buf β) (config : Config) : IO Unit

Execute optimized F16 transposed matmul with shared memory: C = A @ B^T

Uses shared memory for A vector to eliminate redundant global reads. Significantly faster than executeMatMulTransposeF16 for small M (especially M=1).

Requirements: K must be divisible by 8

@param device WebGPU device @param aBuf Matrix A [M, K] in F32 @param bF16Buf Matrix B [N, K] stored as packed F16 ([N, K/2] u32 values) @param cBuf Output matrix C [M, N] in F32 @param config Matrix dimensions (K must be divisible by 8)

Equations

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

def Hesper.WGSL.MatMul.executeScaledMatMulTranspose {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) (scale : Float) : IO Unit

Execute scaled transposed matrix multiply: C = (A @ B^T) / scale

Used for attention scores: scores = (Q @ K^T) / sqrt(d_k)

@param device WebGPU device @param aBuf Matrix A [M, K] @param bBuf Matrix B [N, K] @param cBuf Output matrix C [M, N] @param config Matrix dimensions @param scale Scaling factor

Equations

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

def Hesper.WGSL.MatMul.executeBatchedMatMul {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) (batchSize : Nat) : IO Unit

Execute batched matrix multiply: C[b] = A[b] @ B[b]

@param device WebGPU device @param aBuf Matrix A [batch, M, K] @param bBuf Matrix B [batch, K, N] @param cBuf Output matrix C [batch, M, N] @param config Matrix dimensions (per batch) @param batchSize Number of batches

Equations

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

def Hesper.WGSL.MatMul.executeBatchedScaledMatMulTranspose {β : Type} [GPUBackend β] (ctx : β) (aBuf bBuf cBuf : GPUBackend.Buf β) (config : Config) (batchSize : Nat) (scale : Float) : IO Unit

Execute batched scaled transposed matmul: C[b] = (A[b] @ B[b]^T) * scale

A: [batch, M, K], B: [batch, N, K] → C: [batch, M, N] Used for attention: scores = (Q @ K^T) / sqrt(d_k)

Equations

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