Batched GEMM Math

What a batched GEMM actually is

A plain GEMM computes C = α·A B + β·C for A: [M, K], B: [K, N], C: [M, N]. A batched GEMM computes L of these at once:

for i = 0 .. L-1:
    C_i = α · A_i B_i + β · C_i

A_i: [M, K]   B_i: [K, N]   C_i: [M, N]   i = 0..L-1
total FLOPs = L · 2MNK

The products are independent — no C_i feeds any A_j — which is the whole reason this is expressible as one call. That independence lets the library schedule all L · (M/M_t) · (N/N_t) output tiles as a single flat pool of work instead of L separate, under-parallelised problems. The FLOP count is identical to running the loop yourself; everything batched GEMM buys is on the other side of the ledger — bytes moved, threads idled, per-call overhead paid.

Where the batch dimension comes from in a transformer

Multi-head attention is the canonical source. With batch size B, sequence length S, H heads and per-head dimension d_h = d_model / H, each projection is one big GEMM — X: [B·S, d_model] times W_Q: [d_model, d_model] — but the attention core is not. After reshaping to [B, H, S, d_h], every (sequence, head) pair needs its own pair of products:

scores = Q K^T   [S, d_h] × [d_h, S] → [S, S]
out    = P V     [S, S]   × [S, d_h] → [S, d_h]

batch count L = B · H

For B = 8 and H = 12 that is 96 independent GEMMs per matmul — 192 per layer, over 2,000 for a 12-layer model, in one forward pass. Mixture-of-experts adds a second source: each expert’s feed-forward weights form a separate A_i, applied only to the tokens routed to it.

The strided-batched layout

In practice the batch lives in one contiguous tensor and the library walks it with a constant offset per item — the strided batched form. Each operand needs a leading dimension (row-to-row step inside a matrix) and a batch stride (matrix-to-matrix step):

A_i[m, k] = A[ i*strideA + m*ldA + k ]
B_i[k, n] = B[ i*strideB + k*ldB + n ]
C_i[m, n] = C[ i*strideC + m*ldC + n ]

Q as [B, H, S, d_h], row-major:
  ldA      = d_h            (row to row inside one head)
  strideA  = S * d_h        (head to head)

Two consequences fall out. First, the API is a handful of integers rather than an array of L pointers, so there is no pointer-chasing and no host-side setup scaling with L. Second, strideB = 0 is legal: it broadcasts one shared B across the batch — exactly what grouped-query attention needs when several query heads share one key/value head.

Why a single small matmul is memory-bound

Arithmetic intensity is I = FLOPs / bytes. For one [M, K] × [K, N] product in fp32, assuming a cold cache and one pass over each operand:

FLOPs = 2 M N K
bytes = 4 (MK + KN + MN)
I     = 2MNK / (4(MK + KN + MN))

M = N = K = 4096:  I ≈ 2·6.9e10 / (4·5.0e7)  ≈ 683 FLOP/byte
M = N = K = 64:    I ≈ 2·2.6e5  / (4·1.2e4)  ≈ 10.7 FLOP/byte

A machine sustaining 200 GFLOP/s against 20 GB/s of bandwidth has a roofline ridge point near 10 FLOP/byte. The 4096-cube GEMM sits far to the compute-bound right of that ridge and can approach peak. The 64-cube GEMM lands on it — and that is the optimistic estimate, before per-call overhead, cold TLB entries, or waking the thread pool.

The overhead the loop version pays

Run those small GEMMs as an explicit loop of L library calls and you pay a fixed cost t_ovh every iteration: argument validation, shape dispatch, kernel selection, a thread-pool fork and join, and on an accelerator a kernel launch. Model the total as:

T_loop    = L · (t_ovh + 2MNK / F_eff)
T_batched = t_ovh + L · 2MNK / F_eff'      with F_eff' > F_eff

