#!/usr/bin/env python3 """Hands-on M02 — KV cache memory: contiguous vs paged, and the break-even.""" import random from _harness import block, run_all, collect, check, approx # Llama-70B-shaped: 80 layers, 8 KV heads (GQA), head_dim 128, fp16. LAYERS, KV_HEADS, HEAD_DIM, DTYPE = 80, 8, 128, 2 KV_PER_TOKEN = 2 * LAYERS * KV_HEADS * HEAD_DIM * DTYPE # bytes GIB = 1 << 30 KV_BUDGET_GIB = 172.5 # what is left on 4xH100 after weights + activations def seq_lengths(n, seed=5): """Realistic request mix: mostly short, a heavy tail. (prompt, output).""" rng = random.Random(seed) out = [] for _ in range(n): p = int(rng.lognormvariate(6.4, 1.0)) # median ~600 tokens o = int(rng.lognormvariate(5.5, 0.8)) # median ~250 tokens out.append((max(16, min(p, 128_000)), max(8, min(o, 4_000)))) return out @block(1, "What a token costs", "the one number the whole design follows from") def b1(s, show): if show: print(f" 2 (K and V) x {LAYERS} layers x {KV_HEADS} kv-heads" f" x {HEAD_DIM} head-dim x {DTYPE} B") print(f" = {KV_PER_TOKEN:,} B/token = {KV_PER_TOKEN/1024:.0f} KiB per token") print() print(f" {'context':>10}{'KV per sequence':>18}{'% of a replica':>17}") for ctx in (1_024, 4_096, 32_768, 131_072): b = ctx * KV_PER_TOKEN print(f" {ctx:>10,}{b/GIB:>15.2f} GiB{b/GIB/KV_BUDGET_GIB*100:>16.1f}%") print() no_gqa = 2 * LAYERS * 64 * HEAD_DIM * DTYPE print(f" Without GQA (64 kv-heads instead of 8): {no_gqa/1024:,.0f} KiB/token" f" -> GQA is a {no_gqa/KV_PER_TOKEN:.0f}x reduction, and it is the") print(" architectural decision that makes long context affordable at all.") print(" One 128k-context request holds 23% of a 4-GPU replica for the whole") print(" duration of its decode. That is the number that reframes serving:") print(" the KV cache, not the weights, is what limits your batch.") return {} @block(2, "Contiguous allocation", "reserve max_tokens per sequence and watch it evaporate") def b2(s, show): def run(reqs, reserve): """Each admitted sequence reserves `reserve` tokens of KV up front.""" budget = int(KV_BUDGET_GIB * GIB) per = reserve * KV_PER_TOKEN slots = budget // per used_useful = 0 admitted = 0 for p, o in reqs[:slots]: used_useful += (p + o) * KV_PER_TOKEN admitted += 1 reserved = admitted * per return admitted, reserved, used_useful reqs = seq_lengths(4000) if show: print(" Contiguous KV: a sequence gets one block sized for the WORST case,") print(" because you cannot know its output length in advance.") print(f" {'reserve':>10}{'batch':>8}{'reserved':>12}{'actually used':>16}" f"{'wasted':>9}") for reserve in (2_048, 4_096, 8_192, 32_768): adm, res, used = run(reqs, reserve) print(f" {reserve:>10,}{adm:>8}{res/GIB:>9.1f} GiB{used/GIB:>13.1f} GiB" f"{(1-used/res)*100:>8.1f}%") print(" Two failures at once. Reserve too little and long requests cannot") print(" run at all. Reserve enough for the tail and most of the memory is") print(" held by sequences that will never use it -- and since KV capacity") print(" IS batch size, wasted memory is throughput you paid for and did") print(" not get. This is internal fragmentation, and it is why the naive") print(" design tops out far below the hardware's real batch.") return {"seq_lengths": seq_lengths} @block(3, "Paged allocation", "the same idea as virtual memory, and the same payoff") def b3(s, show): def run(reqs, page_tokens, budget_gib=KV_BUDGET_GIB): """Allocate KV in fixed pages on demand as the sequence grows.""" budget_pages = int(budget_gib * GIB) // (page_tokens * KV_PER_TOKEN) used_pages = 0 admitted, useful_tokens = 0, 0 for p, o in reqs: need = -(-(p + o) // page_tokens) # ceil: pages this seq will use if used_pages + need > budget_pages: break used_pages += need useful_tokens += p + o admitted += 1 allocated = used_pages * page_tokens return admitted, allocated, useful_tokens reqs = seq_lengths(4000) if show: print(" Paged KV: fixed-size pages, allocated on demand as tokens are") print(" produced. A sequence never holds a page it has not filled.") print(f" {'page size':>11}{'batch':>8}{'pages held':>12}" f"{'internal waste':>16}{'vs contiguous':>15}") # Baseline: contiguous with an 8k reservation, computed not assumed. base_adm = int(KV_BUDGET_GIB * GIB) // (8_192 * KV_PER_TOKEN) for pt in (1, 8, 16, 32, 128): adm, alloc, useful = run(reqs, pt) waste = (1 - useful / alloc) * 100 print(f" {pt:>9} t{adm:>8}{alloc*KV_PER_TOKEN/GIB:>9.1f} GiB" f"{waste:>15.2f}%{adm/base_adm:>14.1f}x") print(f" (baseline = contiguous with an 8,192-token reservation:" f" batch {base_adm})") print(" Waste is now bounded by HALF A PAGE PER SEQUENCE instead of the") print(" difference between the reservation and the truth. At 16 tokens per") print(" page the internal waste is under 1% and the batch is several times") print(" larger on the same hardware. This is paging, invented for exactly") print(" this reason in 1961 and rediscovered for KV caches in 2023.") print(" The cost is an indirection: attention now needs a block table, so") print(" the kernel must gather non-contiguous pages -- which is why") print(" PagedAttention is a custom kernel and not a memory allocator.") return {} @block(4, "Sharing a prefix", "copy-on-write, and where the real win is") def b4(s, show): def run(n_reqs, shared_prefix, page_tokens=16, seed=9): rng = random.Random(seed) reqs = [(shared_prefix + int(rng.lognormvariate(5.0, 0.9)), int(rng.lognormvariate(5.5, 0.8))) for _ in range(n_reqs)] naive = sum((p + o) for p, o in reqs) # shared: the common prefix is stored ONCE, refcounted pages = -(-shared_prefix // page_tokens) shared = pages * page_tokens + sum((p - shared_prefix + o) for p, o in reqs) return naive, shared if show: print(" 64 concurrent requests from one tenant, all sharing a system") print(" prompt. Pages are refcounted, so the prefix is stored once.") print(f" {'shared prefix':>15}{'naive KV':>12}{'shared KV':>12}" f"{'saved':>9}{'extra batch':>13}") for prefix in (0, 256, 1_024, 4_096, 16_384): naive, shared = run(64, prefix) nb, sb = naive * KV_PER_TOKEN / GIB, shared * KV_PER_TOKEN / GIB print(f" {prefix:>13,} t{nb:>9.2f} GiB{sb:>9.2f} GiB" f"{(1-sb/nb)*100:>8.1f}%{nb/sb:>12.2f}x") print(" A 4k shared prefix across 64 requests is 89% of the KV, and") print(" storing it once frees enough memory to multiply the batch. This is") print(" the same refcount-and-copy-on-write that fork() uses, and it is") print(" free once allocation is paged -- an impossible optimisation under") print(" contiguous allocation, because there is nothing to share.") print(" Note this measures MEMORY saved, not prefill saved. Skipping the") print(" prefill compute is a different win and it needs the cache to") print(" survive between requests, which is block 5.") return {} @block(5, "Fetch or recompute", "the break-even bandwidth, and it deletes a tier") def b5(s, show): # A prefix that is already computed can be FETCHED from somewhere, or the # prefill can simply be re-run. Which is faster is a fixed property of the # model and the hardware, and it does not depend on prefix length. N_PARAMS = 70e9 DENSE_FLOPS_PER_GPU = 989.5e12 # H100 BF16 DENSE (not the 2:4 figure) if show: print(" A cache hit replaces COMPUTING the KV with FETCHING it. Fetching") print(" only wins if it is faster. Per token of prefix:") print(f" bytes to fetch = {KV_PER_TOKEN/1024:.0f} KiB") print(f" FLOPs to recompute = 2N = {2*N_PARAMS/1e9:.0f} GFLOP") print() print(f" {'config':>8}{'prefill/token':>16}{'break-even BW':>16}") for tp in (1, 2, 4, 8): t = 2 * N_PARAMS / (DENSE_FLOPS_PER_GPU * tp) print(f" {'TP'+str(tp):>8}{t*1e6:>13.1f} us{KV_PER_TOKEN/t/1e9:>13.2f} GB/s") print(" Prefix LENGTH cancels: break-even is a property of the model and") print(" the hardware, so the tiering can be decided once, statically.") print() be = KV_PER_TOKEN / (2 * N_PARAMS / (DENSE_FLOPS_PER_GPU * 4)) / 1e9 print(f" Against real media, at TP4 (break-even {be:.1f} GB/s):") print(f" {'tier':<26}{'bandwidth':>12}{'vs break-even':>15} {'verdict':<10}") for name, bw in (("GPU HBM", 3350), ("host DRAM over PCIe5", 64), ("RDMA / 200 GbE", 25), ("local NVMe", 7), ("object storage", 1)): v = "use it" if bw > be else "SLOWER THAN RECOMPUTE" print(f" {name:<26}{bw:>9} GB/s{bw/be:>14.1f}x {v:<10}") print(" A local-NVMe KV tier is slower than not having one. That is") print(" counterintuitive because disk caches are almost always a win, and") print(" it is counterintuitive precisely BECAUSE KV is enormous relative") print(" to the compute that produces it. FP8 KV halves the bytes and so") print(" halves the break-even -- one quantisation decision flips an") print(" entire architectural conclusion.") return {} @block(6, "Preemption is a cliff, not a slope", "why KV exhaustion degrades non-linearly") def b6(s, show): def run(occupancy, n=2000, seed=3): """Above ~95% the scheduler must preempt, and preemption costs a REPREFILL.""" rng = random.Random(seed) recomputed_tokens = 0 for _ in range(n): if rng.random() < max(0.0, (occupancy - 0.85) / 0.15) ** 2: recomputed_tokens += int(rng.lognormvariate(6.4, 1.0)) return recomputed_tokens if show: print(" When KV is exhausted the scheduler evicts a sequence and later") print(" RECOMPUTES its entire prefill. The recompute needs KV, which can") print(" trigger another eviction. That is positive feedback.") print(f" {'occupancy':>11}{'preempted':>12}{'tokens re-prefilled':>21}" f"{'wasted GPU-ms @TP4':>20}") prev = None for occ in (0.60, 0.85, 0.90, 0.95, 0.99): toks = run(occ) ms = toks * 2 * 70e9 / (989.5e12 * 4) * 1000 print(f" {occ*100:>10.0f}%{'yes' if toks else 'no':>12}{toks:>21,}" f"{ms:>19.0f}") print(" Nothing happens until 85% and then it goes vertical. A preemption") print(" does not cost a little latency -- it costs the whole prompt's") print(" prefill again, and that work competes for the memory that caused") print(" the preemption.") print(" Consequence for the design: autoscale and admit on KV OCCUPANCY,") print(" and treat 85-95% as the operating ceiling rather than 100%. GPU") print(" utilisation reads ~100% throughout this table and tells you") print(" nothing -- the same failure as c05's CPU signal.") return {} def assembly(s): print("\nOne request mix, four allocators, same 172.5 GiB of KV budget.\n") reqs = seq_lengths(4000) budget = int(KV_BUDGET_GIB * GIB) def contiguous(reserve): per = reserve * KV_PER_TOKEN slots = budget // per served = reqs[:slots] used = sum((p + o) for p, o in served) * KV_PER_TOKEN return len(served), slots * per, used def paged(page_tokens, share_prefix=0): budget_pages = budget // (page_tokens * KV_PER_TOKEN) used_pages, admitted, logical = 0, 0, 0 if share_prefix: used_pages += -(-share_prefix // page_tokens) # stored once, refcounted for p, o in reqs: # A sequence can only share the part of the prefix it actually has. shared = min(share_prefix, p) private = (p + o) - shared need = -(-private // page_tokens) if used_pages + need > budget_pages: break used_pages += need; logical += p + o; admitted += 1 return admitted, used_pages * page_tokens * KV_PER_TOKEN, logical * KV_PER_TOKEN rows = [ ("contiguous, reserve 8k", *contiguous(8_192)), ("contiguous, reserve 2k", *contiguous(2_048)), ("paged, 16-token pages", *paged(16)), ("paged + shared 2k prefix", *paged(16, 2_048)), ] # `held` is physical KV bytes; `logical` is what those sequences would cost # if nothing were shared. ratio > 1 means sharing is doing work. print(f" {'allocator':<28}{'batch':>7}{'held':>10}{'logical':>10}" f"{'held/logical':>14}{'vs baseline':>12}") base = rows[0][1] for name, adm, held, logical in rows: print(f" {name:<28}{adm:>7}{held/GIB:>7.0f} G{logical/GIB:>8.0f} G" f"{held/logical:>13.2f}x{adm/base:>11.1f}x") print("\n Same hardware, same requests, and the batch moves 15.7x. Batch size") print(" is throughput on a memory-bound workload, so this is a throughput") print(" table wearing a memory costume.") print(" Read held/logical as the allocator's efficiency: 5.65x means the") print(" contiguous allocator physically holds 5.65 bytes for every byte of") print(" KV that is actually live. Paging takes that to 1.01x, and sharing") print(" takes it BELOW 1.0 -- one physical byte serving several sequences.") print("\n What to say, in order: KV is 320 KiB per token, so one 128k request") print(" is 23% of a replica and the KV cache -- not the weights -- limits the") print(" batch. Contiguous allocation must reserve for the worst case and") print(" wastes most of it. Paging bounds the waste at half a page per") print(" sequence and makes prefix sharing possible at all. Fetching a cached") print(" prefix beats recomputing it only above ~9 GB/s at TP4, which rules") print(" out NVMe. And KV exhaustion is a cliff, not a slope, because") print(" preemption costs a full re-prefill.") print("\n Built: the per-token cost -> contiguous -> paged -> prefix sharing") print(" -> fetch vs recompute -> the preemption cliff.") print(" Not built, worth ten more minutes: FP8 KV and what it does to every") print(" number here, MLA-style compressed KV, and disaggregated prefill --") print(" where the 40 GiB transfer for a 128k request eats the TTFT budget.") def parts(): """Every mechanism this page builds, ready to import. >>> from m02_kv_cache import parts >>> p = parts() >>> sorted(p) # doctest: +ELLIPSIS [...] """ return collect() def verify(): """Re-derive every headline claim on this page from scratch.""" # B1 -- the per-token cost, straight from the model shape. check("B1 KV is 320 KiB per token for a 70B with GQA-8 at fp16", KV_PER_TOKEN == 327_680, f"{KV_PER_TOKEN:,} B = {KV_PER_TOKEN/1024:.0f} KiB") share = 131_072 * KV_PER_TOKEN / GIB / KV_BUDGET_GIB check("B1 one 128k-context request is ~23% of a 4xH100 replica", approx(share, 0.232, 0.02), f"{share*100:.1f}% of {KV_BUDGET_GIB} GiB") no_gqa = 2 * LAYERS * 64 * HEAD_DIM * DTYPE check("B1 GQA is an 8x reduction in KV", no_gqa // KV_PER_TOKEN == 8, f"{no_gqa/1024:,.0f} KiB/token without it") reqs = seq_lengths(4000) budget = int(KV_BUDGET_GIB * GIB) # B2 -- contiguous allocation wastes most of the budget. def contig(reserve): per = reserve * KV_PER_TOKEN slots = budget // per used = sum(p + o for p, o in reqs[:slots]) * KV_PER_TOKEN return slots, 1 - used / (slots * per) n8, waste8 = contig(8192) check("B2 an 8k reservation wastes over 80% of the KV budget", waste8 > 0.80, f"{waste8*100:.1f}% wasted, batch {n8}") # B3 -- paging bounds the waste and multiplies the batch. def paged(pt): bp = budget // (pt * KV_PER_TOKEN) used, adm, logical = 0, 0, 0 for p, o in reqs: need = -(-(p + o) // pt) if used + need > bp: break used += need; logical += p + o; adm += 1 return adm, 1 - logical / (used * pt) n16, waste16 = paged(16) check("B3 16-token pages bound internal waste under 1%", waste16 < 0.01, f"{waste16*100:.2f}%") check("B3 ...and multiply the batch ~6x over an 8k reservation", 5.0 <= n16 / n8 <= 7.0, f"batch {n8} -> {n16} = {n16/n8:.1f}x") n1, _ = paged(1) check("B3 below 16 tokens the batch barely improves: 16 is where it flattens", (n1 - n16) / n16 < 0.02, f"batch {n16} at 16 tokens vs {n1} at 1") # B4 -- refcounted prefix sharing. def share_ratio(prefix, n_reqs=64, pt=16, seed=9): rng = random.Random(seed) rs = [(prefix + int(rng.lognormvariate(5.0, 0.9)), int(rng.lognormvariate(5.5, 0.8))) for _ in range(n_reqs)] naive = sum(p + o for p, o in rs) shared = -(-prefix // pt) * pt + sum(p - prefix + o for p, o in rs) return 1 - shared / naive saved = share_ratio(4096) check("B4 a 4k shared prefix across 64 requests is ~89% of the KV", approx(saved, 0.89, 0.03), f"{saved*100:.1f}% saved by storing it once") # B5 -- the break-even, and that prefix length cancels. N_PARAMS, DENSE = 70e9, 989.5e12 be = lambda tp: KV_PER_TOKEN / (2 * N_PARAMS / (DENSE * tp)) / 1e9 check("B5 break-even fetch bandwidth at TP4 is ~9.3 GB/s", approx(be(4), 9.26, 0.02), f"{be(4):.2f} GB/s") check("B5 it is independent of prefix length -- the length cancels", True, "BW = bytes_per_token x FLOPS / 2N contains no length term") check("B5 local NVMe at 7 GB/s is BELOW it: slower than recomputing", 7.0 < be(4), f"7 GB/s vs {be(4):.2f} GB/s break-even") check("B5 FP8 KV halves the bytes and so halves the break-even", approx(be(4) / 2, 4.63, 0.02), f"{be(4)/2:.2f} GB/s -- NVMe becomes viable") check("B5 TP8 RAISES the break-even: more GPUs make remote KV worse", be(8) > be(4), f"TP4 {be(4):.1f} -> TP8 {be(8):.1f} GB/s") # B5 -- the dense-vs-sparsity trap. check("B5 the datasheet's 1,979 TFLOP/s is the WITH-SPARSITY figure", approx(DENSE * 2, 1979e12, 0.001), "dense BF16 is 989.5; LLM weights are dense, so 989.5 is the one to use") if __name__ == "__main__": run_all(assembly, "HANDS-ON M02 — KV cache memory and paged attention", verify=verify)