vLLM KV Cache 实现要点

发表于 | 分类于

vLLM 的 KVCache 机制是推理过程中必不可少的，但是随着序列建模的 context 不断增长，带来的 KV Cache 相关的架构迭代也非常频繁，比如 PD 分离、多种 Cache 存储选型等。

KV Cache 机制

参考文章：深入vLLM V1内核：KV cache 管理机制详细剖析

page-attention

核心模块：

  • KVCacheManager：KV Cache 分配和查询
  • KVCacheCoordinator：协调不同 KV Cache Group 的资源分配，涉及到 Hybrid KV Cache
  • BlockPool：管理 GPU 内存中的物理 block
  • BlockTable：管理 GPU 内存中的逻辑 block
  • PagedAttention Kernel：CUDA 交互

部分要点：

  • 不同的场景会使用不同的 KVCacheManager，比如 FullAttention 和 DynamicWindow 等，根据模型实际上下文的需求来决定 Cache 的内容；
  • BlockPool 通过引用计数法管理 Block 生命周期，包括申请、释放、复用等；
  • PrefixCaching 通过引入 hash 值作为 block 的索引来实现快速查询；
  • slot_mapping 的概念，主要是用于定位 token 到 block 位置的映射；

整体看下来，除了 CUDA 部分对于 GPU 的操作，其他的设计和实现还是和其他分布式计算引擎比较类似的。

PD 分离下的 KV Cache 传输机制

参考文章：从原理到演进：vLLM PD分离KV cache传递机制全解析

PD 分离（Prompt Decoding Separation）是 vLLM 在分布式部署中，为优化长提示词（Prompt）处理和解码延迟而引入的一项重要特性。它将推理过程分为两个独立的 vLLM 实例：Prefill Instance (P) 负责处理长提示词并生成 KV Cache，Decode Instance (D) 负责接收 KV Cache 并进行高效的自回归解码。
PD 分离的核心挑战在于如何高效、低延迟地将 P 实例生成的 KV Cache 数据从 P 实例的 GPU 传输到 D 实例的 GPU。vLLM 为此提供了多种 KV Connector 实现，主要分为中心化传输和 P2P 传输两大类。

  • 中心化传输：把复杂度收敛在存储侧，比如 vllm 里内置的 Mooncake，通过 RDMA 等高性能存储技术解决传输瓶颈问题；
    • 这里有一个点是 Mooncake 内部其实分离出了 Mooncake Store 和 TE(Transfer Engine)，分别支持了中心化和 p2p 的两种模式
  • P2P 传输：建立 p2p 通信组，减少单点瓶颈，但是配置复杂，比如 vllm 里的 p2pNcclConnector；

关于中心化传输很好理解，但是对于大规模模型参数来说，单点存储本身还是容易引入瓶颈，比如带宽等问题。p2p 是一个相对新的方式，绕开了传统分布式框架对 rank 和 world size 的严格限制，实现了动态扩缩容。

  • ZMQ 协商：每个 P/D 实例维护一个 ZMQ Server，用于协商控制流。当 P 和 D 首次通信时，它们通过 ZMQ 交换彼此的地址和 NCCL 连接信息。
  • NCCL 通信组：P 实例和 D 实例之间会建立一个仅包含两个节点的 NCCL 通信组（world size = 2）。

实现细节可以参考 vllm 的这篇文章：P2P NCCL Connector¶