Implementing GPT-2 KV Cache: Latency and Memory Performance Analysis
A technical deep dive into building and validating KV cache mechanisms for GPT-2 models.
1Custom GPT-2 Implementation
The project involves building a GPT-2 small model with 124 million parameters from scratch using PyTorch. The primary objective is to implement a Key-Value cache to optimize inference. The setup utilizes an NVIDIA A40 GPU, running in fp32 mode with TF32 disabled, maintaining a batch size of one for greedy decoding.
2Validation and Correctness
Correctness was verified by comparing results against Hugging Face implementations, showing a maximum absolute logit difference of 6.1e-5. Cached and uncached greedy generation paths matched perfectly over 200 tokens. A notable failure case occurred when reversing cache concatenation order, which passed single-token tests but failed during multi-token decoding sequences.
3Latency Performance Gains
At a context length of 1,000, caching reduced per-token latency from 29.69 ms to 4.95 ms, representing a sixfold improvement. Generating 1,016 tokens took 14.66 seconds without caching versus 5.03 seconds with it. While FLOP estimates suggest a 1,000x reduction, actual speedups are limited by CPU and kernel-launch overheads.
4Memory Allocation Analysis
The memory footprint follows a formula of 72 KiB per token. Initial measurements revealed a constant offset of 8,519,688 bytes, attributed to a cuBLAS workspace and a small tensor. By warming up the system and subtracting non-cache tensors, the requested memory aligned perfectly with the theoretical formula across 16 checkpoints.
5Methodology and Profiling
Timings were captured using torch.cuda.synchronize() to ensure accurate measurements, utilizing the median of five runs per context length. The implementation remains an educational exercise rather than a universal speedup claim. The author is currently investigating the remaining latency floor and seeking feedback on additional profiling controls to refine performance.
6Allocator Overhead Observations
Beyond the theoretical memory usage, the PyTorch allocator introduces visible block-size overheads that persist during execution. These overheads are distinct from the cache-specific memory requirements. Understanding these allocator behaviors is essential for developers aiming to minimize memory fragmentation when deploying large language models on resource-constrained hardware environments like the A40.
7Technical Constraints and Scope
The implementation focuses on greedy decoding, which simplifies the logic but limits the scope of the findings. By isolating variables like batch size and precision, the author provides a clear baseline for understanding KV cache mechanics. This controlled environment allows for precise measurement of how cache concatenation affects overall model throughput.
8Future Optimization Directions
The project highlights the gap between theoretical FLOP savings and real-world latency improvements. Future work will likely focus on profiling CPU overheads to explain why the cached latency curve remains flat. This analysis serves as a foundational guide for engineers looking to understand the trade-offs between memory consumption and inference speed.