KV Cache Offloading

KV cache offloading implements a hierarchical memory model, treating GPU High-Bandwidth Memory (HBM) as a fast cache and CPU RAM or NVMe storage as a larger, slower backing store. The system must manage:

• Hot vs. Cold Cache Segments: Frequently accessed recent tokens remain in GPU memory, while older, less-accessed context is evicted to CPU/NVMe.
• Eviction Policies: Algorithms like Least-Recently-Used (LRU) determine which cache blocks to move.
• Prefetching: Proactively loading upcoming cache blocks from slow storage based on access prediction to hide I/O latency.

The core trade-off of offloading is exchanging increased I/O latency for greater memory capacity. This is quantified by the access times and bandwidth of different storage tiers:

• GPU HBM: ~1-2 TB/s bandwidth, sub-microsecond latency. Ideal for active decoding.
• CPU RAM (DDR): ~50-100 GB/s bandwidth, ~100 nanosecond latency. First-level offload target.
• NVMe SSD: ~5-7 GB/s bandwidth, ~10-100 microsecond latency. For massive, rarely-accessed context. The system's performance becomes I/O-bound when the cost of transferring cache blocks dominates the compute time for attention.

Modern implementations heavily rely on Unified Virtual Memory architectures like CUDA UVM or AMD's hUMA. UVM provides a single, coherent virtual address space across CPU and GPU, which is essential for offloading because:

• It eliminates the need for explicit, synchronous cudaMemcpy operations.
• Enables on-demand paging: The GPU can trigger a page fault when accessing an offloaded cache block, which the driver handles by asynchronously fetching the data from host memory.
• Simplifies programmer oversight but requires careful management to avoid excessive page fault thrashing.

The primary motivation for offloading is to scale context length beyond GPU memory limits. This enables:

• Document-Level Processing: Analyzing or summarizing multi-hundred-page documents in a single sequence.
• Long Conversational Histories: Maintaining context over extended multi-session dialogues.
• Large Codebase Reasoning: Processing entire software repositories. Without offloading, the KV cache size grows linearly with sequence length (O(n)), making million-token contexts infeasible on standard GPUs. Offloading changes the constraint from fixed GPU memory to the larger, cheaper storage hierarchy.

To mitigate I/O latency, offloading systems use asynchronous and overlapping data transfers. This involves:

• Compute/I/O Overlap: While the GPU computes attention for the current batch, DMA engines concurrently fetch the next required cache blocks from host memory.
• Double Buffering: Using multiple buffers to hold cache data, allowing one buffer to be processed while the other is being filled.
• Non-Blocking Kernels: CUDA kernels are designed to be non-blocking, allowing other streams to execute while waiting for data. Effective overlap is critical to prevent the GPU from idling, which would nullify the benefits of increased cache capacity.

KV cache offloading is highly synergistic with PagedAttention, the memory management algorithm used by vLLM. PagedAttention organizes the cache into fixed-size blocks (pages), which is ideal for offloading because:

• Granular Management: Individual pages can be evicted or fetched independently, enabling fine-grained control over what resides in GPU memory.
• Reduced Fragmentation: Non-contiguous page mapping in virtual memory avoids the need for large, contiguous GPU memory allocations for the entire cache.
• Efficient Sharing: Pages can be shared between sequences in a batch (for shared prefix), and this sharing is preserved even when pages are offloaded, maintaining memory efficiency.

Unified Virtual Memory is a foundational hardware/software feature, such as NVIDIA's CUDA UVM, that enables KV cache offloading. It creates a single, coherent virtual address space spanning both GPU and CPU memory (and sometimes NVMe). This allows the inference engine to seamlessly access the KV cache regardless of its physical location, with data automatically migrated by the driver on-demand. The trade-off is increased I/O latency when accessing off-GPU data, making intelligent prefetching and caching policies essential for performance.

Offloading implementations often use a tiered storage hierarchy:

• CPU RAM (DRAM): Offers higher bandwidth (e.g., ~50-100 GB/s) and lower latency than NVMe, acting as a primary spill destination. It's ideal for active parts of a long context.
• NVMe SSD: Provides vastly higher capacity (multiple terabytes) but lower bandwidth (e.g., ~3-7 GB/s) and higher latency. It serves as a secondary tier for extremely long contexts or archival of inactive session states. The system must implement eviction policies (e.g., LRU) to move data between GPU, CPU, and NVMe tiers based on access patterns.

