缓存与分离

PagedAttention

Woosuk Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023.

vLLM 的核心。PagedAttention 解决的是 KV cache 的存储问题。

关于 Transformer 的 attention：

Attention(Q,K,V)=softmax(QKTd)V \text{Attention}(Q, K, V)=\text{softmax}(\frac{QK^T}{\sqrt{d}})V

Decode 阶段每个 token 的 q 都要 attend 之前 tokens 的 k 和 v，所以维护一个持续增长的 KV cache。一个已经部署的模型常常需要处理并发的请求，对于每个请求都要维护一个 KV cache；这样，KV cache 的内存管理效率就成了推理的一个重要瓶颈。

传统的方法为每个推理请求都预留一段连续的显存。这很容易产生 memory fragmentation。

等等，这不跟 OS 里面的 memory 管理是一个问题吗？直接把 paging 的思想挪过来，把存储分成 logical block 和 physical block 两个抽象层级。KV cache 在逻辑上连续，物理上可能是不连续的。和 OS 一样，这样做的好处是可以按需分配显存，唯一可能产生的浪费是最后一个未填满的 block，所以有效降低了 internal 和 external fragmentation。

还有一个很巧妙的设计是 Copy-on-Write（CoW）。在某些情况比如说 beam search 和 parallel sampling 中，LLM 要对一段相同的 prompt 生成四个不同的回答，有了 PagedAttention，我们直接让它们的 block table 指向相同的 logical blocks 就行。当这些请求各自生成不同的 tokens 时，再给各自需要修改的 block 分配独立空间。

DistServe

Yinmin Zhong et al., DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving, OSDI 2024.

Prefill 和 Decode 两个阶段的计算特征明显是不同的。传统 LLM 推理系统把这两个阶段放在同一批 GPU 上，利用 continuous batching 一起调度。这种设计的好处是可以复用模型权重和 KV cache 从而提高整体 throughput，但是 prefill-decode interference 还是会产生很多阻塞和竞争。

可以看到 Prefill-heavy 的请求 TTFT 会比较长，Decode-heavy 的请求 TPOT 会比较长。如果把两个阶段放在同一张 GPU 上，TTFT 和 TPOT 就处于此消彼长的状态，没法同时优化。解决方法很简单：把 Prefill 和 Decode 放在不同的 GPU 上处理（disaggregation）。完整的处理流程是这样的：

  • 请求进入 prefill instance，处理 prompt，生成首个 token 和 KV cache
  • prefill instance 把 KV cache 传给 decode instance
  • decode instance 接着生成后续 tokens 直到请求结束

DistServe 的成本：prefill 和 decode instance 各自保存模型权重的空间成本，以及 KV cache 的传输成本。DistServe 把 GPU 的分配也建模成一个优化问题（placement），在考虑网络拓扑的情况下把 GPU 分配给 prefill 和 decode，使得 goodput 尽量高。

DistServe 的好处：消除了 prefill-decode interference，可以分别针对 TTFT 和 TPOT 优化两组 GPU，使资源的分配更灵活。比如说 prefill 通常更 compute-intensive，如果 TTFT 要求严格，可以增加 TP 让多个 GPU 同时执行一个 prefill。而 decode 很依赖 batching，如果 TPOT 要求严格，可以先通过 TP 降低单步延迟，再使用 PP 提高 throughput。这在传统的 colocated serving（prefill 和 decode 在同一批 GPU 上处理）中是难以实现的。

Mooncake

Ruoyu Qin et al., Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, FAST 2025

月之暗面给 Kimi 设计的长上下文推理系统。这篇文章看上去也偏系统一点。

Prefill 阶段是非常 compute-intensive 的，但由于 LLM 的自回归特征，对 prefix 的复用能够有效减轻 prefill 的计算量。此前的本地 Prefix Caching 方法主要把可复用的 KV cache 存在 GPU 的显存中，但显存相当珍贵，月饼尝试把 KV cache 从 GPU 中独立出来，利用集群中相对空闲的 CPU DRAM，SSD 和高速网络，构建一个分布式 KV cache pool。这样，KV cache 就不再依附于某张特定的 GPU 而成为集群级别的资源。

月饼是这样处理一个请求的：

  • 在分布式 KV cache pool 中寻找相同 prefix 的缓存，加载到 prefill GPU instance
  • incremental prefill：只计算未命中的 token，生成新的 KV cache
  • 将新生成的 KV cache 传输到 decode GPU instance（传输过程是 layer-wise 的）
  • 利用 KV cache 进行 continuous batching，decode 出回答

在分布式 KV cache pool 之上，有一个全局 cache-aware scheduler，月饼称其为 conductor。一方面，它综合考虑 TQueue,TTransfer,TPrefillT_{\text{Queue}},T_{\text{Transfer}},T_{\text{Prefill}} 与 SLO，把请求分配到合适的 GPU 上；另一方面，它还需要管理 cache 本身，对频繁访问的 KV blocks 进行 hot-spot replication，以缓解热点节点的网络拥塞；还要对不常被访问的 cache 进行 eviction/swapping，腾出存储空间。

除此之外，还有两个机制。

首先是 Chunked Pipeline Parallelism（CPP），它把特别长的 prompt 切成多个 chunks，并把 transformer layer 分配到不同节点上，通过 pipeline 处理。实际上是使用 PP 的思想，以减少 SP 中频繁的跨节点 collective communication。

另一个是 prediction-based early rejection，即在执行 prefill 之前，conductor 会估计后续 decode 阶段是否有足够资源；如果判断负载过高，就拒绝请求，以降低那些注定无法在 SLO 内完成的请求导致的资源浪费。

月饼这套系统很适用于具有长 prompt，重复 prefix，高并发和较高速节点互联的 serving workloads。prefix cache hit rate 和网络带宽都可能成为系统的瓶颈。