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.

What this means on a CPU running a small model

On a CPU-hosted SLM the batch dimension is your main lever, and you get less of it than you want. Prefill is comfortable: hundreds of tokens make every projection a healthy GEMM. Decode is not, because a single interactive user gives R = 1. Three moves follow. Quantise the weights — int8 or int4 cuts the dominant byte term by 4× or 8×, which for a bandwidth-bound GEMV translates almost directly into speedup. Batch whatever you can: fuse W_Q, W_K, W_V into one [d, 3d] matrix so three GEMVs become one, keep speculative draft tokens together so M is 4–8 rather than 1, and serve concurrent requests through one continuous batch. Keep the attention batch contiguous so a strided-batched call sees clean constant strides, not a scattered pointer list.

Pitfalls that quietly cost you

Four recur often enough to name. Layout transposes. Reshaping [B, S, H, d_h] to [B, H, S, d_h] is a real memory-movement pass; done per layer it can rival the attention matmul on a bandwidth-bound machine, so fuse it into the projection’s output write. Aliased strides. If strideC is smaller than M · ldC, batch items overwrite each other — silently, and only for some tile schedules. Load imbalance. A grouped GEMM scheduled per problem rather than per tile leaves cores idle when one expert dominates. Batching too little. If L is below the core count you have moved the under-utilisation, not fixed it. Parallelise over the batch dimension first: L independent items need no reduction split, which beats splitting K and paying for cross-thread accumulation.

Putting the numbers together

The argument compresses to one inequality. Let I_ridge be your machine’s ridge point — peak FLOP/s divided by peak bytes/s:

I(M, N, K) = 2MNK / (bytes_per_elem · (MK + KN + MN))

compute-bound  ⇔  I ≥ I_ridge

So the design question is never ‘is this matmul big?’ but ‘after batching and reuse, does its intensity clear the ridge?’ Attention heads clear it only when batched, decode projections only when many requests share the weight read, MoE experts only when grouped instead of padded. The batch dimension is not bookkeeping — it decides which side of the roofline you land on.