KV Cache Implementation in nanoVLM

Hugging Face has implemented KV (Key-Value) Caching from scratch in nanoVLM, a concise PyTorch codebase for training Vision Language Models. This optimization resulted in a 38% speedup in generation by eliminating redundant computations during the autoregressive inference process.

The Computational Redundancy in Autoregressive Generation

Autoregressive language models generate text one token at a time. In a standard transformer implementation without caching, the model must process the entire sequence—including all previously generated tokens—to predict the next token.

Because transformers are internally parallel, each new token prediction requires a full forward pass through all layers. This creates a quadratic increase in memory and compute requirements relative to the sequence length. Specifically, the model recomputes the Key (K) and Value (V) tensors for all previous tokens at every step, even though those tokens and their corresponding projections have not changed.

How KV Caching Optimizes Inference

KV Caching mitigates this inefficiency by storing the computed keys and values for each layer after the initial prompt is processed. Instead of re-processing the entire sequence, the model follows this incremental workflow:

Cache Initial State: After the first pass, the computed $K$ and $V$ for each layer are cached.
Incremental Computation: During subsequent generation steps, the model only computes $K$ and $V$ for the newest token.
Cache Update: The new $K$ and $V$ are appended to the existing cache.
Attention Calculation: The Query ($Q$) for the current token is used in conjunction with the cached $K$ and $V$ to produce the output.

In practice, this cache is maintained as a per-layer dictionary containing "key" and "value" tensors with the shape (batch_size, num_heads, seq_len_cached, head_dim).

Technical Implementation in nanoVLM

The implementation in nanoVLM involves modifications across three primary components to transition from full-sequence re-computation to an incremental update system.

1. Attention Block Updates

In the LanguageModelGroupedAttention class, the forward function was modified to accept a block_kv_cache. If a cache exists (indicating the model is not in the prefill phase), the model computes $K_{new}$ and $V_{new}$ for the current token and concatenates them with the cached tensors. If no cache exists, it performs the initial computation for the prompt.

2. Layer-Wise Cache Tracking

The LanguageModel class now implements layer-wise cache tracking. It utilizes a start_pos argument to ensure that rotary positional encodings are correctly aligned with the current generation index, ensuring the model knows the absolute position of the newly generated token relative to the sequence.

3. Bifurcation of the Generation Loop

The generate() method in the VisionLanguageModel was split into two distinct phases:

Prefill Phase: The model encodes the full input prompt and constructs the initial KV cache for all layers.
Decode Phase: The model generates tokens sequentially, using the cached keys and values to avoid re-processing the prompt and previously generated tokens.

Summary of Architectural Changes

Module	Original Behaviour	New Behaviour
`LanguageModelGroupedAttention.forward`	Recomputes $Q$, $K$, $V$ on every step	Uses and updates KV cache
`LanguageModel.forward`	No memory of previous state	Tracks per-layer KV cache, handles `start_pos`
`VisionLanguageModel.generate`	One-phase generation loop	Split into prefill and decode phases

Trade-offs and Implications

KV caching reduces per-token inference complexity from quadratic to $O(\text{seq len})$, enabling faster inference and the ability to run large models on consumer hardware. However, this efficiency comes with a trade-off: it increases memory usage to store the cache and increases code complexity. Additionally, it can restrict certain inference schemes, such as beam search, which may require more complex cache management.