执行摘要

• 一句话：修复 DCP 下 MLA KV 越界检查误报
• 推荐动作：建议阅读该 PR 以理解 DCP 下逻辑索引与物理索引空间的口径差异。修复本身非常直接，但思路值得借鉴——在 DCP 或任何跨卡并行场景中，OOB 检查必须与内核实际的索引映射保持一致。可顺带排查 memory_pool.py 中其他使用物理容量做上界的断言（如 get 路径）是否存在类似误报风险。

功能与动机

PR body 明确指出：set_mla_kv_buffer receives a widened logical loc under DCP — set_mla_kv_buffer_kernel keeps loc % DCP_WORLD_SIZE == DCP_RANK and then divides by the world size to reach the physical row — but the bounds check compared that logical index against the physical capacity, so it fires spuriously once the pool is small enough for logical indices to exceed it. 复现实验显示，在 4x GB300 上使用 test/registered/dcp/test_kimi_linear_dcp4.py 加 --max-total-tokens 8192 时出现 index >= 8256 的误报，且越界索引落在物理边界的 [2x, 4x) 区间，正好对应 dcp_size=4 的逻辑空间，证明是指纹空间口径不一致而非真正越界。

实现拆解

实现步骤

1. 定位问题：python/sglang/srt/mem_cache/memory_pool.py 的 set_mla_kv_buffer 是 MLA KV 写入入口，写入前调用 maybe_detect_oob(loc, 0, self.size + self.page_size, ...) 对 loc 做越界断言，上界取的是物理容量。

2. 对齐 DCP 索引语义：在 DCP 下，set_mla_kv_buffer_kernel 会先按 loc % DCP_WORLD_SIZE == DCP_RANK 过滤出本 rank 对应的逻辑索引，再除以 world size 得到物理行。因此传给该函数的 loc 取值范围是物理容量的 attn_dcp_size 倍，物理容量不能直接作为逻辑索引的上界。

3. 修改断言上限：将 maybe_detect_oob 的上界改为 (self.size + self.page_size) * get_parallel().attn_dcp_size，并新增注释说明 loc 在 DCP 下是加宽后的逻辑索引，内核自身会完成取模与缩放。

4. 兼容性保障：attn_dcp_size 在非 DCP 场景下为 1，因此非 DCP 的断言行为完全不变；改动不影响写入路径的实际数据布局，只影响越界检测的判定口径。

5. 验证方式：在 4x GB300 上使用 test/registered/dcp/test_kimi_linear_dcp4.py，配合 SGLANG_ENABLE_ASYNC_ASSERT=1 和 --max-total-tokens 8192 可复现原误报，修复后 3/3 通过；PR 未新增独立测试文件，依赖既有 DCP 测试覆盖。

关键文件：

• python/sglang/srt/mem_cache/memory_pool.py（模块 KV 缓存；类别 source；类型 core-logic；符号 set_mla_kv_buffer）: 唯一变更文件。set_mla_kv_buffer 是 MLA KV 写入的公共入口，原 OOB 检查在 DCP 下使用物理容量做上界，与内核按 world size 取整的逻辑索引空间不一致，导致小池场景误报。修正后上界乘以 attn_dcp_size，逻辑索引口径与内核保持一致。

关键符号：set_mla_kv_buffer

关键源码片段

python/sglang/srt/mem_cache/memory_pool.py

唯一变更文件。set_mla_kv_buffer 是 MLA KV 写入的公共入口，原 OOB 检查在 DCP 下使用物理容量做上界，与内核按 world size 取整的逻辑索引空间不一致，导致小池场景误报。修正后上界乘以 attn_dcp_size，逻辑索引口径与内核保持一致。

def set_mla_kv_buffer(
    self,
    layer: RadixAttention,
    loc: torch.Tensor,
    cache_k_nope: torch.Tensor,
    cache_k_rope: torch.Tensor,
):
    # DCP 下 loc 是加宽的逻辑索引：内核 set_mla_kv_buffer_kernel 会先过滤
    # loc % DCP_WORLD_SIZE == DCP_RANK，再除以 world size 得到物理行，
    # 因此逻辑索引上限是物理容量的 attn_dcp_size 倍。
    maybe_detect_oob(
        loc,
        0,
        (self.size + self.page_size) * get_parallel().attn_dcp_size,
        "set_mla_kv_buffer (MLA)",
    )
    layer_id = layer.layer_id
    self._write_mla_kv_buffer(
        self.kv_buffer[layer_id - self.start_layer],
        loc,
        cache_k_nope,
        cache_k_rope,
    )

评论区精华

该 PR 没有实质性的 review 讨论。审核人 Fridge003 直接批准（未附说明）。Issue 评论中仅有一条 gemini-code-assist 机器人提醒其代码审查服务已停止（无技术内容），以及作者本人触发的 /tag-and-rerun-ci 重跑指令。整体上这是一个改动明确、审核迅速的小修复。

• 暂无高价值评论线程

风险与影响

• 风险：
  • DCP 执行路径变更：修改的是 DCP 模式下 MLA KV 写入前的断言逻辑，虽然不影响实际写入正确性，但改变了 OOB 检测的判定口径；若未来内核索引映射方式变化（例如不再按 world size 简单取模缩放），该上限需要同步更新。
  • 依赖并行状态初始化：新增对 get_parallel().attn_dcp_size 的读取，若在并行状态未正确初始化的上下文调用 set_mla_kv_buffer，可能引入新的异常点；实际调用路径中并行状态应已就绪，风险较低。
  • 缺少新增测试覆盖：PR 未添加独立回归测试，仅依赖既有 DCP 测试在特定参数下复现验证，后续难以防止类似口径问题再次引入。
• 影响：
  • 影响范围：仅影响 attn_dcp_size > 1（DCP 开启）且启用了 SGLANG_ENABLE_ASYNC_ASSERT 的 MLA 模型（如 DeepSeek 系列）KV 写入断言路径；非 DCP 场景行为完全不变。
  • 用户影响：修复了 DCP 小 KV 池下异步断言误报导致的运行失败，使该配置可以正常开启断言用于调试和 CI 检查。
  • 系统影响：每次 set_mla_kv_buffer 调用多一次 get_parallel() 属性访问，性能开销可忽略。
  • 团队影响：改动极小、风险低，维护者可直接合入，无需额外回归成本。
  • 风险标记：DCP 执行路径变更, 缺少新增测试覆盖

关联脉络

• PR #33448 [DCP] Bound a request by the aggregate KV pool, not one rank's share: 同属 DCP 下 KV 池边界语义的正确性修正：该 PR 从请求长度维度约束到聚合 KV 池，本 PR 修正 MLA 写路径 OOB 断言使用物理容量而非 DCP 逻辑容量的问题，两者都涉及 DCP 索引空间与物理空间的映射。