Prhub

#50493 [Kimi-K3] support DCP partial prefix cache hit

原始 PR 作者 GirasoleY 合并时间 2026-08-18 08:05 文件变更 8 提交数 6 评论 7 代码增减 +523 / -37

执行摘要

Kimi-K3 开启 DCP 部分前缀缓存命中,修复 Mamba 块表越界

PR body 明确说明这是 Kimi-K3 DCP 支持(#50484)的后续:"Enables hash-aligned partial-prefix reuse under DCP for FullAttention plus aligned Mamba while preserving sharded-attention and replicated-Mamba correctness"。同时修复了一个高危几何错误:"MRV2 previously divided every cache group's block-table width by DCP size—correct for attention, but wrong for replicated Mamba. For a 1M context, DCP8, and 16-token Mamba blocks, it allocated 8,192 entries instead of 65,536, potentially causing metadata corruption and CUDA illegal-address failures"。

值得精读。核心设计决策是"cache identity 而非 virtual-page 或 interleave 几何决定 hash 语义"以及"每种 cache 类型自己负责 CP 感知的块表宽度",这两个思路对后续 hybrid 模型(attention + mamba)在并行化下的缓存设计有直接借鉴价值。建议按 kv_cache_coordinator.pymodel_runner.pyblock_table.py 的顺序阅读三个源码文件,再配合 test_gpu_model_runner_v2.py 中捕获块表宽度的测试技巧(monkeypatch get_block_table_width 并断言两次调用值)理解修复前后差异。

讨论亮点

本 PR 的实质代码 review 讨论较少,主要流程记录如下:

  • claude[bot] 指出该 PR 来自 fork,自动 review 被关闭,需维护者手动触发 @claude review
  • 维护者 ivanium 最终给出 LGTM,并致歉 "Sorry for the delayed review.",属于延期但无争议的审批。
  • mergify[bot] 多次提示 merge conflict 并要求 rebase,期间还有一次 pre-commit 失败;提交历史上可见 3 次 merge main 分支的操作,属于常规的冲突消解,没有引发设计变更。

综合来看,评审意见没有对实现方案提出质疑或替代设计,核心决策(DCP 下放宽 hash 对齐条件、按 cache 类型计算块表宽度)由作者在 commit message 中自行论证后直接落地。

实现拆解

实现按 4 个步骤拆解:

  1. 缓存协调器放宽 partial 命中门控vllm/v1/core/kv_cache_coordinator.py
    - 原逻辑 enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group,且非 DCP 下要求 Mamba block_size > hash_block_size(strictly greater)。
    - 新逻辑允许 DCP 下 block_size >= hash_block_size,因为 DCP 会把 full-attention 的 effective block 放大为 block_size * dcp_world_size,两种 cache 组在 hash 边界对齐的条件与 TP 不同。
    - 保留 manager 级降级保护:任一 cache manager 不支持细粒度 hash 查找且其块大小不等于 hash_block_size 时,整体退回 block-aligned 命中并打印 warning。

  2. 块表几何改为按 cache 类型计算vllm/v1/worker/gpu/model_runner.py
    - GPUModelRunner.initialize_kv_cache 不再统一执行 cdiv(max_model_len, spec.block_size * dcp_size),改调 spec.max_num_blocks_per_req(self.vllm_config, block_table_max_model_len),让 FullAttentionSpecMambaSpec 各自表达拓扑语义(attention 分片 vs Mamba 复制)。
    - 修复效果:1M context、DCP8、16-token Mamba 块场景下,Mamba 块表从错误的 8,192 项恢复为 65,536 项。

  3. 块表新增越界写保护vllm/v1/worker/gpu/block_table.py
    - BlockTables.append_block_idsstage_write 之前检查 end > row_capacity,失败路径从静默越界写(可能破坏相邻行元数据)改为显式 RuntimeError,同时把 num_blocks.np 的更新统一为 end 变量。

  4. 测试配套(5 个测试文件)
    - tests/v1/worker/test_gpu_model_runner_v2.py(新增):验证 Mamba 块表宽度不被 DCP shard,以及越界写被拒绝。
    - tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py:DCP world size 参数化下验证 partial hit 的 CoW、replicated Mamba snapshot 恢复、joint hit 上界、EAGLE 场景回退一个 hash unit。
    - tests/distributed/test_kimi_linear_context_parallel.py:端到端对比 DCP size 1 与 2 在 K3 小模型上的前缀复用输出一致性与缓存命中数。
    - tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.pytests/v1/attention/test_mla_backends.py:分别覆盖块池 primitive 的 DCP 参数化和 MLA 后端对非 virtual-block-aligned 前缀的接受。

文件 模块 状态 重要度
vllm/v1/core/kv_cache_coordinator.py 缓存协调 modified 6.37
vllm/v1/worker/gpu/model_runner.py 模型运行 modified 6.73
vllm/v1/worker/gpu/block_table.py 块表管理 modified 6.05
tests/v1/worker/test_gpu_model_runner_v2.py 模型运行 added 7.43
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py 前缀缓存 modified 7.38
tests/distributed/test_kimi_linear_context_parallel.py 分布式测试 modified 6.33
tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py 前缀缓存 modified 5.92
tests/v1/attention/test_mla_backends.py 注意力后端 modified 5.15

关键符号

KVCacheCoordinator.__init__ GPUModelRunner.initialize_kv_cache BlockTables.append_block_ids KVCacheSpec.max_num_blocks_per_req Scheduler._mamba_block_aligned_split build_mla_chunked_context_metadata

关键源码片段

vllm/v1/core/kv_cache_coordinator.py core-logic

partial prefix cache hit 的总开关。将 enable_partial_hash_hits 从 DCP 下强制关闭改为按 hash 对齐条件判定,并细化 TP 与 DCP 对 Mamba block_size 的不同要求,是本次功能开启的核心逻辑。

# 判定是否启用 partial hash hit(细粒度前缀缓存命中)
# 前置条件:存在 Mamba "align" 缓存组
has_partial_mamba_group = any(
    isinstance(g.kv_cache_spec, MambaSpec)
    and g.kv_cache_spec.mamba_cache_mode == "align"
    and (
        # TP 场景:Mamba block 必须严格大于 hash_block_size,
        # 保证 hash 边界能落在 block 内部,形成可复用的部分块
        (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size)
        or (
            # DCP 场景:full-attention 的 effective block 被放大为
            # block_size * dcp_world_size,因此 hash_block_size 只要
            # 不大于 Mamba block 即可让两种 cache 组在 hash 边界对齐
            dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size
        )
    )
    for g in kv_cache_config.kv_cache_groups
)
# 之前是 dcp_world_size == 1 才允许,现在 DCP 下同样启用
self.enable_partial_hash_hits = has_partial_mamba_group
if self.enable_partial_hash_hits:
    # 若存在不支持细粒度查找的 manager 且其块大小不等于 hash_block_size,
    # 则统一降级为 block-aligned 命中,保证正确性优先
    unsupported_partial_hit_managers = {
        type(manager).__name__
        for manager in self.single_type_managers
        if not manager.supports_fine_grained_hash_lookup
        and manager.block_size != hash_block_size
    }
    if unsupported_partial_hit_managers:
        self.enable_partial_hash_hits = False
        logger.warning_once(
            "Disabling fine-grained prefix-cache hits because these KV "
            "cache managers require block-aligned lookups: %s.",
            ", ".join(sorted(unsupported_partial_hit_managers)),
        )
vllm/v1/worker/gpu/model_runner.py data-contract

MRV2 块表几何修复的主路径。将统一按 DCP size 整除的宽度计算改为委托给 spec.max_num_blocks_per_req(),区分 attention 分片与 Mamba 复制,是修复 1M context 下元数据损坏的关键。

block_sizes = []
max_num_blocks_per_group = []
for kv_cache_group in kv_cache_config.kv_cache_groups:
    spec = kv_cache_group.kv_cache_spec
    block_sizes.append(spec.block_size)
    # 让每种 cache 类型自行计算 CP 感知的块表宽度:
    # - FullAttention 的 KV 在 DCP rank 间分片,单个块覆盖
    # block_size * dcp_size 个全局 token
    # - Mamba/GDN 的循环状态在 DCP rank 间复制,
    # 因此块表仍需覆盖完整全局序列位置,不能被 DCP size 整除
    max_num_blocks = spec.max_num_blocks_per_req(
        self.vllm_config, block_table_max_model_len
    )
    # 应用拓扑感知宽度后,再保留各 cache 类型的块对齐要求
    if isinstance(spec, MambaSpec):
        max_num_blocks = get_block_table_width(
            max_num_blocks, spec.block_size, token_alignment=None
        )
    else:
        max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
    max_num_blocks_per_group.append(max_num_blocks)
vllm/v1/worker/gpu/block_table.py core-logic

为 append_block_ids 增加行容量边界检查,阻止越界写覆盖相邻行元数据,将静默内存破坏转换为显式 RuntimeError,是本次 bug 的防御性收口。

# BlockTables.append_block_ids 的逐 group 写入循环
# (省略了 start 由 request 当前已用块数推导等前置逻辑)
for i, block_ids in enumerate(new_block_ids):
    # blocks_per_kv_block > 1 时,单个逻辑块横跨多个物理块,需先展开
    bpk = self.blocks_per_kv_block[i]
    if bpk > 1:
        block_ids = [b * bpk + k for b in block_ids for k in range(bpk)]
    end = start + len(block_ids)
    row_capacity = self.block_tables[i].gpu.shape[1]
    # 在 stage_write 之前拦截越界写入:一旦 end 超过行容量,
    # 未检查的写会覆盖相邻行的元数据,这正是 DCP 大 context 场景下
    # Mamba 块表尺寸错误导致 CUDA illegal-address 的根因之一
    if end > row_capacity:
        raise RuntimeError(
            f"Block table write for request {req_index}, group {i} exceeds "
            f"row capacity ({end} > {row_capacity})"
        )
    self.block_tables[i].stage_write(req_index, start, block_ids)
    self.num_blocks.np[i, req_index] = end

评论区精华

fork PR 自动 review 关闭 other

claude[bot] 指出该 PR 来自 fork,automated review 被禁用,需维护者手动触发 @claude review。

结论:未触发额外 review,由维护者 ivanium 手工审核。 · 已解决

维护者延迟审批 other

ivanium 审批时致歉 "Sorry for the delayed review." 并给出 LGTM,全程未提出代码质疑。

结论:直接 approve,无修改要求。 · 已解决

merge conflict 与 pre-commit 流程 other

mergify[bot] 多次提示需要 rebase 解决冲突;pre-commit 曾失败一次。提交历史显示 3 次 merge main 操作完成消解。

结论:通过多次 merge main 解决冲突,CI 最终通过。 · 已解决

风险与影响

  1. 门控放宽的影响面kv_cache_coordinator.pyenable_partial_hash_hits 不再要求 dcp_world_size == 1,只要满足 FullAttention + align Mamba 且 DCP 开启就会启用 partial hash hits,影响范围不限于 Kimi-K3,其它 hybrid 模型同样被覆盖。
  2. 数据契约变更model_runner.py 引入 spec.max_num_blocks_per_req() 作为块表宽度的唯一入口,所有 KVCacheSpec 子类(包括非 DCP 场景)的 attention 宽度计算路径都被改写,属于 data-contract 级别变化,需要关注自定义 spec 的兼容性。
  3. 异常路径行为变化block_table.py 新增的越界 RuntimeError 会把原先"内存被静默破坏、偶发 illegal-address"转换为确定性崩溃。错误配置下用户体验变化明显,但方向是安全的(fail fast)。
  4. 分布式语义复杂度:DCP 下 partial hit 的 joint-hit 上界依赖 replicated Mamba snapshot 数量(test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots 专门验证),多请求并发交错时缓存分配与 CoW 行为更难推理,端到端分布式测试仅覆盖 DCP size 2。

对用户:Kimi-K3 在 DCP 部署下获得从前缀缓存复用;1M context 场景避免块表元数据损坏导致的进程崩溃,属于直接可用性修复。
对系统:缓存命中粒度从 scheduler block 细化到 hash block,提高缓存利用率;块表内存分配按 cache 类型精确化,attention 部分不变、replicated Mamba 部分显著增大(1M/DCP8 场景为 8 倍),属于预期内的内存回退。
对团队:为 K3 DCP 后续功能(如 DSpark 推测解码、RecoverSSM)扫清了块表与缓存协调的前置障碍,MRV2 与 prefix cache 逻辑的耦合关系也因此更加清晰。影响程度中等偏上,集中在 v1 执行路径,不涉及旧版 V0 路径。

核心路径变更 数据契约调整 分布式行为依赖测试 影响范围超出 K3 模型

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论