To mitigate the latency penalty of off-GPU access, sophisticated prefetching algorithms are critical. As the model generates tokens autoregressively, the system predicts which blocks of the KV cache will be needed next (often the most recent tokens and attention sinks) and loads them into GPU memory ahead of time. This overlaps computation with data transfer. Effective policies are cache-aware, considering the model's attention pattern (e.g., sliding window) to minimize unnecessary I/O and keep the performance-critical working set in fast memory.

KV cache offloading is most powerful when combined with continuous batching. In a batched inference scenario, different requests have varying context lengths and generation states. Offloading allows the system to maintain a much larger total logical KV cache across all requests than the physical GPU memory could hold. The scheduler can offload the cache for completed or paused sequences to make room for active ones, dramatically improving GPU utilization and overall system throughput for workloads with long and diverse contexts.

The KV cache is the foundational memory buffer that stores computed key and value tensors from previous tokens during the autoregressive decoding phase of a transformer. Its primary purpose is to eliminate redundant computation: without it, each new token generation would require recomputing attention over the entire preceding sequence, making inference impractically slow for long contexts. The size of the KV cache scales linearly with batch size, sequence length, and the model's number of layers and attention heads, making it the primary memory bottleneck for large language model inference.

PagedAttention is a memory management algorithm, central to the vLLM inference engine, that organizes the KV cache into non-contiguous, fixed-size blocks or 'pages'. This design is analogous to virtual memory in operating systems and solves critical inefficiencies in continuous batching:

• Eliminates Internal Fragmentation: By using blocks, it avoids wasted memory from padding variable-length sequences to a fixed size.
• Enables Efficient Memory Sharing: Pages can be shared between sequences in scenarios like beam search or parallel sampling, reducing duplicate storage.
• Facilitates Dynamic Management: Pages can be allocated, deallocated, and swapped (e.g., to CPU) on-demand, providing the granular control needed for advanced techniques like offloading.

Continuous batching (or iterative batching) is a scheduling paradigm that dynamically groups incoming inference requests into a single, continuously executing batch. Unlike static batching, where the entire batch must finish before new requests are processed, this technique allows:

• Higher GPU Utilization: The compute unit is kept constantly occupied.
• Improved Throughput: More requests are processed per unit time.
• Reduced Latency: New requests can join the batch immediately. This method creates a highly dynamic and shared KV cache environment, where efficient memory management (like PagedAttention) is essential to handle sequences that start and end at different times without waste.

These are attention architecture variants designed explicitly to reduce the memory footprint of the KV cache.

• Multi-Query Attention (MQA): All query heads share a single key head and a single value head. This drastically reduces the size of the KV cache, often by a factor of 8-32x compared to standard multi-head attention, at a potential cost to model quality.
• Grouped-Query Attention (GQA): A tunable hybrid where groups of query heads share a single key and value head. It offers a practical trade-off, providing most of the KV cache compression of MQA while preserving closer-to-original model performance. These architectures directly lessen the memory pressure that techniques like offloading aim to address.

KV cache quantization is a compression technique that stores the key and value tensors in a lower numerical precision format (e.g., FP8, INT8, or even INT4) instead of the standard FP16 or BF16. This approach:

• Reduces Memory Bandwidth: Lower-precision data transfers are faster, alleviating memory-bound bottlenecks.
• Increases Effective Cache Capacity: More context can be stored in the same physical GPU memory.
• Introduces Trade-offs: Aggressive quantization can introduce noise that may degrade output quality or require calibration. It is often used in conjunction with, or as an alternative to, offloading, depending on the latency/quality/context-length requirements.

Unified Virtual Memory is a hardware/software architecture (e.g., CUDA UVM) that provides a single, coherent virtual address space across CPU and GPU memory. It is a critical enabling technology for practical KV cache offloading because it:

• Simplifies Data Movement: Allows the GPU to access CPU memory directly via page faults, without requiring explicit cudaMemcpy calls for every access.
• Enables On-Demand Paging: Only the actively used 'pages' of the KV cache need to reside in GPU memory; idle portions can reside in CPU RAM or NVMe, and are automatically fetched when needed.
• Adds Overhead: This convenience comes with performance cost due to page fault handling and lower-bandwidth access to CPU memory, which is the fundamental trade-off of offloading.