Key Implementation Details of vLLM’s KV Cache

KV caching is essential to vLLM inference. As sequence-modeling contexts grow, KV-cache architecture evolves frequently, including prefill/decode separation and different cache storage choices.

KV Cache Mechanism

Reference: Inside vLLM V1: A Detailed Analysis of KV Cache Management

Core modules:

• KVCacheManager: KV cache allocation and lookup
• KVCacheCoordinator: Coordinates resource allocation across KV cache groups, including hybrid KV caching
• BlockPool: Manages physical blocks in GPU memory
• BlockTable: Manages logical blocks in GPU memory
• PagedAttention Kernel: CUDA interaction

Some key points:

• Different scenarios use different KVCacheManagers, such as FullAttention and DynamicWindow, deciding what to cache according to the model’s actual context requirements;
• BlockPool uses reference counting to manage block allocation, release, and reuse;
• PrefixCaching introduces hashes as block indexes for fast lookup;
• slot_mapping primarily maps tokens to positions within blocks;

Apart from the CUDA-level GPU operations, the design and implementation are fairly similar to those of other distributed computing engines.

KV Cache Transfer with Prefill/Decode Separation

Reference: From Principles to Evolution: vLLM’s KV Cache Transfer Under PD Separation

• How PD separation works:

  • Prefill: Calculate k, v, and attention for the prompt in parallel, with L^2 computation. This includes QKV for each prompt token and the LxL attention matrix, producing logits and the next token; full self-attention.
  • Decode: Based on the prefill result, concatenate new tokens and generate each next token, repeatedly loading the KV cache; single-query attention.
  • Similarities: Both produce one next_token_ids per request, optionally include log probabilities, and update or extend the KV cache. Their API and data structures represent the same kind of result, for example GenerationBatchResult with next_token_ids and extend_logprob_start_len_per_req.
  • Differences: Inputs differ, a full sequence versus a single token, so computation, data volume, and latency differ. Prefill can also produce log probabilities for the full input; decode generally produces them only for the current next token.

• PD separation, or Prefill Decoding Separation, is an important vLLM feature for distributed deployments that optimizes long-prompt processing and decoding latency. It splits inference between two independent vLLM instances: a Prefill Instance (P) processes long prompts and generates the KV cache, while a Decode Instance (D) receives that cache and performs efficient autoregressive decoding.

• The central challenge is transferring KV cache data efficiently and with low latency from P’s GPU to D’s GPU. vLLM offers several KV Connector implementations, broadly divided into centralized and peer-to-peer transfer.

• Centralized transfer: Concentrate complexity in storage. For example, vLLM’s built-in Mooncake uses high-performance technologies such as RDMA to address transfer bottlenecks;

  • Mooncake internally separates Mooncake Store and TE (Transfer Engine), supporting centralized and peer-to-peer modes respectively

• P2P transfer: Establish peer-to-peer communication groups to reduce single-point bottlenecks, at the cost of complex configuration; vLLM’s p2pNcclConnector is an example;

Centralized transfer is easy to understand, but with large model parameter sets, a single storage point can still become a bottleneck, for example in bandwidth. P2P is a relatively new approach that bypasses traditional distributed frameworks’ strict rank and world-size constraints to enable dynamic scaling.

• ZMQ negotiation: Each P/D instance maintains a ZMQ server for control-plane negotiation. On their first communication, P and D exchange addresses and NCCL connection information through ZMQ.
• NCCL communication group: P and D establish an NCCL group containing only those two nodes, with world size = 2.

For implementation details, see vLLM’s P2P NCCL Connector¶.