从零实现 GPT-2 KV cache：正确性、延迟与内存实测

I built GPT-2 small and a KV cache in PyTorch as a learning project, then compared the mathematical savings with measured latency and memory. I have published the implementation and a detailed write-up, and would appreciate feedback on the methodology.

Setup: pretrained GPT-2 small (124M), NVIDIA A40, fp32 with TF32 off, batch size 1, greedy decoding. Timings are bracketed by torch.cuda.synchronize(), after warmup, using the median of five runs per context length.

Correctness: maximum absolute logit difference versus Hugging Face was about 6.1e-5; cached and uncached greedy generation matched over 200 tokens. One useful failure case: reversing the cache concatenation order still passed a single-token attention test, but failed when multiple new tokens were decoded together.

Latency: at context length 1,000, the FLOP estimate is about 1,000x smaller with caching, but the per-token time goes from 29.69 ms to 4.95 ms (~6x). Generating 1,016 tokens takes 14.66 s uncached versus 5.03 s cached (~2.9x). The cached curve is nearly flat in this implementation. CPU/kernel-launch overhead is a hypothesis, not a conclusion supported by a profiler yet.

Memory: the formula is 72 KiB/token. My initial requested-byte measurement had a constant 8,519,688-byte offset: a lazily allocated 8,519,680-byte cuBLAS workspace plus an 8-byte next-token tensor. Warming up before the baseline and subtracting live non-cache tensors brought requested-minus-formula to 0 B at all 16 checkpoints. The allocator's block-size overhead remains visible separately.

Code: https://github.com/angry-kratos/angry-kratos.github.io/tree/main/code/kv-cache

Write-up and figures: https://angry-kratos.github.io/posts/2026-10-kv-cache/

This is an educational implementation, not a new algorithm or a universal speedup claim. What additional correctness cases or profiling controls would you add before investigating the remaining latency floor?

Disclosure: I am the author. The code and article are freely accessible; AI assistance was used for editing and tooling. I also shared a short introduction in this month's self-promotion thread.

submitted by /u/atronos_kronios
[link] [留言]