GPT-2 KV cache from scratch: 1000x fewer FLOPs but only 6x latency gain, measured
atronos_kronios · reddit · 2026-10-18
The author built GPT-2 small (124M) with a KV cache in PyTorch and compared theoretical savings against measured latency and memory, publishing code and a detailed write-up.
- Correctness: max logit difference vs HuggingFace 6.1e-5; cached vs uncached greedy generation matched over 200 tokens. Notable failure case: reversing cache concatenation order passes single-token attention tests but fails when decoding multiple new tokens together.
- Latency: at context length 1,000, FLOPs drop 1000x but per-token time only improves 6x (29.69ms → 4.95ms) on an A40, fp32; generating 1,016 tokens takes 14.66s vs 5.03s (2.9x). Kernel-launch/CPU overhead is hypothesized for the remaining floor.
- Memory: formula says 72 KiB/token. An initial constant 8,519,688-byte offset was traced to a lazily allocated cuBLAS workspace plus an 8-byte tensor; after warmup and subtracting live tensors, measured minus formula was 0 B at all 16 checkpoints.
Educational implementation, not a new algorithm; the author requests methodology feedback.
More from Infra
- FT: 47 major US AI data center projects delayed or canceled amid $10T buildout — TansuYegen · 2026-10-18
- Google rewrites SRE book for the AI era, announces second edition 10 years later — mgill25 · 2026-10-18
- Agens Volundr 32B v1 Kills Repetition Loops; MIT Engine Fits 262K Context in 16GB — ComfortableKindly507 · 2026-10-18
- Korea Chip Stocks Fall 7.57% as Investors Price the Gap Between AI Demand and Financing Capacity — shashib · 2026-10-18
- Three receipts in 14 days: Anthropic's $15B SpaceX compute deal fuels the 'last window' thesis — elonmusk · 2026-10-18
- Cloudflare Now Scores Verified Bots on robots.txt Compliance and 429 Response — shashib · 2026-10-18