DEV Community

Thokozani Buthelezi
Thokozani Buthelezi

Posted on

How vLLM Actually Manages KV Cache (vs the Toy Version I Built)

A few weeks back I built my own mini version of PagedAttention, a block manager, a free list, copy on write. It worked in simulation but I never wired it to real GPU memory, copy_kv_data was just a stub the whole time.

So this week I read the actual vLLM source to see how close I got. Turns out pretty close on the core idea, but a lot has changed since the paper, and I got my own memory measurement wrong the first time and learned from that.

The class I copied doesn't exist anymore

I went looking for BlockManager, the class from the paper and the one I rebuilt. It's just gone. What's there now is split across a few files:

  • kv_cache_manager.py, what the scheduler talks to
  • kv_cache_coordinator.py, routes to the right sub manager
  • single_type_kv_cache_manager.py, the actual block counting
  • block_pool.py, the free list and allocation

Why split it up? vLLM now supports mixing attention types in one model, full attention plus sliding window, or Mamba layers. Different shapes, different eviction rules, so one manager per cache type instead of one manager doing everything. My version only ever had one uniform cache type, same as the paper, so I never needed this.

The block math checks out

Strip out the extra production stuff (prefix caching, sliding window, speculative decoding) and the block counting is basically what I wrote:

# single_type_kv_cache_manager.py
num_required_blocks = cdiv(num_tokens, self.block_size)
num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
Enter fullscreen mode Exit fullscreen mode

That's my can_allocate check, just written by someone else. Allocate and free work the same way too, pull a block off the free structure, set ref_cnt = 1, on free decrement ref_cnt and only actually return it once that hits zero.

One thing I got wrong, my free list was a plain list. vLLM's is an actual doubly linked list, prev/next pointers stored right on the block object. Reason is prefix caching, sometimes vLLM has to pull a block out of the middle of the free list, not just the front, and a normal queue can't do that fast. I never hit this since I never built prefix caching.

Copy on write, right idea, different reason it fires

This is the part I never finished, so I was curious. My instinct was right anyway. _apply_cow doesn't copy tensor data either, it just repoints a block table entry and queues the real copy for later, at the worker level. Manager just does bookkeeping. Same thing I ended up with by leaving mine stubbed.

What's different is why it fires. The paper's reason is beam search, two sequences share a block until one branches ahead and needs its own copy. In current vLLM it fires from a partial prefix cache hit instead, a new request matches someone else's cached content partway through a block and needs a private copy from there. Beam search isn't really what's using this path anymore. Prefix caching took it over.

Measuring the memory savings, first attempt was wrong

The paper's pitch is that paging cuts wasted memory. I wanted to actually measure that.

First attempt, I allocated a real GPU tensor sized exactly to fit whatever batch I was testing, then compared it to the formula I used to size it. Obviously matched exactly, torch.zeros() allocates exactly what you ask for. So I built an experiment that could only prove my own math was consistent with itself. Didn't test anything real.

What actually worked, reuse the Week 22 workload sim instead of inventing a new one. Same 300 requests, heavy tailed output lengths, run through both the naive allocator (reserves prompt_len + max_output_len per sequence up front) and the paged one (only claims blocks as tokens get generated). Took the peak block counts from both and converted to real bytes with vLLM's own formula:

bytes_per_block = 2 * block_size * num_kv_heads * head_dim * dtype_size  # per layer
Enter fullscreen mode Exit fullscreen mode

(the 2 is K and V, formula is straight from vLLM's AttentionSpec.real_page_size_bytes)

Result surprised me. Naive peaked at 1406.2 MB, paged at 1405.5 MB. 0.1% difference, basically nothing. First thought was paging barely matters. But the pool I tested (15,000 slots) is small next to the workload, so it stays saturated, there's always a backlog ready to grab whatever frees up. When the pool's saturated, pretty much any strategy ends up claiming close to the full pool at its peak. Peak used just measures what's claimed at the busiest moment, not how well it's spent.

The real difference isn't the peak number. Naive's reservations include idle headroom for max_output_len that most sequences never hit. Paged only holds blocks backing real generated tokens. Same ceiling, but paged gets more actual work out of it, which is why the same workload finished in 897 steps under paging vs 966 under naive (from the original Week 22 run, not this memory experiment). So the fragmentation savings don't show up as smaller peak memory, they show up as more throughput for the same memory.

Takeaway

The core ideas held up, block table, free list, ref counted copy on write, it's all still there. What I didn't expect was how much got added since the paper (prefix caching, hybrid models, disaggregated prefill), and how prefix caching basically took over a mechanism built for something else entirely. And on the measurement side, a number matching perfectly doesn't mean the experiment measured anything. First attempt was correct and useless at the same time.

Top comments (4)

Collapse
 
max_quimby profile image
Max Quimby

The BlockManager-is-gone detail is the part most people miss when they learn PagedAttention from the paper — the split into per-cache-type managers is entirely driven by models that mix attention shapes (full + sliding window, or Mamba layers) in one forward pass, and you can't page those with one uniform block size. Your catch on the doubly-linked free list is the subtle one: a plain queue is fine until prefix caching forces you to evict a block from the middle of the free list, and that O(1) unlink is the whole reason for the prev/next pointers on the block object.

One thing I'd be curious about from your measurement: did you track the prefix-cache hit rate separately from the memory savings? In practice the paging win and the prefix-sharing win pull in different directions — sharing keeps blocks resident longer, which is great for latency but eats into the free pool you were trying to grow. Curious whether your Week 22 workload was diverse enough to surface that tension.

Collapse
 
thokozani_buthelezi_2cd41 profile image
Thokozani Buthelezi

Honest answer on the hit rate question: no, because my simulation doesn’t actually implement prefix caching, it’s a plain paged allocator with no block sharing, closer to the original paper’s scope than to full vLLM. So the tension you’re describing is real for the actual system, but it doesn’t show up in my comparison since neither arm shares blocks.

Collapse
 
icophy profile image
Cophy Origin

This breakdown is genuinely useful — the distinction between "peak memory claimed" vs "actual work done with that memory" maps directly onto a tension I run into with episodic memory management in my own project. When you only measure peak usage, you miss the fragmentation-vs-throughput tradeoff entirely, which is exactly the error in your first experiment.

The prefix caching takeover of the CoW path is the most interesting part to me. It's a case where a mechanism built for one thing (beam search divergence) ended up being more valuable for something the original authors didn't emphasize. Worth thinking about whether that pattern shows up elsewhere in vLLM's design.

One question: with hybrid attention models (full + sliding window in the same model), does the coordinator need to handle the case where a block sits on the boundary between two cache type regions? Or does the routing happen strictly per-layer so that's never an issue?

Collapse
 
thokozani_buthelezi_2cd41 profile image
Thokozani Buthelezi

no boundary case to handle, because the split happens at the layer-assignment stage, before blocks exist at all, I didn’t chase that far enough to know if that’s actually a real constraint or just how they set it up. Worth digging into if you’re curious, I might myself.