KV Cache Memory Management for Concurrent Requests

Hello! Welcome back to your course on preparing for the Anthropic software engineering interview.

In our last lesson, we designed a sophisticated dynamic batching strategy using chunked prefills. We established that this approach is excellent for balancing throughput and latency, but it comes at a cost: it creates a highly dynamic and fragmented memory landscape. Requests of varying lengths are constantly entering and leaving the batch, making traditional contiguous memory allocation for their KV caches incredibly inefficient.

This lesson tackles that exact problem head-on. Our learning outcome is to design a memory management strategy for the KV cache to handle concurrent, variable-length requests. We will explore how modern LLM serving systems, like vLLM, solve this by borrowing a classic, powerful idea from operating systems. Your background in distributed systems and OS-level concepts will be particularly useful here, as we'll be drawing direct parallels to virtual memory and paging.

1. The KV Cache Memory Problem

As we've discussed, the KV cache is essential for efficient decoding, but it's also a major memory consumer. For a large model, the cache for a single request can require gigabytes of VRAM. When you use dynamic batching to serve many concurrent requests, each with an unpredictable output length, naive memory management leads to massive waste.

Let's watch a short video that clearly identifies the three main sources of this memory inefficiency in older systems.

These issues combined mean that in traditional systems, only 20-40% of the allocated KV cache memory is actually used to store token states. The rest is wasted, severely limiting the number of requests that can be batched and thus capping throughput.

2. The Solution: Virtual Memory and Paging for the GPU

The problem of memory fragmentation is not new; operating systems solved it decades ago with virtual memory and paging. The core idea is to decouple a program's logical address space from the physical memory layout. Instead of allocating one large contiguous chunk, memory is divided into fixed-size "pages" (logical) which can be mapped to any available "frames" (physical).

Modern LLM serving systems apply this exact same principle to the KV cache.

This approach, called PagedAttention, revolves around a few key components:

• KV Blocks: The KV cache is partitioned into fixed-size blocks, analogous to memory pages. Each block can store the key and value vectors for a small, fixed number of tokens (e.g., 16).
• Physical Block Pool: The serving engine pre-allocates a large pool of physical KV blocks on the GPU.
• Block Tables: Each request has its own block table, analogous to a page table. This table maps the request's logical sequence of blocks to the physical blocks in the pool.

This design brilliantly solves the fragmentation problem:

• External fragmentation is eliminated because all blocks are the same size. Any free block can satisfy any allocation request.
• Internal fragmentation is minimized. Waste only occurs in the very last block of a sequence and is bounded by the block size, which is negligible compared to the total sequence length.
• Reservation waste is gone because blocks are allocated on-demand, just-in-time, as tokens are generated.

To get a clear, first-principles understanding of this mechanism, the following blog post is an excellent resource.

3. Handling Concurrency and Memory Sharing

The true power of PagedAttention becomes apparent when handling concurrent requests with complex decoding strategies, such as parallel sampling (generating multiple outputs from one prompt) or beam search. In these scenarios, different output sequences often share a common prefix (the prompt) and can even share intermediate generated tokens.

With contiguous memory allocation, sharing is difficult and often requires expensive memory copies. With PagedAttention, it's elegant and nearly free.

• Sharing via Block Tables: If multiple sequences share a prefix, their block tables simply point to the same physical blocks for that prefix. No data is duplicated.
• Reference Counting: Each physical block maintains a reference count. When a block is shared, its count is incremented. When a sequence finishes and releases the block, the count is decremented. The block is only returned to the free pool when its count reaches zero.
• Copy-on-Write (CoW): What if two sequences share a block, but one needs to append a new token to it? The system performs a copy-on-write. It allocates a new block, copies the contents of the shared block, and updates the block table of the modifying sequence to point to the new copy. The original block's reference count is decremented.

This video demonstrates this CoW mechanism in action.

The vLLM paper provides the definitive technical details on this.

4. Managing Memory Pressure: Preemption and Swapping

A critical question remains: what happens when the system is under heavy load and runs out of free physical blocks in the GPU pool? This is directly analogous to the problem you've solved in Rudi with load shedding and QoS.

The scheduler must preempt some requests to free up memory. Given the paged memory model, vLLM can employ sophisticated strategies:

• Eviction Policy: The scheduler typically uses a policy like "evict the newest arriving requests first" (a variation of LIFO) to be fair to requests that have been waiting longer. A whole sequence group (e.g., all beams in a beam search) is evicted at once.
• Recovery Strategy: Once a sequence is evicted, how is its state recovered when it's rescheduled?
  1. Swapping: The classic OS approach. The evicted KV blocks are copied from GPU VRAM to cheaper CPU RAM. When the sequence is rescheduled, the blocks are copied back. This is viable if the CPU-GPU interconnect (PCIe) bandwidth is high.
  2. Recomputation: The evicted KV cache is simply discarded. When the sequence is rescheduled, its state is recomputed from the original prompt and previously generated tokens. This is viable if GPU compute is fast relative to the PCIe transfer speed.

These mechanisms ensure the system can gracefully handle overload, prioritizing active requests while managing a queue of preempted ones.

5. Complementary Memory Optimization Techniques

PagedAttention is a brilliant memory management strategy. It's often combined with techniques that reduce the memory footprint of the KV cache itself.

The CTO of Mistral AI provides a concise overview of several such techniques.

These methods are complementary to PagedAttention. For instance, you can have a quantized KV cache that is managed using paging, combining both a smaller footprint and efficient dynamic allocation.

Conclusion

In this lesson, we designed a complete memory management strategy for the LLM KV cache to support concurrent, variable-length requests. We saw that by adopting the principles of virtual memory from operating systems, we can overcome the severe fragmentation issues of naive allocation.

Key Takeaways:

• The Problem: Contiguous memory allocation for dynamic, variable-length KV caches leads to massive internal and external fragmentation, wasting up to 80% of VRAM.
• The Solution is Paging: PagedAttention virtualizes the KV cache by partitioning it into fixed-size blocks. Per-request block tables map logical sequences to non-contiguous physical blocks.
• Benefits of Paging: This design eliminates external fragmentation, minimizes internal fragmentation, and allows for just-in-time memory allocation, achieving near-optimal memory utilization (>96%).
• Efficient Concurrency: Paging naturally enables complex memory sharing for concurrent decoding strategies like parallel sampling and beam search, using reference counting and copy-on-write to avoid expensive memory copies.
• Handling Overload: When memory pressure is high, the system can preempt requests by either swapping their KV blocks to CPU RAM or recomputing them later, ensuring graceful degradation.

Next Lesson Preview

We now have a scheduler that can dynamically batch requests and a memory manager that can efficiently handle the resulting state. The final piece of the serving puzzle is delivering the output. How do we get the generated tokens back to the user in a way that feels interactive and responsive?

In our next lesson, we will focus on the client-facing side: Design a system for streaming token-by-token responses, comparing WebSocket, Server-Sent Events (SSE), and chunked transfer.