Implementing GPT-2 KV Cache from Scratch: Correctness, Latency, and Memory Benchmarks

A Teaching Project Reveals the Real Gains and Hidden Costs of KV Cache

KV Cache is a core technique for accelerating autoregressive inference, but how big is the gap between its theoretical benefits and actual performance? A developer implemented GPT-2 small's KV Cache from scratch in PyTorch and measured three sets of data: correctness, latency, and memory. The conclusions are unsurprising but honest enough: theoretical FLOP savings of 1000x translate to only a 6x speedup in measured single-token latency; get the cache concatenation order wrong, and single-token tests pass while multi-token decoding breaks down.

The project's code and full analysis are publicly available on GitHub and the author's blog. The author explicitly states this is a teaching implementation, not a new algorithm, and makes no claims of general-purpose acceleration. This restrained framing actually makes the experimental data more valuable as a reference.

Experimental Setup: Why These Numbers Are Worth Examining

The experimental environment was chosen cleanly: pretrained GPT-2 small (124M parameters), NVIDIA A40, fp32 precision with TF32 disabled, batch size 1, greedy decoding. Timing was wrapped with torch.cuda.synchronize(), with the median of five runs taken after warmup.

This setup eliminates interference from mixed precision and batching, focusing KV Cache gains on the single-sequence inference scenario. For developers wanting to understand caching mechanisms, this is a more controllable starting point than running benchmarks.

KV Cache is an inference optimization technique in autoregressive generation: as the model generates token by token, each step only needs to compute the Query for the new token, while Key and Value can reuse results from all previous tokens. Without caching, each new token requires recomputing K and V for the entire sequence, causing computation to grow quadratically with sequence length; with caching, each step only computes K and V for the new token and appends them to the cache, reducing computation to linear growth. This is the source of the massive theoretical FLOP savings. However, caching also introduces memory usage that grows linearly with sequence length, and actual latency is constrained by factors like memory bandwidth, kernel launch overhead, and CPU scheduling.

Correctness: An Easily Overlooked Failure Case

The author first verified numerical consistency: compared to the Hugging Face implementation, the maximum logit absolute error is approximately 6.1e-5; in 200-token greedy generation, results with and without caching match exactly.

The truly valuable finding came from a failure case. When the cache concatenation order was reversed, the single-token attention test still passed, but decoding multiple new tokens together failed. This shows that single-step tests have limited coverage, and cache logic validation must include multi-token decoding scenarios. The author solicited additional correctness test cases from the community in discussion—this failure mode itself is a good reminder.

Latency: 1000x FLOP Savings Buys Only 6x Speedup

At a context length of 1000, the estimated FLOPs with caching are about 1000x smaller, but single-token latency drops from 29.69ms to 4.95ms—only about 6x improvement. Total time to generate 1016 tokens drops from 14.66 seconds to 5.03 seconds, about 2.9x.

The author attributes the gap to CPU overhead and kernel launch overhead, but explicitly states this is a hypothesis without profiler data to support it. In this implementation, the latency curve with caching is nearly flat, meaning the bottleneck is no longer in matrix multiplication computation.

This is also the watershed between teaching implementations and engineering optimization: theoretical FLOPs are just an upper bound; scheduling, memory access, and launch overhead set the floor for actual latency. The author asks the community at the end which profiling controls to add—this question itself points to the next direction.

FLOP (floating-point operations) is often used to measure computational workload, but it only reflects the number of arithmetic operations and does not directly equal latency. On actual hardware, latency also depends on memory access, kernel launch overhead, CPU-GPU synchronization, and more. Especially in small-batch autoregressive decoding, matrix multiplications are small, GPU utilization is low, and kernel launch and scheduling overhead dominate, so FLOP savings often cannot translate linearly into speedups. This is the fundamental reason behind "1000x FLOP savings buys only 6x latency reduction" in the article.

Memory: Where Does the 8.5MB Offset Beyond the Formula Come From

The theoretical memory formula for KV Cache is 72 KiB/token. But the author's initial measurements of requested bytes showed a fixed offset beyond the formula: 8,519,688 bytes. Breaking this down revealed that 8,519,680 bytes were cuBLAS lazily allocated workspace, with the remaining 8 bytes being the next-token tensor.

After warming up before baseline and subtracting live non-cache tensors, the "requested minus formula" difference was zero across 16 checkpoints. The allocator's own block size overhead was visible separately. The value of this experiment lies in methodology: measuring memory cannot rely solely on requested bytes—lazy allocation and workspace cause systematic overestimation.

cuBLAS is NVIDIA's GPU linear algebra library, which many deep learning frameworks call under the hood for matrix multiplication. cuBLAS lazily allocates internal workspace on first invocation of certain operations, used to store intermediate results or data required by optimization algorithms. This memory does not belong to user-explicitly-requested tensors but counts toward process memory usage, causing measured "requested bytes" to exceed the theoretical formula by a fixed offset. When analyzing memory, you typically need to warm up to trigger lazy allocation, then subtract non-cache tensor usage, to get clean cache overhead.

Takeaways for Practitioners

This project proposes no new algorithm, but provides three sets of reproducible measurements and one failure mode worth remembering. For teams doing inference optimization, at least three things can be borrowed directly: correctness tests must cover multi-token decoding; leave budget for scheduling overhead between theoretical FLOPs and measured latency; memory measurements require warmup first, then stripping non-cache tensors.

The author also candidly used AI-assisted editing and tools. The code and analysis are publicly accessible, suitable as a starting point for understanding KV Cache internals rather than as a performance benchmark.