KVPP 设计¶

1. 背景¶

MLA/SFA 模型会在张量并行（TP）rank 之间复制历史 KV 缓存，从而限制长上下文和并发请求的容量。KVPP 在每个 TP 组内按层分配持久缓存，并在计算需要时广播该层的缓存。

模型执行保留其 TP/EP/PP 配置。调度器保留完整的逻辑缓存规格和块编号。

2. 设计¶

KVPP 结合层所有权、两个可复用的临时缓冲区，以及提前预取一层：

TP group within one PP stage
    Owner rank: persistent main KV + indexer KV for its layers
         |
         +-- Full-layer broadcast --> Other ranks: scratch buffer
                                          |
                                          +-- Wait --> Compute current layer
                                                       Prefetch next layer

2.1 层所有权¶

每个流水线并行（PP）阶段会按层索引对其本地目标层进行排序，并在其 TP 组内将连续且均衡的分区分配给所有者 rank。一个层的主 KV、indexer KV 和量化组件构成一个 bundle，并拥有单一所有者。MTP 缓存仍由需要它们的每个 rank 独立分配，并且被排除在目标层分区之外。

每个 rank 为其拥有的层分配持久存储，并复用两个临时缓冲区用于其他层。目标层的执行序号会选择交替的临时缓冲区，使当前层和下一次预取可以使用独立存储。

2.2 物理布局与内存预算¶

bundle 内的组件是连续的，并且每个组件包含该层的所有逻辑块。层不需要一次全局连续分配；只有每个广播 bundle 必须是连续的。布局遵循实际缓存规格，包括主 KV、indexer 数据和 scale。

每个 rank 上的物理预算为：

Physical bytes per block
  = Sum of owned target-bundle bytes per block
  + Local MTP-cache bytes per block
  + 2 * Largest target-bundle bytes per block

Available blocks = floor(Available KV memory / Physical bytes per block)

worker 会将得到的块容量转换为 planner 的逻辑内存预算。物理分配使用最终块数，在按 KVPP 布局分配存储的同时保留逻辑缓存管理。

2.3 整层广播与预取¶

任何包含已计算 token 的实际请求的前向传播都会从预取第一层开始。每个注意力层在投影之后、首次 KV 缓存访问或写入之前等待其广播，然后开始预取下一层。

每次广播都保持在当前 PP 阶段的 KVPP 组内。所有者会一次性广播完整的 bundle，包括所有组件和所有已分配的逻辑块。其他 rank 将其接收至对应的临时缓冲区。计算完成后，所有者会保留更新后的缓存，供后续前向传播使用。

预取使用独立的传输流。ready 事件在广播前建立数据依赖，completion 事件确保在 wait 返回之前设备传输已完成。下一层的广播可以与当前层计算重叠；重叠程度取决于通信和计算开销。没有历史的前向传播、dummy 运行和 profiling 路径不会广播历史缓存。

3. 支持与约束¶

KVPP 以通信换取持久缓存容量。每个 rank 仍然需要两个按最大目标层大小设置的临时缓冲区，外加其独立的 MTP 缓存。因此，内存节省取决于层数、TP 大小和每层缓存大小。广播流量随已分配块数增长；负载不会按请求或活动块进行过滤。

4. 用法¶

通过模型启动参数启用 KVPP。KVPP 组大小跟随 TP 大小：

vllm serve <model-path> \
  --tensor-parallel-size 2 \
  --enforce-eager \
  --additional-config '{"enable_kvpp": true}'

enable_kvpp 默认为 false。启用 PP 后，每个阶段会在其自己的 TP 组内分配所有者并进行广播。