9. Batching and Bandwidth: Why Decoding Is Slow, and What Helps¶
About the animals in this picture and where they are standing
- Domestic goat (Capra hircus) descends from the bezoar ibex and is among the oldest domesticated animals. There are on the order of a billion of them, kept for milk, meat, fibre and as pack animals.
- Alpine ibex (Capra ibex) lives in the European Alps, usually above the tree line on steep rock, and is a famously sure-footed climber. Males carry horns that can reach about a metre. Hunting left only a small remnant, around the Gran Paradiso in Italy, by the early 1800s; reintroductions have since restored it across the Alps.
- Markhor (Capra falconeri) lives in the mountains of Pakistan, Afghanistan and neighbouring regions. The males' horns twist in a corkscrew; the name is often explained as Persian for 'snake-eater'. It is Pakistan's national animal, and community-based conservation has helped some populations recover.
The scene is imaginary: it is set on an imaginary planet with two suns and a purple moon. The animals are stylised drawings, shown where they could not live, for effect.
What you will understand: the one idea that explains most of the economics of running a language model: decoding is limited by memory traffic, not by arithmetic. You will measure it on the chip, build the standard remedy (batching several requests so they share the weights), see exactly how much of the problem it fixes and what it cannot touch, state the limit with a roofline, and test the batched program builder with the oracle that batching makes possible: a request must get the same answer whether or not other requests ride along.
What you need to know first: Chapter 7 (the chip and its counters) and Chapter 8 (the decode step). Appendix F (memory systems and DRAM) explains where the bandwidth limit comes from.
What this chapter builds: model/batch.py (the batched decode-step builder and its invariance test), tools/ch09_run.py, tools/mut_ch09.py, and for the running examples tools/ch09_example_a.py (a capacity planner) and tools/ch09_example_b.py (static versus continuous batching).
Why this chapter: the economics of running a model¶
Chapter 8 ended with a puzzle. The cached decode step is much cheaper than recomputing, and yet it keeps the matrix array mostly idle, because each step multiplies a single vector by a big matrix. This chapter explains why that is not a quirk of GA-2 but a property of the whole business of serving language models, gives the standard remedy, and measures how far the remedy goes.
Think of a restaurant kitchen whose cook is very fast but whose ingredients are in a warehouse across town. Cooking one meal at a time means a trip across town per meal, and the cook mostly waits. Cooking four meals that use the same ingredients means one trip. Batching is that: the "ingredients" are the model's weights, which must be read from memory for every token, and the "meals" are the requests. The twist the chapter will show is that some ingredients cannot be shared: each meal also needs its own side dish, and that is the KV cache.
Arithmetic intensity: how much work per byte moved¶
Every program does some number of multiply-accumulates (MACs) and moves some number of bytes to and from memory. The ratio is its arithmetic intensity, in MACs per byte. A chip has two ceilings:
- a compute ceiling: the most MACs per cycle it can do (here 16, a 4 x 4 array, at best);
- a memory ceiling: bytes per cycle it can move times the program's intensity.
The attainable speed is the lower of the two. This is the roofline model. The ridge point is the intensity at which the two ceilings meet: below it a program is memory-bound (faster arithmetic would not help), above it compute-bound. For GA-2 the memory moves one word per cycle (one int8 element per word in the design's logical accounting) and the array peaks at 16 MACs per cycle, so the ridge point is 16 MACs per word. A program needs to do 16 MACs for every word it moves just to keep the array busy.
Figure 9.1 (measured points on a derived roofline): the dark line is the best any program could do at each intensity. All of the circuit's decode steps (blue) sit well left of the ridge and below the memory ceiling (orange).
Now look at what a decode step does. Each weight is used once per token: a matrix-vector product reads every element of the matrix and uses it in one multiply. Intensity is about 1 MAC per word, 16 times below the ridge. That is true of every chip and every model: for a single sequence, a decode step is memory-bound by a wide margin.
Batching: reuse the weights¶
If B requests are decoded together, their B new tokens can be stacked as the rows of one matrix, and one matrix product with M = B rows does the projection for all of them. The weights are read once for the whole batch, so intensity rises by up to B. On a 4 x 4 array M = 4 costs almost the same as M = 1, so a batch of four is nearly free in the projections.
Figure 9.2: batching stacks the requests' vectors as rows of one matrix, so the weights stream through once.
What batching cannot share is attention: each request has its own cache, and its scores and weighted sum read only that cache. Those costs add up per request.
The batched program builder: model/batch.py¶
#!/usr/bin/env python3
"""Chapter 9: decode steps for several sequences at once ("batching"), as GA-2 programs.
The B sequences share one set of weights but each has its own hidden states and its own K and V caches, possibly at different lengths (as in a server where requests arrive at different times).
What batching shares: the weights. The projections of the B new tokens are ONE matrix product with M = B rows instead of B products with M = 1, so the weights are loaded once and the 4 x 4 array has B rows to work on.
What it cannot share: attention. Each sequence attends over its own cache, so scores and weighted sums stay per sequence.
Memory plan (words). EXTERNAL (sequence b, 0..3): X_b at 2048 + 256 b, K cache 3072 + 256 b, V cache 4096 + 256 b, output 5120 + 16 b; the shared weights at 512, 768, 1024. Contexts of up to 16 tokens.
SCRATCHPAD: see SP_* (batch slot j holds the j-th sequence of the batch)."""
import math, os, random, sys
os.environ.setdefault("GA2_EXT", "8192")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ga2_isa as I
from quant import quantize, scale_for, mantissa_shift, requant, matmul_int8
from softmax_gold import softmax_fixed
DM, D, LMAX = 16, 16, 16
EX_WQ = 512
def EX_X(b): return 2048 + 256 * b
def EX_KC(b): return 3072 + 256 * b
def EX_VC(b): return 4096 + 256 * b
def EX_OUT(b): return 5120 + 16 * b
SP_X, SP_W, SP_QA, SP_KA, SP_VA, SP_Q8 = 0, 64, 832, 896, 960, 1024
SP_S, SP_S8, SP_P, SP_WT, SP_O, SP_O8 = 1088, 1120, 1152, 1184, 1216, 1232
def SP_KC(j): return 1280 + 512 * j
def SP_VC(j): return 1280 + 512 * j + 256
B_ = I.build
def make_batch(nseq, L, seed):
"""nseq sequences of L tokens with shared weights; ONE set of scales for the whole batch (a model has one set of scales, not one per request)."""
R = random.Random(seed); g = lambda r, c, sd: [[R.gauss(0, sd) for _ in range(c)] for _ in range(r)]
W = [g(DM, D, 1 / math.sqrt(DM)) for _ in range(3)]; Xs = [g(L, DM, 1.0) for _ in range(nseq)]
mm = lambda A, Bm: [[sum(A[i][t] * Bm[t][j] for t in range(len(Bm))) for j in range(len(Bm[0]))] for i in range(len(A))]
flat = lambda Ms: [v for M in Ms for r in M for v in r]
Qs, Ks, Vs = ([mm(X, w) for X in Xs] for w in W)
outs = []
for b in range(nseq):
for t in range(1, L + 1):
q = Qs[b][t - 1]; sc = [sum(q[i] * Ks[b][p][i] for i in range(D)) / math.sqrt(D) for p in range(t)]; mx = max(sc); e = [math.exp(v - mx) for v in sc]; z = sum(e)
outs.append([sum(e[p] / z * Vs[b][p][i] for p in range(t)) for i in range(D)])
sc_ = {"x": scale_for(flat(Xs)), "wq": scale_for(flat([W[0]])), "wk": scale_for(flat([W[1]])), "wv": scale_for(flat([W[2]])), "q": scale_for(flat(Qs)), "k": scale_for(flat(Ks)), "v": scale_for(flat(Vs)), "o": scale_for([v for o in outs for v in o])}
return {"nseq": nseq, "L": L, "Xs": Xs, "W": W, "scales": sc_}
def rq_params(ex):
s = ex["scales"]; ms = mantissa_shift
return {"q": ms(s["x"] * s["wq"] / s["q"]), "k": ms(s["x"] * s["wk"] / s["k"]), "v": ms(s["x"] * s["wv"] / s["v"]), "scores": ms(s["q"] * s["k"] / math.sqrt(D) * 16), "weights": ms(127 / 65536), "out": ms(s["v"] / (127 * s["o"]))}
def ext_image(ex):
s = ex["scales"]; img = [0] * I.EXT; w32 = lambda v: v & 0xFFFFFFFF
for wi, nm in enumerate(("wq", "wk", "wv")):
for r in range(DM):
for c in range(D): img[EX_WQ + 256 * wi + r * D + c] = w32(quantize(ex["W"][wi][r][c], s[nm]))
for b in range(ex["nseq"]):
for t in range(ex["L"]):
for c in range(DM): img[EX_X(b) + t * DM + c] = w32(quantize(ex["Xs"][b][t][c], s["x"]))
return img
def program(seqs, ts, P):
"""One decode step for the sequences `seqs` (indices into the external layout), which have `ts[i]` tokens each including the new one. Batch slot i holds sequence seqs[i]."""
n = len(seqs); prog = []
for j, b in enumerate(seqs): prog.append(B_("LD", dst=SP_X + j * DM, src=EX_X(b) + (ts[j] - 1) * DM, len=DM))
prog.append(B_("LD", dst=SP_W, src=EX_WQ, len=3 * DM * D))
for j, b in enumerate(seqs):
if ts[j] > 1: prog += [B_("LD", dst=SP_KC(j), src=EX_KC(b), len=(ts[j] - 1) * D), B_("LD", dst=SP_VC(j), src=EX_VC(b), len=(ts[j] - 1) * D)]
for wi, acc in enumerate((SP_QA, SP_KA, SP_VA)): # the projections: ONE matrix product with M = n rows per column block
for c in range(D // 4): prog.append(B_("MM", dst=acc + 4 * c, A=SP_X, B=SP_W + 256 * wi + 4 * c, M=n, K=DM, N=4, tb=0, lda=DM, ldb=D, ldc=D))
prog.append(B_("RQ", dst=SP_Q8, src=SP_QA, len=n * D, m=P["q"][0], s=P["q"][1]))
for j, b in enumerate(seqs):
r = (ts[j] - 1) * D
prog += [B_("RQ", dst=SP_KC(j) + r, src=SP_KA + j * D, len=D, m=P["k"][0], s=P["k"][1]), B_("RQ", dst=SP_VC(j) + r, src=SP_VA + j * D, len=D, m=P["v"][0], s=P["v"][1]),
B_("ST", src=SP_KC(j) + r, dst=EX_KC(b) + r, len=D), B_("ST", src=SP_VC(j) + r, dst=EX_VC(b) + r, len=D)]
for j, b in enumerate(seqs): # attention, one sequence at a time
L = ts[j]
for c in range((L + 3) // 4): prog.append(B_("MM", dst=SP_S + 4 * c, A=SP_Q8 + j * D, B=SP_KC(j) + 4 * c * D, M=1, K=D, N=min(4, L - 4 * c), tb=1, lda=D, ldb=D, ldc=4))
prog += [B_("RQ", dst=SP_S8, src=SP_S, len=L, m=P["scores"][0], s=P["scores"][1]), B_("SM", dst=SP_P, src=SP_S8, len=L), B_("RQ", dst=SP_WT, src=SP_P, len=L, m=P["weights"][0], s=P["weights"][1])]
for c in range(D // 4): prog.append(B_("MM", dst=SP_O + 4 * c, A=SP_WT, B=SP_VC(j) + 4 * c, M=1, K=L, N=4, tb=0, lda=L, ldb=D, ldc=4))
prog += [B_("RQ", dst=SP_O8, src=SP_O, len=D, m=P["out"][0], s=P["out"][1]), B_("ST", src=SP_O8, dst=EX_OUT(b), len=D)]
return prog + [B_("HALT")]
def outputs(ext, seqs): return {b: [I.s8(ext[EX_OUT(b) + i]) for i in range(D)] for b in seqs}
def int_reference(ex, b, t):
"""Plain-Python integer result of sequence b at its t-th token (no instructions, no addresses)."""
s = ex["scales"]; P = rq_params(ex)
Xq = [[quantize(v, s["x"]) for v in r] for r in ex["Xs"][b][:t]]; Wq = [[[quantize(v, s[nm]) for v in r] for r in W] for W, nm in zip(ex["W"], ("wq", "wk", "wv"))]
rqm = lambda M, key: [[requant(a, *P[key]) for a in r] for r in M]
q8 = rqm(matmul_int8([Xq[-1]], Wq[0]), "q")[0]; K8 = rqm(matmul_int8(Xq, Wq[1]), "k"); V8 = rqm(matmul_int8(Xq, Wq[2]), "v")
sc = [requant(sum(q8[i] * K8[p][i] for i in range(D)), *P["scores"]) for p in range(t)]; w = [requant(v, *P["weights"]) for v in softmax_fixed(sc)]
return [requant(sum(w[p] * V8[p][c] for p in range(t)), *P["out"]) for c in range(D)]
def run(ext, prog): return I.Machine(ext).run(prog)
def traffic(prog):
"""(MACs, words moved to or from external memory) of a program. One word is one int8 element in this design's logical accounting."""
macs = sum(i["f"] * i["g"] * i["h"] for i in prog if i["op"] == "MM"); words = sum(i["c"] for i in prog if i["op"] in ("LD", "ST")); return macs, words
def check_batching(nseq_list=(1, 2, 3, 4), lengths=(1, 2, 5, 9, 16), seeds=range(2), exact=True, ragged=True):
"""Batch invariance: at every step, each sequence's output and cache rows must equal what the same sequence produces when decoded ALONE (batch of one). Also ragged batches (different lengths).
And each sequence decoded alone must equal the integer reference exactly (batch invariance alone is blind to a bug that the batch and the single run share). Returns the list of problems."""
problems = []
for nseq in nseq_list:
for L in lengths:
for seed in seeds:
ex = make_batch(nseq, L, 7000 + 31 * nseq + 5 * L + seed); P = rq_params(ex)
ext_b = ext_image(ex); ext_s = [list(ext_image(ex)) for _ in range(nseq)]; seqs = list(range(nseq))
for t in range(1, L + 1):
m = run(ext_b, program(seqs, [t] * nseq, P)); ext_b = m.ext
for b in seqs:
ext_s[b] = run(ext_s[b], program([b], [t], P)).ext
if outputs(ext_b, [b])[b] != outputs(ext_s[b], [b])[b]: problems.append(f"B={nseq} L={L} seed={seed} step {t}: sequence {b} output differs from decoding it alone")
if ext_b[EX_KC(b):EX_KC(b) + 256] != ext_s[b][EX_KC(b):EX_KC(b) + 256] or ext_b[EX_VC(b):EX_VC(b) + 256] != ext_s[b][EX_VC(b):EX_VC(b) + 256]: problems.append(f"B={nseq} L={L} seed={seed} step {t}: sequence {b} cache differs")
if exact and outputs(ext_s[b], [b])[b] != int_reference(ex, b, t): problems.append(f"B={nseq} L={L} seed={seed} step {t}: sequence {b} decoded alone differs from the integer reference")
# ragged: sequences at different positions share one step
for seed in (seeds if ragged else []):
ex = make_batch(4, 16, 9100 + seed); P = rq_params(ex); lens = [3, 9, 1, 14]; ext_b = ext_image(ex); seqs = [0, 1, 2, 3]
for b in seqs:
for t in range(1, lens[b]): ext_b = run(ext_b, program([b], [t], P)).ext
ext_s = [list(ext_b) for _ in seqs]; ext_b = run(ext_b, program(seqs, lens, P)).ext
for b in seqs:
ext_s[b] = run(ext_s[b], program([b], [lens[b]], P)).ext
if outputs(ext_b, [b])[b] != outputs(ext_s[b], [b])[b]: problems.append(f"ragged seed={seed}: sequence {b} (length {lens[b]}) differs from decoding it alone")
if exact and outputs(ext_s[b], [b])[b] != int_reference(ex, b, lens[b]): problems.append(f"ragged seed={seed}: sequence {b} differs from the integer reference")
return problems
if __name__ == "__main__":
print("problems:", check_batching())
program(seqs, ts, P) builds one decode step for a list of sequences, with ts[i] tokens each (the lengths may differ: a server's requests arrive at different times). The projections are three loops of four MM instructions with M = n; the caches are per sequence; attention is a loop over sequences. check_batching is the test, explained below.
Commentary. Read program from the top. First one LD per sequence brings in its hidden state (one row of SP_X), then one LD brings in all three weight matrices: that is the shared traffic. Then each sequence loads its own K and V cache rows (SP_KC(j), SP_VC(j)): the per-sequence traffic. The three projections are loops of MM with M = n, the number of sequences, and are the only instructions where batching shows up as a different M. The new key and value rows are appended to each sequence's own cache at its own position ts[j] - 1. The final loop does attention for each sequence separately, with L = ts[j]. Because ts can differ per sequence, the program handles requests at different points of their generation, which is the normal case in a server.
Running it on the chip¶
#!/usr/bin/env python3
"""Chapter 9: batching and the bandwidth analysis. Batched decode steps on the RTL (both simulators), cycles and traffic per batch size, the roofline of GA-2, and a derived table for a large model. Usage: ch09_run.py"""
import os, re, sys
os.environ["GA2_EXT"] = "8192"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "model"))
import hw, ga2_isa as I, ga2_progs as G, batch as Bt
F = ["rtl/sram.v", "rtl/dma.v", "rtl/systolic.v", "rtl/requant.v", "rtl/exp_lut.v", "rtl/divu.v", "rtl/rowmem.v", "rtl/ga2_vec.v", "rtl/ga2_sm.v", "rtl/ga2_mm.v", "rtl/ga2.v", "tb/extmem_rw.v", "tb/ga2_tb.v"]
def line(out): return ([l for l in out.splitlines() if l.startswith(("PASS", "FAIL", "MISMATCH"))] or out.strip().splitlines()[-2:])[0]
print("== 1. batch invariance (Python): each sequence's result and cache are identical whether it is decoded alone or in a batch of 2, 3, 4, or in a ragged batch")
print("problems:", Bt.check_batching() or "none")
print("\n== 2. batched decode steps on the circuit: B sequences, all at their 16th token, then a ragged batch (lengths 3, 9, 1, 14)")
progs = []; meta = []
for nseq in (1, 2, 3, 4):
ex = Bt.make_batch(nseq, 16, 4242); P = Bt.rq_params(ex); ext = Bt.ext_image(ex); seqs = list(range(nseq))
for t in range(1, 16): ext = Bt.run(ext, Bt.program(seqs, [t] * nseq, P)).ext
prog = Bt.program(seqs, [16] * nseq, P); progs.append((prog, ext)); meta.append((f"B={nseq} (all at token 16)", nseq, prog))
ex = Bt.make_batch(4, 16, 4343); P = Bt.rq_params(ex); ext = Bt.ext_image(ex); lens = [3, 9, 1, 14]
for b in range(4):
for t in range(1, lens[b]): ext = Bt.run(ext, Bt.program([b], [t], P)).ext
prog = Bt.program([0, 1, 2, 3], lens, P); progs.append((prog, ext)); meta.append(("B=4 ragged (3, 9, 1, 14)", 4, prog))
words = [len(progs)]
for prog, ext in progs: words += G.record(prog, ext)
open(os.path.join(hw.ROOT, "out", "ga2_programs.hex"), "w").write("\n".join("%08x" % (w & 0xFFFFFFFF) for w in words) + "\n")
res = {}
for name, fn in (("Icarus", hw.sim_icarus), ("Verilator", hw.sim_verilator)):
rc, out = fn(F, "ga2_tb", defines=("EXTW=8192", "DUMP")); print(f" {name:10s}{line(out)} (exit {rc})")
for m in re.finditer(r"COUNTERS (\d+) (\d+) (\d+) (\d+) (\d+) (\d+)", out): res[int(m.group(1))] = [int(g) for g in m.groups()[1:]]
print("\n== 3. cycles, traffic and the roofline. Peak compute = 16 MAC/cycle (a 4x4 array). Memory = 1 word per cycle once a transfer is under way (one word = one int8 element).")
print(f"{'batch':26s} {'cycles':>7s} {'per token':>9s} {'MACs':>6s} {'words':>6s} {'MAC/word':>9s} {'achieved':>9s} {'ceiling':>8s} {'of ceiling':>10s} {'if loads overlapped':>20s}")
for i, (name, nseq, prog) in enumerate(meta):
tot, mm, dma, vec, n = res[i]; macs, wds = Bt.traffic(prog); inten = macs / wds; ceil = min(16.0, inten * 1.0); ach = macs / tot
ov = tot / max(dma, tot - dma)
print(f"{name:26s} {tot:7d} {tot/nseq:9.0f} {macs:6d} {wds:6d} {inten:9.2f} {ach:9.3f} {ceil:8.2f} {100*ach/ceil:9.1f}% {ov:15.2f}x faster")
print("\n ('per token' = cycles / number of sequences in the batch; 'achieved' = MACs / cycles; 'ceiling' = min(peak, MAC/word x words/cycle); the last column is the best case if every load ran in the shadow of compute, derived from the counters, not built)")
print("\n== 4. what batching can and cannot share (one step, L = 16 tokens per sequence; words from the programs, MACs counted from the MM instructions)")
macs_proj = 3 * 16 * 16; macs_att = 2 * 16 * 16; words_w = 3 * 16 * 16; words_seq = 16 + 2 * 15 * 16 + 2 * 16 + 16
print(f"shared by the batch : weights {words_w} words; projections {macs_proj} MACs per sequence (these reuse the weights)")
print(f"per sequence : cache and I/O {words_seq} words; attention {macs_att} MACs")
for nseq in (1, 2, 4, 16, 64, 1000000): print(f" batch {nseq:>7d}: {(nseq*(macs_proj+macs_att))/(words_w+nseq*words_seq):5.2f} MAC/word")
print(f" limit as the batch grows: {(macs_proj+macs_att)/words_seq:.2f} MAC/word at L = 16; as L grows the limit tends to 1: every cached word is used once per step.")
print("\n== 5. derived, for a 7B-class model (7e9 weights; int8 cache of 256 KiB per token, 4096-token context = 1.07e9 bytes per sequence), ASSUMING a chip with 1 TB/s of memory bandwidth and ignoring compute and overheads")
print(" tokens per second per chip = B x BW / (weight bytes + B x cache bytes)")
KV = 2 * 32 * 32 * 128 * 1 * 4096; BW = 1e12
print(f"{'weights':14s} {'bytes':>8s} " + " ".join(f"{'B='+str(b):>8s}" for b in (1, 4, 16, 64, 256)) + f" {'cache = weights at B':>22s}")
for name, bpw in (("fp16", 2.0), ("int8", 1.0), ("int4", 0.5)):
Wb = 7e9 * bpw; row = " ".join(f"{b*BW/(Wb+b*KV):8.0f}" for b in (1, 4, 16, 64, 256)); print(f"{name:14s} {Wb/1e9:6.1f}GB {row} {Wb/KV:22.1f}")
print(f" upper limit as B grows (cache-bound): {BW/KV:.0f} tokens/s per chip at a 4096-token int8 cache")
Output (cloud sandbox -- live-executed; Icarus Verilog 12.0, Verilator 5.020, Yosys 0.33)
== 1. batch invariance (Python): each sequence's result and cache are identical whether it is decoded alone or in a batch of 2, 3, 4, or in a ragged batch
problems: none
== 2. batched decode steps on the circuit: B sequences, all at their 16th token, then a ragged batch (lengths 3, 9, 1, 14)
Icarus PASS: 5 programs match the Python reference (memory and all five cycle counters), 33460 cycles in total (exit status 0) (exit 0)
Verilator PASS: 5 programs match the Python reference (memory and all five cycle counters), 33460 cycles in total (exit status 0) (exit 0)
== 3. cycles, traffic and the roofline. Peak compute = 16 MAC/cycle (a 4x4 array). Memory = 1 word per cycle once a transfer is under way (one word = one int8 element).
batch cycles per token MACs words MAC/word achieved ceiling of ceiling if loads overlapped
B=1 (all at token 16) 3827 3827 1280 1312 0.98 0.334 0.98 34.3% 1.55x faster
B=2 (all at token 16) 5777 2888 2560 1856 1.38 0.443 1.38 32.1% 1.50x faster
B=3 (all at token 16) 7727 2576 3840 2400 1.60 0.497 1.60 31.1% 1.48x faster
B=4 (all at token 16) 9677 2419 5120 2944 1.74 0.529 1.74 30.4% 1.47x faster
B=4 ragged (3, 9, 1, 14) 6452 1613 3936 1760 2.24 0.610 2.24 27.3% 1.41x faster
('per token' = cycles / number of sequences in the batch; 'achieved' = MACs / cycles; 'ceiling' = min(peak, MAC/word x words/cycle); the last column is the best case if every load ran in the shadow of compute, derived from the counters, not built)
== 4. what batching can and cannot share (one step, L = 16 tokens per sequence; words from the programs, MACs counted from the MM instructions)
shared by the batch : weights 768 words; projections 768 MACs per sequence (these reuse the weights)
per sequence : cache and I/O 544 words; attention 512 MACs
batch 1: 0.98 MAC/word
batch 2: 1.38 MAC/word
batch 4: 1.74 MAC/word
batch 16: 2.16 MAC/word
batch 64: 2.30 MAC/word
batch 1000000: 2.35 MAC/word
limit as the batch grows: 2.35 MAC/word at L = 16; as L grows the limit tends to 1: every cached word is used once per step.
== 5. derived, for a 7B-class model (7e9 weights; int8 cache of 256 KiB per token, 4096-token context = 1.07e9 bytes per sequence), ASSUMING a chip with 1 TB/s of memory bandwidth and ignoring compute and overheads
tokens per second per chip = B x BW / (weight bytes + B x cache bytes)
weights bytes B=1 B=4 B=16 B=64 B=256 cache = weights at B
fp16 14.0GB 66 219 513 774 886 13.0
int8 7.0GB 124 354 662 845 908 6.5
int4 3.5GB 219 513 774 886 920 3.3
upper limit as B grows (cache-bound): 931 tokens/s per chip at a 4096-token int8 cache
Reading it¶
- The batched programs run correctly on the chip (five programs, both simulators, memory and cycle counters exactly as the reference), including a ragged batch where four requests are at tokens 3, 9, 1 and 14 in the same step.
-
Batching helps, by less than you might hope. Cycles per token fall from 3,827 (batch of 1) to 2,419 (batch of 4), a 1.58x gain, and the MACs per word rise from 0.98 to 1.74. The projections are shared, but the attention work, which does not shrink, becomes the bulk of the step.
Figure 9.3: the traffic is a fixed 768 words of weights plus 544 words per request. -
The words moved per step follow a formula, and the measured programs match it exactly:
768 + 544 B(768 words of weights, loaded once; 544 words of cache and input/output per sequence at L = 16): 1,312 for one request, 2,944 for four, as measured. Section 4 extends the formula: with an unlimited batch the intensity tends to 2.35 MACs per word at L = 16, and to 1 as the context grows, because each cached key and value is used for exactly one multiply per step. No batch size gets attention across the ridge point.
Figure 9.4 (derived from the formula, matching the circuit at B = 1, 2, 4): the intensity of the whole step rises with the batch but flattens at 2.35 MACs per word at L = 16.
- The achieved speed is about a third of the memory ceiling and rises no higher. achieved (MACs per cycle) is 0.33 to 0.53; the ceiling is the intensity itself, 0.98 to 1.74. About 30% of it is reached, whatever the batch. The reasons are in Chapter 7: instructions run one at a time (a load never overlaps a compute), and the matrix unit stages operands before streaming them. The last column estimates what overlapping would give if every load ran in the shadow of compute: 1.4x to 1.55x faster. That is derived from the counters (total divided by the larger of the DMA and non-DMA times), not built, and the gain would be the double buffering of Chapter 5.
What it means for a real model (derived, with an assumed chip)¶
The last table is arithmetic, not a measurement, and it assumes a chip with 1 TB/s of memory bandwidth that is limited only by memory traffic. Tokens per second per chip is B x BW / (weight bytes + B x cache bytes) for a 7B-class model with an int8 cache of 256 KiB per token and a 4,096-token context (1.07 GB per request):
- At batch 1, the speed is set by the weights alone: 66, 124 and 219 tokens/s for fp16, int8 and int4 weights. Halving the bytes per weight doubles the speed, which is why inference chips push to 8 and 4 bits.
- As the batch grows, throughput rises and then flattens at 931 tokens/s per chip, whatever the weights: the cache of each request must be read once per token, and that traffic does not amortize.
- The cache equals the weights in size at a batch of 13 (fp16), 6.5 (int8) and 3.3 (int4). Quantizing the weights more aggressively makes batching run out of road sooner. This is why long-context serving is limited by the size and bandwidth of the KV cache, and why techniques that shrink the cache (fewer key-value heads shared by many query heads, a smaller cache dtype, evicting old tokens) matter as much as ones that shrink the weights.
Running example A: a capacity planner¶
The point of this example: the table above assumes a chip with unlimited memory. Real chips have a capacity as well as a bandwidth, and capacity decides how many requests can be in the batch at all. This script is a small planner: give it a model, a context length and a chip, and it says how many requests fit, how fast the chip can go at each batch size, and which limit you hit first. Everything it prints is derived arithmetic for assumed chips (labelled A, B and C; none is a real product), meant to be edited.
#!/usr/bin/env python3
"""Chapter 9, running example A: a capacity planner for decode serving. Given a model (weight bytes, KV-cache bytes per token), a context length, and a chip's memory capacity and bandwidth, it answers: how many requests fit, how fast can the chip go at each batch size, and which resource is the limit? All figures are DERIVED arithmetic for ASSUMED chips (no real product is meant); change the numbers at the top to plan for your own. Writes out/ch09_example_a.json. Usage: ch09_example_a.py"""
import json, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); import hw
GiB = 2 ** 30
MODEL = {"name": "7B-class, int8 weights, int8 KV cache", "weight_bytes": 7e9, "kv_per_token": 2 * 32 * 32 * 128 * 1} # 256 KiB per token
CHIPS = [("chip A: 24 GB, 0.5 TB/s", 24e9, 0.5e12), ("chip B: 80 GB, 2 TB/s", 80e9, 2.0e12), ("chip C: 192 GB, 5 TB/s", 192e9, 5.0e12)]
def plan(chip_cap, bw, ctx, model=MODEL):
kv = model["kv_per_token"] * ctx; free = chip_cap - model["weight_bytes"]; maxb = max(0, int(free // kv))
def tps(b): return b * bw / (model["weight_bytes"] + b * kv)
return kv, maxb, tps
res = {"chips": [], "ctx": [], "tps": {}, "maxb": {}}
print(f"model: {MODEL['name']}: weights {MODEL['weight_bytes']/1e9:.1f} GB, KV cache {MODEL['kv_per_token']/1024:.0f} KiB per token\n")
print("tokens per second = B x BW / (weight bytes + B x KV bytes per request) [memory-bound, one full read of the weights and every cache per step]")
print("batch size is limited by capacity: weights + B x KV <= memory\n")
for ctx in (2048, 8192, 32768):
kv_ = MODEL["kv_per_token"] * ctx
print(f"-- context {ctx:,} tokens: KV cache {kv_/GiB:.2f} GiB per request --")
print(f"{'chip':26s} {'max batch (capacity)':>21s} " + " ".join(f"{'B='+str(b):>7s}" for b in (1, 4, 16, 64)) + f" {'best tokens/s':>14s} {'limit as B grows':>17s}")
for name, cap, bw in CHIPS:
kv, maxb, tps = plan(cap, bw, ctx); row = [(f"{tps(b):7.0f}" if b <= maxb else f"{'--':>7s}") for b in (1, 4, 16, 64)]
best = tps(maxb) if maxb else 0; lim = bw / kv
print(f"{name:26s} {maxb:21d} " + " ".join(row) + f" {best:14.0f} {lim:17.0f}")
res["tps"][f"{name}|{ctx}"] = [tps(b) if b <= maxb else None for b in range(1, 129)]; res["maxb"][f"{name}|{ctx}"] = maxb
print()
res["chips"] = [c[0] for c in CHIPS]
print("reading it:")
print(" * at batch 1 the speed is bandwidth / weight bytes: 0.5e12 / 7e9 = 71 tokens/s on chip A, 286 on chip B, 714 on chip C (before the cache read).")
print(" * chip A holds exactly one 32,768-token request next to the weights (7 GB + 8.6 GB = 15.6 GB; a second request would need 24.2 GB): look at the 'max batch' column.")
print(" * past a point more batch adds almost nothing: the curve flattens at bandwidth / KV bytes per request, the 'limit' column. Long contexts lower that limit and the max batch together.")
print(" * on these assumed chips capacity ends the curve before it has flattened (chip A at 2,048 tokens reaches 656 of its 931 limit): a bigger memory buys more batch, and more batch buys throughput only until the curve flattens.")
json.dump(res, open(os.path.join(hw.ROOT, "out", "ch09_example_a.json"), "w"))
To compile and run: python3 tools/ch09_example_a.py (pure Python, instant).
Output (cloud sandbox -- live-executed)
model: 7B-class, int8 weights, int8 KV cache: weights 7.0 GB, KV cache 256 KiB per token
tokens per second = B x BW / (weight bytes + B x KV bytes per request) [memory-bound, one full read of the weights and every cache per step]
batch size is limited by capacity: weights + B x KV <= memory
-- context 2,048 tokens: KV cache 0.50 GiB per request --
chip max batch (capacity) B=1 B=4 B=16 B=64 best tokens/s limit as B grows
chip A: 24 GB, 0.5 TB/s 31 66 219 513 -- 656 931
chip B: 80 GB, 2 TB/s 135 265 875 2053 3095 3397 3725
chip C: 192 GB, 5 TB/s 344 663 2186 5132 7737 8973 9313
-- context 8,192 tokens: KV cache 2.00 GiB per request --
chip max batch (capacity) B=1 B=4 B=16 B=64 best tokens/s limit as B grows
chip A: 24 GB, 0.5 TB/s 7 55 128 -- -- 159 233
chip B: 80 GB, 2 TB/s 33 219 513 774 -- 848 931
chip C: 192 GB, 5 TB/s 86 547 1283 1934 2215 2243 2328
-- context 32,768 tokens: KV cache 8.00 GiB per request --
chip max batch (capacity) B=1 B=4 B=16 B=64 best tokens/s limit as B grows
chip A: 24 GB, 0.5 TB/s 1 32 -- -- -- 32 58
chip B: 80 GB, 2 TB/s 8 128 193 -- -- 211 233
chip C: 192 GB, 5 TB/s 21 321 484 554 -- 560 582
reading it:
* at batch 1 the speed is bandwidth / weight bytes: 0.5e12 / 7e9 = 71 tokens/s on chip A, 286 on chip B, 714 on chip C (before the cache read).
* chip A holds exactly one 32,768-token request next to the weights (7 GB + 8.6 GB = 15.6 GB; a second request would need 24.2 GB): look at the 'max batch' column.
* past a point more batch adds almost nothing: the curve flattens at bandwidth / KV bytes per request, the 'limit' column. Long contexts lower that limit and the max batch together.
* on these assumed chips capacity ends the curve before it has flattened (chip A at 2,048 tokens reaches 656 of its 931 limit): a bigger memory buys more batch, and more batch buys throughput only until the curve flattens.
Figure 9.5 (derived): throughput against batch size for the three assumed chips. Each curve stops where the memory is full.
Walkthrough.
- The formula. A step reads the weights once and every request's cache once, so a chip with bandwidth BW does
B x BW / (weights + B x KV)tokens per second. The KV term is the context length times 256 KiB per token (this model, int8 cache). - The capacity limit.
weights + B x KV <= memorygives the largest batch. At a 2,048-token context chip B can hold 135 requests; at 32,768 tokens only 8: the cache grows 16 times, so the batch shrinks about 16 times. - Where the curve flattens. As B grows the weights become negligible and the speed tends to
BW / KV(the "limit" column): 3,725 tokens/s for chip B at 2,048 tokens, 233 at 32,768. Long context lowers both the ceiling and the batch that can reach it. - The practical lessons. (a) At small batch the speed is set by the weights: 66, 265, 663 tokens/s on chips A, B, C at B = 1. (b) Memory capacity is what lets you batch, and bandwidth is what you get for it. (c) A larger context costs throughput twice: fewer requests fit and each step reads more cache. This is the quantitative reason the techniques of Chapters 13 to 15 (smaller weights, shared KV heads, evicting old cache) exist.
Running example B: static against continuous batching¶
The point of this example: so far every batch had requests at the same point. A real server's requests start and end at different times. There are two ways to run a batch. Static batching takes four requests and runs the group until the longest one finishes; shorter requests finish early but their slots ride along idle. Continuous batching refills a slot the moment its request ends, so the batch stays full. This example runs both on the cycle model of the real batched program of this chapter (the step costs are the program's cycle counts, equal to the circuit's counters as the run above showed), for eight requests that want 14, 3, 9, 5, 12, 2, 7 and 10 tokens, with room for four at a time.
#!/usr/bin/env python3
"""Chapter 9, running example B: static against continuous batching. Eight requests with different output lengths are served on GA-2 with a batch capacity of four. STATIC batching takes four requests and runs until the longest finishes; CONTINUOUS batching refills a slot the moment its request ends. The cost of every decode step is the cycle count of the REAL batched program (model/batch.py) for the sequences present, from the cycle model (equal to the circuit's counters, Chapter 9's run). Writes out/ch09_example_b.json. Usage: ch09_example_b.py"""
import json, os, sys
os.environ["GA2_EXT"] = "8192"
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); import hw
sys.path.insert(0, os.path.join(hw.ROOT, "model")); import ga2_isa as I, batch as Bt
P = Bt.rq_params(Bt.make_batch(4, 16, 1)); LENS = [14, 3, 9, 5, 12, 2, 7, 10]; CAP = 4; cache = {}
def step_cost(ts):
key = tuple(ts)
if key not in cache: cache[key] = I.cycle_counts(Bt.program(list(range(len(ts))), list(ts), P))[0]
return cache[key]
def static():
t = 0; done = {}; log = []
for g in range(0, len(LENS), CAP):
grp = list(range(g, min(g + CAP, len(LENS)))); steps = max(LENS[i] for i in grp)
for s in range(1, steps + 1):
ts = [s for _ in grp] # finished requests keep riding (padding): the batch is rigid until the longest ends
c = step_cost(ts); log.append((t, c, len(grp), [i for i in grp if s <= LENS[i]])); t += c
for i in grp:
if s == LENS[i]: done[i] = t
return t, done, log
def continuous():
t = 0; waiting = list(range(len(LENS))); slots = []; prog = {}; done = {}; log = []
while waiting or slots:
while waiting and len(slots) < CAP: i = waiting.pop(0); slots.append(i); prog[i] = 0
for i in slots: prog[i] += 1
c = step_cost([prog[i] for i in slots]); log.append((t, c, len(slots), list(slots))); t += c
for i in list(slots):
if prog[i] == LENS[i]: done[i] = t; slots.remove(i)
return t, done, log
ts_, ds, ls = static(); tc, dc, lc = continuous(); tokens = sum(LENS)
print(f"requests (output tokens each): {LENS} batch capacity {CAP} total tokens {tokens}\n")
print("cost of one decode step on GA-2 (cycles), from the batched program, for a few batch shapes (every sequence at the same token t):")
for n in (1, 2, 3, 4): print(f" {n} sequence(s) at token 8: {step_cost([8]*n):6d} cycles ({step_cost([8]*n)/n:6.0f} per token)")
print()
for name, tot, d in (("static", ts_, ds), ("continuous", tc, dc)):
lat = [d[i] for i in range(len(LENS))]
print(f"{name:11s} total {tot:7d} cycles for {tokens} tokens = {tot/tokens:7.0f} cycles per token; request completion time: mean {sum(lat)/len(lat):8.0f}, worst {max(lat):7d}")
print(f"\ncontinuous batching finishes the same work {ts_/tc:.2f}x faster and the mean request completes {sum(ds.values())/sum(dc.values()):.2f}x sooner.")
print("steps with how many sequences really generating a token (static, including the idle padded ones in brackets):")
print(" static :", [f"{len(x[3])}/{x[2]}" for x in ls])
print(" continuous:", [f"{len(x[3])}" for x in lc])
useful_s = sum(len(x[3]) for x in ls); slots_s = sum(x[2] for x in ls)
print(f"\nstatic batching spent {slots_s - useful_s} of {slots_s} sequence-slots on requests that had already finished ({100*(slots_s-useful_s)/slots_s:.0f}% waste); continuous batching spent none (every slot is a live request).")
print("The ragged batch of Chapter 9's circuit run (tokens 3, 9, 1, 14 at the same step) is exactly what continuous batching creates all the time.")
json.dump({"lens": LENS, "static": {"total": ts_, "done": ds, "log": ls}, "continuous": {"total": tc, "done": dc, "log": lc}}, open(os.path.join(hw.ROOT, "out", "ch09_example_b.json"), "w"))
To compile and run: python3 tools/ch09_example_b.py (needs nothing but Python; the cycle model builds the batched programs but does not simulate them).
Output (cloud sandbox -- live-executed)
requests (output tokens each): [14, 3, 9, 5, 12, 2, 7, 10] batch capacity 4 total tokens 62
cost of one decode step on GA-2 (cycles), from the batched program, for a few batch shapes (every sequence at the same token t):
1 sequence(s) at token 8: 3115 cycles ( 3115 per token)
2 sequence(s) at token 8: 4353 cycles ( 2176 per token)
3 sequence(s) at token 8: 5591 cycles ( 1864 per token)
4 sequence(s) at token 8: 6829 cycles ( 1707 per token)
static total 170118 cycles for 62 tokens = 2744 cycles per token; request completion time: mean 93170, worst 170118
continuous total 98964 cycles for 62 tokens = 1596 cycles per token; request completion time: mean 59092, worst 98964
continuous batching finishes the same work 1.72x faster and the mean request completes 1.58x sooner.
steps with how many sequences really generating a token (static, including the idle padded ones in brackets):
static : ['4/4', '4/4', '4/4', '3/4', '3/4', '2/4', '2/4', '2/4', '2/4', '1/4', '1/4', '1/4', '1/4', '1/4', '4/4', '4/4', '3/4', '3/4', '3/4', '3/4', '3/4', '2/4', '2/4', '2/4', '1/4', '1/4']
continuous: ['4', '4', '4', '4', '4', '4', '4', '4', '4', '4', '4', '4', '4', '4', '2', '1', '1', '1', '1']
static batching spent 42 of 104 sequence-slots on requests that had already finished (40% waste); continuous batching spent none (every slot is a live request).
The ragged batch of Chapter 9's circuit run (tokens 3, 9, 1, 14 at the same step) is exactly what continuous batching creates all the time.
Figure 9.6: static batching (top) leaves slots idle (pink) while it waits for the longest request of the group; continuous batching (bottom) keeps every slot busy.
Walkthrough.
- The cost of a step. The first table shows that GA-2 batching works as the chapter says: one sequence at token 8 costs 3,115 cycles; four of them cost 6,829, i.e. 1,707 per token. The batch shares the projections, not the attention.
- Static. The first group of four runs 14 steps (until the 14-token request ends) but the 3-, 5- and 9-token requests are done after 3, 5 and 9 steps; their slots ride along. The second group does the same with a 12-token request. 170,118 cycles in total.
- Continuous. Whenever a request finishes, the next waiting one takes the slot. The batch is four-wide for 14 steps, then drains as the queue empties. 98,964 cycles in total: 1.72 times faster for the same 62 tokens.
- Latency too. The mean completion time falls 1.58 times, because a short request no longer waits for the group's longest.
- The waste. Static batching spends 42 of its 104 sequence-slots on finished requests (40%). Continuous batching spends none.
- Why the chip's ragged batch matters. The circuit run in this chapter (lengths 3, 9, 1, 14 in one step) is not an exotic test: it is the normal state of a continuously batched server, and the test suite exercises it because the program had to be correct there.
Testing the tests: batch invariance¶
A batched program is easy to get subtly wrong: a row index off, a cache in the wrong slot, a length taken from the wrong request. The test that finds all of these has a beautiful form that needs no knowledge of the arithmetic: a request must produce the same output and the same cache rows whether it is decoded alone or in a batch with others. check_batching runs batches of 1 to 4 and a ragged batch, compares every sequence with the same sequence decoded alone at every step, and also compares each sequence decoded alone with the integer reference.
#!/usr/bin/env python3
"""Chapter 9: test the tests of the batched program builder (model/batch.py). The check is BATCH INVARIANCE: each sequence's output and cache rows must equal what the same sequence gives when decoded alone, for batches of 1 to 4 and for a ragged batch;
a batch of one must also equal the integer reference. Each mutant breaks one line of the builder; 'caught' means the check reports a problem or crashes."""
import concurrent.futures, os, shutil, subprocess, sys, tempfile
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))); SRC = open(os.path.join(ROOT, "model", "batch.py")).read()
M = [
("projection runs on one row (M = 1) instead of the whole batch", "A=SP_X, B=SP_W + 256 * wi + 4 * c, M=n,", "A=SP_X, B=SP_W + 256 * wi + 4 * c, M=1,"),
("projection results are stored with row stride 4, not d", "tb=0, lda=DM, ldb=D, ldc=D))", "tb=0, lda=DM, ldb=D, ldc=4))"),
("every new key is taken from row 0 of the projection", "src=SP_KA + j * D, len=D", "src=SP_KA, len=D"),
("every new value is taken from row 0 of the projection", "src=SP_VA + j * D, len=D", "src=SP_VA, len=D"),
("every sequence's cache lives in batch slot 0", "def SP_KC(j): return 1280 + 512 * j", "def SP_KC(j): return 1280"),
("the value caches of all sequences overlap", "def SP_VC(j): return 1280 + 512 * j + 256", "def SP_VC(j): return 1280 + 256 * j + 256"),
("sequence j attends with the query of row 0", "A=SP_Q8 + j * D, B=SP_KC(j)", "A=SP_Q8, B=SP_KC(j)"),
("hidden states are all loaded into slot 0", "dst=SP_X + j * DM, src=EX_X(b)", "dst=SP_X, src=EX_X(b)"),
("every sequence reads the hidden state of sequence 0", "src=EX_X(b) + (ts[j] - 1) * DM", "src=EX_X(0) + (ts[j] - 1) * DM"),
("outputs are all stored in sequence 0's output slot", "B_(\"ST\", src=SP_O8, dst=EX_OUT(b), len=D)", "B_(\"ST\", src=SP_O8, dst=EX_OUT(0), len=D)"),
("the new key is appended at the first sequence's position", "r = (ts[j] - 1) * D", "r = (ts[0] - 1) * D"),
("the value cache is loaded from the key cache", "B_(\"LD\", dst=SP_VC(j), src=EX_VC(b)", "B_(\"LD\", dst=SP_VC(j), src=EX_KC(b)"),
("the new value is appended to the key cache in external memory", "B_(\"ST\", src=SP_VC(j) + r, dst=EX_VC(b) + r", "B_(\"ST\", src=SP_VC(j) + r, dst=EX_KC(b) + r"),
("attention length of every sequence is the first sequence's", "L = ts[j]\n", "L = ts[0]\n"),
("the q requantization covers only one row", "src=SP_QA, len=n * D", "src=SP_QA, len=D"),
("a sequence's cache is loaded with the first sequence's length", "len=(ts[j] - 1) * D), B_(\"LD\", dst=SP_VC(j), src=EX_VC(b), len=(ts[j] - 1) * D)", "len=(ts[0] - 1) * D), B_(\"LD\", dst=SP_VC(j), src=EX_VC(b), len=(ts[j] - 1) * D)"),
]
def check(label, old, new, ragged=True):
if SRC.count(old) != 1: return label, "BAD ANCHOR (%d occurrences)" % SRC.count(old)
d = tempfile.mkdtemp(prefix="mut9_"); shutil.copytree(os.path.join(ROOT, "model"), os.path.join(d, "model")); open(os.path.join(d, "model", "batch.py"), "w").write(SRC.replace(old, new))
r = subprocess.run([sys.executable, "-c", "import sys; sys.path.insert(0,'model'); import batch; p=batch.check_batching(ragged=%s); sys.exit(1 if p else 0)" % ragged], cwd=d, capture_output=True, text=True, timeout=600); shutil.rmtree(d, ignore_errors=True)
return label, r.returncode != 0
print("NOTE: this script deliberately breaks copies of the batched program builder. 'caught' is EXPECTED: it shows batch invariance notices the mistake.")
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
rows = list(pool.map(lambda m: check(*m), M)); rows2 = list(pool.map(lambda m: check(m[0], m[1], m[2], False), M))
n = 0; n2 = 0
print(f"{'':75s} {'full check':>11s} {'no ragged batch':>16s}")
for (label, r), (_, r2) in zip(rows, rows2):
print(f"{label:75s} {'caught' if r is True else 'NOT CAUGHT':>11s} {'caught' if r2 is True else 'NOT CAUGHT':>16s}"); n += r is True; n2 += r2 is True
print(f"\nbatched builder: {n} of {len(M)} broken programs caught by the full check; {n2} of {len(M)} if the ragged batch (different lengths in one step) is left out")
sys.exit(0 if n == len(M) else 1)
Output (cloud sandbox -- live-executed)
NOTE: this script deliberately breaks copies of the batched program builder. 'caught' is EXPECTED: it shows batch invariance notices the mistake.
full check no ragged batch
projection runs on one row (M = 1) instead of the whole batch caught caught
projection results are stored with row stride 4, not d caught caught
every new key is taken from row 0 of the projection caught caught
every new value is taken from row 0 of the projection caught caught
every sequence's cache lives in batch slot 0 caught caught
the value caches of all sequences overlap caught caught
sequence j attends with the query of row 0 caught caught
hidden states are all loaded into slot 0 caught caught
every sequence reads the hidden state of sequence 0 caught caught
outputs are all stored in sequence 0's output slot caught caught
the new key is appended at the first sequence's position caught NOT CAUGHT
the value cache is loaded from the key cache caught caught
the new value is appended to the key cache in external memory caught caught
attention length of every sequence is the first sequence's caught NOT CAUGHT
the q requantization covers only one row caught caught
a sequence's cache is loaded with the first sequence's length caught NOT CAUGHT
batched builder: 16 of 16 broken programs caught by the full check; 13 of 16 if the ragged batch (different lengths in one step) is left out
All 16 broken builders are caught, and the second column shows what each part of the test is worth:
- Without the ragged batch (sequences of different lengths in one step) three mutants survive: "attention length of every sequence is the first sequence's", "the new key is appended at the first sequence's position" and "a sequence's cache is loaded with the first sequence's length". In a batch where every sequence has the same length, the first sequence's length is everybody's length, so these bugs are invisible. Ragged batches are what real servers have.
- An earlier version of the test checked the integer reference only for a batch of one, i.e. only for sequence 0. The mutant "every sequence reads the hidden state of sequence 0" then passed in lock-step batches: the batched run and the single run shared the bug, so they agreed, and only sequence 0 was compared with the reference. Checking every sequence against the reference closed the hole (it is the version in
model/batch.pyabove). Batch invariance is a strong oracle, but like any comparison between two things built from the same code it is blind to a bug that both contain.
Common mistakes¶
- Assuming batching helps attention. It shares the weights, not the caches. Quoting a batching speed-up without separating the two overstates it.
- Choosing the batch by arithmetic and forgetting capacity. The batch is limited by memory (Example A), long before arithmetic is the limit.
- Testing batches whose members are all alike. A bug that mixes up the first sequence with the others is invisible if all sequences have the same length. The ragged batch is the test.
- Trusting a comparison between two things built from the same code. The batched and single runs shared a bug in the first version of the test; only the exact reference for every sequence caught it.
- Reading roofline "achieved" as a verdict on the arithmetic. The chip reaches a third of the memory ceiling because instructions do not overlap (Chapter 7), not because the array is slow.
- Quoting tokens per second without the context length. The same chip gives 3,725 tokens/s at 2,048 tokens and 233 at 32,768 in Example A.
What this chapter does and does not establish¶
- Measured on the chip: batched decode steps for 1 to 4 sequences (L = 16) and a ragged batch, with exact agreement with the reference simulator; cycles per token, words moved, and MACs per word.
- Derived, not built: the estimate of 1.4x to 1.55x from overlapping loads with compute; the limit of 2.35 MACs per word (and 1 for long contexts); the tokens-per-second table for a 7B-class model on an assumed 1 TB/s chip.
- Not covered: more than four sequences (the scratchpad layout holds four), contexts above 16 tokens in the batched layout, compute-bound phases (prefill, with long matrix products, is the opposite case and sits near the ridge), and real DRAM behaviour.
Chapter summary¶
A decode step uses each weight once, so its arithmetic intensity is about 1 MAC per byte against a ridge point of 16: memory-bound, on this chip and any other. Batching stacks requests into one matrix product so the weights are read once, which raised the speed per token by 1.58x at a batch of 4; it cannot touch attention, whose traffic is each request's own KV cache. The cache, not the weights, is what bounds throughput at large batch and long context. A batched program is tested by the invariance of each request to its company, with a ragged batch, plus an exact reference for each sequence alone.
Self-check questions¶
- What is GA-2's ridge point, and which side of it is a decode step on?
- The words moved per step are
768 + 544 B. Where do 768 and 544 come from? - Why does the batch's MAC-per-word figure approach 2.35 but not 16?
- Using the table, at what batch size does the KV cache equal the weights for int8 weights, and what does that say about adding more batch beyond it?
- Why did the first version of the batch test miss "every sequence reads the hidden state of sequence 0", and what fixed it?
- In Running example A, chip B has 80 GB and the model's weights are 7 GB. At a 8,192-token context a request needs 2.0 GiB of cache. Verify the printed maximum batch of 33.
- In Running example B, how many sequence-slots does static batching spend on a group of four requests of lengths 14, 3, 9 and 5? How many are useful?
- The chapter says batching cannot lift attention across the ridge point. Using the 2.35 limit, what is the best fraction of the array's peak that attention-dominated decoding could ever reach on GA-2 at L = 16?
Exercises¶
- Your own chip. Edit
CHIPSinch09_example_a.pyto a memory and bandwidth of your choice (say 16 GB and 0.3 TB/s) and a context of 4,096; predict the maximum batch and the batch-1 speed before running. - Int4 weights. Change
weight_bytesto 3.5e9 (int4) and rerun. What happens to the speed at batch 1, and at the limit? Which table column does not change, and why? - A different mix. Change
LENSinch09_example_b.pyso that all eight requests want 8 tokens. How do static and continuous compare now? What does that say about when continuous batching matters? - Capacity 8. Set
CAP = 8in Example B (the program supports up to four sequences, but the cycle model does not mind for this experiment; note that this goes beyond the layout ofmodel/batch.py). Is the gain larger or smaller than with four slots? - A new mutant. Add a mutant to
tools/mut_ch09.pythat gives every sequence the second sequence's cache base instead of its own. Predict which part of the test (lock-step batches, ragged batch, exact reference) catches it.