KV Cache

vLLM's defining innovation is PagedAttention, which manages the KV cache using a virtual memory paging system analogous to an operating system.

• Block-Based Allocation: KV cache is allocated in fixed-size blocks (e.g., 16 tokens). Each request's sequence is stored across non-contiguous blocks.
• Eliminates External Fragmentation: This allows for near-zero waste from internal fragmentation, enabling high GPU memory utilization even with highly variable sequence lengths.
• Shared Prefix Caching: Blocks can be shared between requests that share a common prompt prefix, a critical optimization for multi-tenant serving.
• Performance Impact: This design is the primary reason for vLLM's industry-leading throughput, especially in scenarios with long sequences and many concurrent users.

Text Generation Inference (TGI) integrates KV cache management with its continuous batching (also called iterative or dynamic batching) engine.

• Dynamic Cache Allocation: Memory for the KV cache is allocated on-the-fly as new tokens are generated for each request in the batch.
• Batch State Management: The system maintains the KV cache state for all active requests in the batch, evicting finished sequences and adding new ones without stopping the engine.
• Flash-Attention Integration: TGI leverages optimized attention kernels like FlashAttention-2 to compute attention over the cached KV states efficiently.
• Framework Support: Primarily serves models from the Hugging Face Transformers library, providing a seamless path from training to optimized serving.

TensorRT-LLM uses a technique called inflight batching to manage the KV cache, tightly coupled with its compiler-based optimizations for NVIDIA GPUs.

• Static & Dynamic Planning: The TensorRT compiler plans KV cache memory usage during model compilation, but the runtime handles dynamic allocation within those plans.
• Memory Reuse: Aggressively reuses memory buffers between different phases of execution and between finished and new requests to minimize overhead.
• Kernel Fusion: The KV cache operations are fused with attention computation kernels, reducing data movement and improving latency.
• Hardware-Specific Optimization: The entire KV cache pipeline is optimized for the target GPU architecture (e.g., Hopper, Ada Lovelace), leveraging features like FP8 support and new attention instructions.

Triton Inference Server provides a flexible, backend-agnostic platform where KV cache implementation is delegated to the specific backend or model framework being used.

• Backend Responsibility: A PyTorch, TensorFlow, or ONNX Runtime backend within Triton is responsible for managing its own KV cache. Triton orchestrates the requests and responses.
• Dynamic Batching: Triton's dynamic batcher can group inference requests, which the backend must then handle, including managing the KV cache for the batched sequences.
• Multi-Framework Support: This allows teams to serve models from different frameworks with a consistent API, but KV cache efficiency depends on the chosen backend's capabilities.
• Use Case: Often used for deploying custom or proprietary model serving pipelines where control over the entire stack is required.

For maximum control, engineers can implement a basic KV cache directly in PyTorch, which reveals the underlying mechanics.

• Manual Tensor Management: The cache is typically implemented as a pair of growing tensors (cache_k, cache_v) that are concatenated with new key/value states for each generated token.
• Attention Masking: A causal attention mask is applied to ensure each token only attends to itself and prior tokens in the cache.
• The Naive Limitation: A simple implementation suffers from internal fragmentation as all sequences in a batch must be padded to the length of the longest sequence, wasting memory and compute.
• Foundation for Optimization: This approach is the starting point for understanding more advanced techniques like PagedAttention or continuous batching, which solve the fragmentation problem.

An inference optimization technique that dynamically groups incoming requests into a single batch for processing, even if they arrive at different times. This maximizes GPU utilization and throughput by keeping the hardware constantly occupied.

• Key Mechanism: Unlike static batching, it doesn't wait for a fixed batch size. It continuously adds new requests and evicts completed ones from the batch.
• Synergy with KV Cache: Efficiently manages the KV cache for multiple concurrent sequences within a batch, sharing compute where possible while isolating per-sequence state.
• Impact: Can increase throughput by 5-10x compared to request-by-request processing, directly reducing inference cost per token.

An algorithm introduced by the vLLM serving engine that manages the KV cache using concepts from operating system virtual memory. It treats the KV cache as non-contiguous "pages" that can be allocated and freed dynamically.

• Solves Fragmentation: Traditional KV cache allocation for variable-length sequences leads to significant memory fragmentation. PagedAttention eliminates this by using blocks.
• Enables Memory Sharing: Allows for efficient sharing of KV cache blocks between sequences in scenarios like beam search or parallel sampling, drastically reducing memory overhead.
• Result: Enables higher batch sizes and longer context windows within the same GPU memory, a foundational optimization for modern high-throughput servers.

A high-throughput and memory-efficient inference and serving engine for LLMs. Its performance is largely attributed to its implementation of PagedAttention for optimal KV cache management.

• Core Innovation: Re-frames KV cache management as a virtual memory paging problem, solving the fragmentation issue inherent in autoregressive decoding.
• Performance Claims: Achieves up to 24x higher throughput than Hugging Face Transformers in common serving scenarios.
• Ecosystem Role: Has become a de facto standard for research and production serving due to its simplicity, performance, and broad model support.

An IO-aware, exact attention algorithm that reduces memory reads/writes between GPU memory hierarchies. While primarily a training optimization, it has major implications for inference and the nature of the KV cache.

• Principle: Minimizes costly HBM (High Bandwidth Memory) accesses by performing more computations on-chip (SRAM), following a tiling strategy.
• Impact on Inference: Makes the initial computation of the KV cache for the prompt phase much faster and more memory-efficient. Some inference kernels integrate FlashAttention-like ideas for the prefill stage.
• Distinction: Optimizes the computation of attention (creating K/V). KV Cache optimizes the reuse of already-computed K/V during generation.

Attention architecture variants designed to reduce the memory footprint and bandwidth cost of the KV cache.

• Multi-Query Attention (MQA): Uses a single key and value head shared across all query heads. This dramatically shrinks the KV cache size.
• Grouped-Query Attention (GQA): A middle ground between MHA and MQA. Groups query heads share a single key and value head. Provides a tunable trade-off between model quality (closer to MHA) and inference efficiency (closer to MQA).
• Inference Benefit: For a 7B model, MQA/GQA can reduce the KV cache memory by 4-8x, allowing for much larger batch sizes or context lengths. Used in models like Llama 2 (GQA) and Falcon (MQA).

An inference technique that uses a smaller, faster draft model to propose a sequence of tokens, which are then verified in parallel by the larger target model using its KV cache.

• Process: The draft model runs autoregressively for K steps. The target model then takes this proposed sequence, uses its KV cache to compute logits for all K tokens in a single, parallel forward pass, and accepts the correct prefix.
• KV Cache Role: The target model's ability to score the entire draft sequence in one pass is only efficient because it can attend to its own previously generated KV cache for the context, plus the new draft tokens.
• Outcome: Can achieve 2-3x latency reduction for the same model, effectively making the larger model "think faster" by amortizing the verification cost.