When a client asked why a 7-billion-parameter model that ran fine on a 24 GB GPU suddenly started throwing CUDA out of memory after a few dozen generations, the answer wasn’t the model size. It was the key-value (KV) cache that the transformer builds for every token. The cache lives in VRAM, grows with context length, and, if you look closely, it behaves like a heap that fragments over time.

The KV cache in a nutshell

Every self-attention layer stores two matrices per token: the keys and the values. During generation the model appends a new row for each incoming token, then re-uses the existing rows for the next attention lookup. In PyTorch or TensorFlow this is usually a single large tensor of shape [seq_len, heads, head_dim] that gets sliced as seq_len grows.

The problem appears when you try to truncate or slide the window. Some pipelines keep a fixed-size buffer and copy the tail of the tensor to the front when the limit is reached. That copy operation leaves a gap in VRAM because the original allocation is never released, the allocator thinks the memory is still in use. Over many generations those gaps add up and the effective usable VRAM shrinks, even though the logical sequence length stays the same.

Why fragmentation matters more than bandwidth

A few weeks back I wrote an article about VRAM bandwidth, which is still relevant: the faster you can feed data to the GPU, the more tokens per second you get. KV cache fragmentation is orthogonal, it reduces the capacity you can address. Even a GPU with plenty of bandwidth will choke if the allocator can’t find a contiguous block for the next slice of the cache.

In practice I’ve seen a 12 GB RTX 3060 fill up after about 2 k tokens, while the same model on a 24 GB card survives 4 k tokens. The difference isn’t the raw memory; it’s the amount of fragmentation left after each truncation.

PagedAttention: the idea

PagedAttention was introduced as a way to treat the KV cache like a virtual memory system. Instead of allocating one gigantic tensor, the implementation breaks the cache into fixed-size pages (e.g., 256 tokens per page). When the context exceeds the current pages, a new page is allocated; when you slide the window, the oldest page is simply released back to the allocator.

The key insight is that releasing whole pages avoids the tiny gaps left by copying rows. The GPU’s memory manager can now reuse those pages for future allocations, keeping the overall VRAM footprint stable.

A minimal PyTorch sketch

Below is a stripped-down example that demonstrates page-based KV management. It isn’t a full transformer, just the part that stores and retrieves keys/values.

import torch

PAGE_SIZE = 256  # tokens per page
HEADS = 12
HEAD_DIM = 64

def allocate_page():
    # Allocate a fresh page on the GPU
    return torch.empty(PAGE_SIZE, HEADS, HEAD_DIM, device='cuda')

class PagedKVCache:
    def __init__(self):
        self.pages = []          # list of (keys, values) tuples
        self.length = 0          # total tokens stored
    def append(self, new_keys, new_vals):
        # new_keys/vals shape: [t, HEADS, HEAD_DIM]
        t = new_keys.size(0)
        while t > 0:
            if not self.pages or self.pages[-1][0].size(0) - self.pages[-1][0].size(0) == PAGE_SIZE:
                # current page full, allocate a new one
                self.pages.append((allocate_page(), allocate_page()))
            page_keys, page_vals = self.pages[-1]
            offset = self.length % PAGE_SIZE
            space = PAGE_SIZE - offset
            take = min(t, space)
            page_keys[offset:offset+take] = new_keys[:take]
            page_vals[offset:offset+take] = new_vals[:take]
            new_keys = new_keys[take:]
            new_vals = new_vals[take:]
            self.length += take
            t -= take
    def truncate(self, keep):
        # keep most recent `keep` tokens, drop the rest
        if keep >= self.length:
            return
        drop = self.length - keep
        # drop whole pages first
        while drop >= PAGE_SIZE and self.pages:
            self.pages.pop(0)
            drop -= PAGE_SIZE
            self.length -= PAGE_SIZE
        # if anything left, just adjust the first page's offset (no copy needed)
        if drop > 0 and self.pages:
            # we can't shrink a tensor, so we just move the pointer logically
            self.length -= drop
    def get_all(self):
        # concatenate pages for attention lookup
        keys = torch.cat([p[0] for p in self.pages], dim=0)[:self.length]
        vals = torch.cat([p[1] for p in self.pages], dim=0)[:self.length]
        return keys, vals
  • allocate_page creates a fresh GPU tensor of the fixed page size.
  • append writes new KV rows into the current page, allocating a new page when the old one fills.
  • truncate drops whole pages first, avoiding any memory copy; the remaining drop is handled by adjusting the logical length.
  • get_all stitches the active pages together for the attention operation.

The code is intentionally simple; a production implementation would reuse pages from a pool and handle multi-GPU sharding, but the principle stays the same.

Practical pitfalls

  1. Page size choice, Too small and you end up with many allocations, hurting launch overhead. Too large and you re-introduce fragmentation because the tail of a page may stay allocated forever. In my tests 256-512 tokens per page hits a sweet spot on RTX 4090.
  2. Allocator behavior, Some older CUDA versions fragment even whole-page allocations when the driver compacts memory lazily. Upgrading the driver and using torch.cuda.memory_snapshot to inspect free blocks can reveal hidden fragmentation.
  3. Mixed precision, If you store KV in float16 but later switch to bfloat16 for a different model, the page size in bytes changes. Keep the page size expressed in tokens, not bytes, to stay consistent.
  4. Debugging, The most common symptom is an unexpected CUDA out of memory after a long generation run. Print torch.cuda.memory_allocated() before and after a truncation; if the number doesn’t drop, you’re still holding onto pages.

When not to use PagedAttention

If your use-case never exceeds a modest context window (say <1 k tokens) the overhead of page management isn’t worth it. Also, if you’re running on CPU-only inference, the fragmentation issue disappears because the OS handles virtual memory more gracefully.

How this fits into a larger pipeline

PagedAttention is just one piece of a memory-efficient inference stack. Pair it with:

  • FlashAttention or xFormers for reduced compute per token.
  • Quantization (e.g., 4-bit GPTQ) to shrink each KV entry.
  • Sliding-window attention that discards older tokens at the algorithmic level, not just the cache.

When you combine these, you can push a 13 B model to 32 k token contexts on a single 24 GB GPU, something that would have been impossible a year ago.

FAQ

Why does the KV cache grow even when I limit the context length?

The cache grows because each new token adds a row to every layer’s KV tensor. If you only slice the tensor without releasing the underlying memory, the allocator still sees the original allocation as in use, leading to hidden growth.

Can I use PagedAttention with existing Hugging Face models?

Yes. Hugging Face’s generate API accepts a custom past_key_values object. You can wrap the model’s KV tensors in a PagedKVCache instance and feed the concatenated view back into the model each step.

Does PagedAttention affect inference speed?

The extra concatenation step adds a few microseconds per token, which is negligible compared to the attention compute. In most cases you actually see a speedup because the GPU no longer stalls waiting for a large contiguous allocation.

Wrapping up

If you’re tired of watching your GPU run out of memory after a few hundred generations, stop treating the KV cache as a monolithic block. Break it into pages, release whole pages when you slide the window, and you’ll keep the VRAM footprint steady. It’s a small architectural change that pays off big when you push context windows to the limits.