L = 96, MNK-work = 2·64·64·64 ≈ 524 kFLOP
at 50 GFLOP/s effective: compute ≈ 10.5 µs per item
t_ovh ≈ 5 µs  →  loop spends ≈ 32% of wall time on overhead

Batching amortises t_ovh across all L items, so it falls from a third of runtime to a rounding error. The second term matters as much: F_eff improves because the batched kernel keeps its packed panels and threads alive across items instead of rebuilding them L times.

How batching restores arithmetic intensity

Batching does not change FLOPs, and with entirely distinct operands it would not change bytes either — so why does intensity rise? Three effects. Shared operands. Whenever a matrix is reused across the batch (a broadcast B, a shared KV head, one expert’s weights applied to many tokens), it is loaded once and amortised over L products, dividing its byte contribution by L. Cache residency. One kernel can hold packed panels in L2 across consecutive items; a loop of calls evicts and re-packs. Parallel efficiency. A [64, 64] output offers one or two tiles — not enough to fill sixteen cores — whereas L × those tiles fills them comfortably, so throughput per item rises though the arithmetic is unchanged. With a shared B the bytes term becomes 4(L·MK + KN + L·MN), and intensity climbs toward the large-GEMM regime as L grows.

Batching across requests: the decode-time case

Autoregressive decoding is the pathological case. Each step processes exactly one new token, so M = 1 and every projection degenerates from a GEMM into a GEMV:

M = 1:  FLOPs = 2NK,  bytes ≈ 4KN  (the weight matrix dominates)
        I ≈ 2NK / 4KN = 0.5 FLOP/byte   (fp32)
        I ≈ 2 FLOP/byte                 (int8 weights)

That is one to two orders of magnitude below any sane ridge point: decode is pure weight streaming and the ALUs idle. The fix is batching across requests. Serving R concurrent sequences turns the GEMV back into a [R, K] × [K, N] GEMM with the weight bytes read once for all R tokens, so intensity scales roughly linearly in R until the compute roof is hit. That is the whole argument for continuous batching — and equally why a single-user CPU SLM feels bandwidth-starved no matter how many cores it has.

Grouped GEMM: when the shapes differ

Strided batching assumes every item has identical M, N, K and a constant stride. Plenty of transformer workloads violate that. A grouped GEMM takes a list of problem descriptors and schedules them as one fused kernel:

problem[i] = (M_i, N_i, K_i, ptrA_i, ptrB_i, ptrC_i)
total tiles = Σ_i ceil(M_i/M_t) · ceil(N_i/N_t)
each worker pulls the next tile from one global queue

The key idea is that scheduling happens over the flattened tile space, not over problems, so a group of one large and twenty tiny matmuls still balances across cores. Three workloads want this: mixture-of-experts, where expert e receives a different token count M_e; ragged batches of unpadded variable-length sequences; and layers with heterogeneous projection widths. The cost is a descriptor table and slightly more complex tile indexing — cheap next to padding everything to the maximum.

A worked MoE grouped-GEMM example

Take an MoE layer with E = 8 experts, top-1 routing, 1,024 tokens, d_model = 512, d_ff = 2048. Routing is never uniform; suppose the counts land as:

M_e = [310, 205, 160, 120,  95,  70,  40,  24]   Σ = 1024

grouped:  Σ_e 2 · M_e · 512 · 2048 = 2 · 1024 · 512 · 2048 ≈ 2.15 GFLOP
padded to M_max = 310:  8 · 2 · 310 · 512 · 2048 ≈ 5.20 GFLOP
waste factor = 5.20 / 2.15 ≈ 2.4×

Padding to a uniform batch burns 2.4× the arithmetic on tokens that do not exist. Note too that the grouped total is invariant to the routing distribution — it depends only on the 1,024 tokens — while the padded cost is set entirely by the most loaded expert. The more skewed the router, the worse padding gets, which is why load-balancing losses and grouped kernels tend to appear together.