# PR #39487 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Fix][DCP] Localize widened KV ids in MLA retraction CPU backup/restore
- 合并时间：2026-09-15 13:33
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/39487

---

# 执行摘要

- 一句话：修复 DCP 下 MLA 重删 CPU 备份 / 恢复的 KV 索引越界崩溃
- 推荐动作：值得精读，尤其对维护 DCP / PD 分离 / HiCache / KV 内存池的同学。重点关注三处设计决策：① 把索引换算从 host pool 方法提升为 `layers/dcp/layout.py` 纯函数的动机与收益；② `_write_loc_dcp_span` + `write_loc_is_dcp_resolved` 这对契约如何避免统一池上的双重换算；③ 为什么可以安全删除「整加宽页」断言——理解 `seqlen - 1` 的 ragged 前缀与「起始页对齐」这个更弱前提的差别。若只关心结论，可记住两点：崩溃触发条件是 DCP>1 时用加宽 id 直接索引 rank 局部 KV buffer；修复方式是 `indices[dcp_rank::dcp_size] // dcp_size`。合并后建议跟进两件事：为 `maybe_dcp_kernel_indices` 补一个起始对齐的显式校验或更强的测试守护，以及排查作者提到的「恢复请求重复 token」独立缺陷。

# 功能与动机

Issue#38645 记录了线上的具体崩溃：在 PD-disaggregation decode 节点、DCP8、Kimi-K3（MLA + KDA 混合）配置下，decode 容量压力触发 retract_decode 后，调度器在 Req.offload_kv_cache -> TokenToKVPoolAllocator.get_cpu_copy -> MLATokenToKVPool.get_cpu_copy 的 kv_buffer[layer_id][chunk_indices] 处抛 CUDA device-side assert。Issue 明确给出的根因是索引空间不一致：「req_to_token stores DCP-widened logical KV indices」而「per-rank MLA kv_buffer stores only rank-local physical rows」，正常 MLA 读写路径都会先做 owner 过滤与 // dcp_size 折叠，只有重删 CPU 备份 / 恢复路径漏掉了这一步；Issue 还提示即使 id 仍在范围内，也可能「address the wrong physical row and silently copy/restore incorrect KV data」。PR body 补充：本 PR 同时是 Draft 验证 PR#38790 的正式落地版本。

# 实现拆解

1. **抽取统一的 DCP 索引换算入口（layers/dcp/layout.py）**：把原 `HostKVCache.maybe_dcp_kernel_indices` 方法（位于 `mem_cache/pool_host/base.py`）改写成模块级纯函数 `maybe_dcp_kernel_indices(indices, dcp_size, dcp_rank)`，与 `get_dcp_lens`、`filter_dcp_local_kv_indices`、`filter_dcp_local_chunk_kv_indices` 等既有 DCP 索引数学并列，`dcp_size == 1` 时恒等返回。这样做的原因有两个：一是新调用方 `mem_cache/memory_pool.py` 不属于 HiCache host pool 体系，方法式 API 无法复用；二是纯函数把依赖显式化为参数，避免全局状态耦合。

2. **放宽页对齐断言**：原实现带 `assert indices.numel() % self.dcp_size == 0`（要求整加宽页），但重删前缀只有 `seqlen - 1` 个 token、本身是 ragged 的，断言会误杀；新函数只依赖「run 起始页对齐」这一更弱前提，用 `indices[dcp_rank::dcp_size] // dcp_size` 的步长视图选择本 rank 槽位并折叠为物理行。相应地 `pool_host/mla.py` 的 `load_to_device_per_layer` 与 `backup_from_device_all_layer` 改为传入 `self.dcp_size`、`self.dcp_rank`。

3. **修复 PD 重删备份 / 恢复主路径（mem_cache/memory_pool.py）**：`MLATokenToKVPool.get_cpu_copy` 与 `load_cpu_copy` 在分块拷贝前各做一次 `maybe_dcp_kernel_indices(indices, self._write_loc_dcp_span, get_parallel().attn_dcp_rank)`。这里用 `_write_loc_dcp_span` 而不是裸 `dcp_size`，是为了沿用池上已有的 `write_loc_is_dcp_resolved` 契约：统一池里 `req_to_token` 存的已是物理行，此时 span 为 1、函数透传，天然规避二次换算；备份与恢复各自独立换算，因为恢复阶段 `req_to_token` 里已是重新分配后的新加宽 id，不能复用备份时的物理行。KDA/Mamba 的 `mamba_indices` 明确不做转换（其索引空间独立）。

4. **测试配套**：`test/registered/unit/mem_cache/test_decode_retraction_backup.py` 新增 `TestDcpRetractionBackup`，用 `_bare_mla_pool()` 构造只含 `kv_buffer` 的裸 MLA 池、用 `get_parallel().override(...)` 模拟 DCP=4/rank=1，覆盖 `test_restore_lands_on_new_owned_rows`（备份 + 恢复往返落到新的本 rank 行，未涉及行不变）与 `test_resolved_pool_takes_ids_as_rows`（`write_loc_is_dcp_resolved=True` 时 id 原样透传，防双重换算）；`test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py` 改为直接调用纯函数，删除 `_bare_pool` 与随断言一起失效的 `test_ragged_run_is_rejected`。

5. **演进与验证**：9 个 commit 几乎全是作者自迭代——先把布尔掩码换成步长视图（1fe7d91），再尝试在两个 cpu-copy 调用点内联（6397112、747510c），随后改为共享 helper（dec6560、3d8bd25），最后把测试用例从 `test_retraction_mamba_backup` 迁到 `test_decode_retraction_backup`（60efcfe、40dba53）。CI 侧 `/rerun-test` 覆盖 `test_hicache_dcp_host_pool.py`、`test_decode_retraction_backup.py`、`test_unified_radix_cache_kl_dcp.py`、`test_kimi_linear_pd_dcp4.py`、`test_kimi_linear_pd_dcp4_dspark.py`，五个目标全部通过。

关键文件：
- `python/sglang/srt/mem_cache/memory_pool.py`（模块 内存池；类别 source；类型 core-logic；符号 MLATokenToKVPool.get_cpu_copy, MLATokenToKVPool.load_cpu_copy）: 修复落点：MLATokenToKVPool.get_cpu_copy / load_cpu_copy 现在在索引物理 kv_buffer 前先做 DCP 本地化，并靠 _write_loc_dcp_span 兼容统一池的已物理化 id。
- `python/sglang/srt/layers/dcp/layout.py`（模块 索引换算；类别 source；类型 core-logic；符号 maybe_dcp_kernel_indices）: 新增共享纯函数 maybe_dcp_kernel_indices，把「加宽逻辑槽位 -> 本 rank 物理行」的换算从 HiCache host pool 提升为 DCP 索引数学模块的公共能力。
- `test/registered/unit/mem_cache/test_decode_retraction_backup.py`（模块 重删测试；类别 test；类型 test-coverage；符号 _bare_mla_pool, _dcp, TestDcpRetractionBackup, test_restore_lands_on_new_owned_rows）: 新增 TestDcpRetractionBackup，用裸 MLA 池 + DCP override 复现「加宽 id 当行号」的两个方向：往返落到新行、以及已物理化池的透传。
- `python/sglang/srt/mem_cache/pool_host/base.py`（模块 宿主缓存；类别 source；类型 refactor；符号 HostKVCache.maybe_dcp_kernel_indices）: HostKVCache.maybe_dcp_kernel_indices 被删除（迁往 layers/dcp/layout.py），连带移除「整加宽页」断言，是本 PR 语义变化的来源之一。
- `python/sglang/srt/mem_cache/pool_host/mla.py`（模块 宿主缓存；类别 source；类型 dependency-wiring；符号 MLATokenToKVPoolHost.load_to_device_per_layer, MLATokenToKVPoolHost.backup_from_device_all_layer）: MLATokenToKVPoolHost 的两个传输入口 load_to_device_per_layer、backup_from_device_all_layer 改为调用模块级纯函数并显式传入 self.dcp_size / self.dcp_rank。
- `test/registered/unit/mem_cache/test_hicache_dcp_host_pool.py`（模块 缓存测试；类别 test；类型 test-coverage；符号 _bare_pool, test_ragged_run_is_rejected, test_matches_owner_rule_on_merged_unordered_pages）: HiCache DCP host pool 单测改为直接调用纯函数，删除 _bare_pool 与随断言一起失效的 test_ragged_run_is_rejected，是断言移除的测试侧证据。

关键符号：maybe_dcp_kernel_indices, MLATokenToKVPool.get_cpu_copy, MLATokenToKVPool.load_cpu_copy, MLATokenToKVPoolHost.load_to_device_per_layer, MLATokenToKVPoolHost.backup_from_device_all_layer, TestDcpRetractionBackup.test_restore_lands_on_new_owned_rows, TestDcpRetractionBackup.test_resolved_pool_takes_ids_as_rows

## 关键源码片段

### `python/sglang/srt/mem_cache/memory_pool.py`

修复落点：MLATokenToKVPool.get_cpu_copy / load_cpu_copy 现在在索引物理 kv_buffer 前先做 DCP 本地化，并靠 _write_loc_dcp_span 兼容统一池的已物理化 id。

```python
# python/sglang/srt/mem_cache/memory_pool.py :: MLATokenToKVPool

def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
    # 重删备份：req_to_token 给的是 DCP 加宽逻辑 id，kv_buffer 只存本 rank 的物理行。
    # _write_loc_dcp_span 在统一池已物理化时为 1（write_loc_is_dcp_resolved），
    # 此时 maybe_dcp_kernel_indices 原样透传，天然避免二次换算。
    indices = maybe_dcp_kernel_indices(
        indices, self._write_loc_dcp_span, get_parallel().attn_dcp_rank
    )
    current_platform.synchronize()
    kv_cache_cpu = []
    chunk_size = self.cpu_offloading_chunk_size
    for layer_id in range(self.layer_num):
        kv_cache_cpu.append([])
        if self.kv_buffer[layer_id].shape[0] == 0:
            continue
        # 分块拷贝，chunk_indices 这里已经是本 rank 的物理行号
        for i in range(0, len(indices), chunk_size):
            chunk_indices = indices[i : i + chunk_size]
            kv_cpu = self.kv_buffer[layer_id][chunk_indices].to(
                "cpu", non_blocking=True
            )
            kv_cache_cpu[-1].append(kv_cpu)
    current_platform.synchronize()
    return kv_cache_cpu


def load_cpu_copy(
    self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
    # 恢复同样独立换算：重删后重新分配的是「新的加宽 id」，
    # 不能复用备份时算出的物理行号。KDA/Mamba 状态的 mamba_indices 不做转换。
    indices = maybe_dcp_kernel_indices(
        indices, self._write_loc_dcp_span, get_parallel().attn_dcp_rank
    )
    current_platform.synchronize()
    chunk_size = self.cpu_offloading_chunk_size
    for layer_id in range(self.layer_num):
        if self.kv_buffer[layer_id].shape[0] == 0:
            continue
        for i in range(0, len(indices), chunk_size):
            chunk_indices = indices[i : i + chunk_size]
            kv_cpu = kv_cache_cpu[layer_id][i // chunk_size]
            # 换算对称性的显式检查：备份与恢复两侧必须得到同样多的行
            assert kv_cpu.shape[0] == len(chunk_indices)
            kv_chunk = kv_cpu.to(self.kv_buffer[0].device, non_blocking=True)
            self.kv_buffer[layer_id][chunk_indices] = kv_chunk
    current_platform.synchronize()

```

### `python/sglang/srt/layers/dcp/layout.py`

新增共享纯函数 maybe_dcp_kernel_indices，把「加宽逻辑槽位 -> 本 rank 物理行」的换算从 HiCache host pool 提升为 DCP 索引数学模块的公共能力。

```python
# python/sglang/srt/layers/dcp/layout.py
# 该模块原本只放 DCP 的纯索引数学（get_dcp_lens、filter_dcp_local_kv_indices 等）。
# 本 PR 把原本是 HostKVCache 方法的换算提升到这里，
# 让 HiCache host pool 与 MLA 池的重删 CPU 备份 / 恢复共用同一份语义。

def maybe_dcp_kernel_indices(
    indices: torch.Tensor, dcp_size: int, dcp_rank: int
) -> torch.Tensor:
    """Widened logical slots -> this rank's physical rows.

    Owner rule: slot % dcp_size == dcp_rank, row = slot // dcp_size. The run
    starts page-aligned, so a strided view selects the owned slots without a mask.
    """
    if dcp_size == 1:
        return indices  # DCP=1 是恒等情形，直接返回原张量避免多余拷贝
    # 关键取舍：去掉原实现的「必须整加宽页」断言（重删前缀是 seqlen - 1，天然 ragged），
    # 只依赖更弱的前提「run 起始页对齐」，用步长视图代替布尔掩码：
    # 选中 indices[dcp_rank]、indices[dcp_rank + dcp_size]、...，再整体 // dcp_size
    return indices[dcp_rank::dcp_size] // dcp_size


# 同模块中的对照实现：读取侧只做选择、仍保持加宽 id，由调用方再折叠
# def filter_dcp_local_kv_indices(kv_indices):
# parallel = get_parallel()
# if parallel.dcp_enabled:
# kv_indices = kv_indices[kv_indices % parallel.dcp_size == parallel.dcp_rank]
# return kv_indices

```

### `test/registered/unit/mem_cache/test_decode_retraction_backup.py`

新增 TestDcpRetractionBackup，用裸 MLA 池 + DCP override 复现「加宽 id 当行号」的两个方向：往返落到新行、以及已物理化池的透传。

```python
# test/registered/unit/mem_cache/test_decode_retraction_backup.py
DCP_SIZE = 4
DCP_RANK = 1
DCP_ROWS = 8


def _bare_mla_pool() -> MLATokenToKVPool:
    # 绕过构造函数：只挂上 kv_buffer / layer_num / 分块大小，
    # 每行以 arange + 100 * layer 填充，方便断言「KV 被搬到了哪一行」。
    pool = object.__new__(MLATokenToKVPool)
    pool.layer_num = 2
    pool.cpu_offloading_chunk_size = 3
    pool.kv_buffer = [
        (torch.arange(DCP_ROWS, dtype=torch.float32) + 100 * layer).view(DCP_ROWS, 1, 1)
        for layer in range(pool.layer_num)
    ]
    return pool


def _dcp():
    # 通过 runtime_context 覆盖并行配置，模拟 DCP=4 且本 rank 为 1
    return get_parallel().override(
        dcp_enabled=True, attn_dcp_size=DCP_SIZE, attn_dcp_rank=DCP_RANK
    )


class TestDcpRetractionBackup(unittest.TestCase):
    """`req_to_token` names KV slots in the widened DCP id space while
    `kv_buffer` holds only this rank's rows; a widened id used as a row index
    reads past the buffer or copies another token's row."""

    def test_restore_lands_on_new_owned_rows(self):
        pool = _bare_mla_pool()
        before = [buf.clone() for buf in pool.kv_buffer]
        old_widened = torch.arange(0, 12, dtype=torch.int64)   # 备份来源的加宽 id
        new_widened = torch.arange(12, 24, dtype=torch.int64)  # 重删后新分配的加宽 id

        with _dcp():
            # 备份与恢复各自独立换算：旧 id 读出本 rank 行，写回新 id 对应的本 rank 行
            pool.load_cpu_copy(pool.get_cpu_copy(old_widened), new_widened)

        old_rows = old_widened[DCP_RANK::DCP_SIZE] // DCP_SIZE
        new_rows = new_widened[DCP_RANK::DCP_SIZE] // DCP_SIZE
        self.assertEqual(new_rows.tolist(), [3, 4, 5])
        untouched = torch.tensor(
            [r for r in range(DCP_ROWS) if r not in new_rows.tolist()]
        )
        for layer in range(pool.layer_num):
            torch.testing.assert_close(
                pool.kv_buffer[layer][new_rows], before[layer][old_rows]
            )  # 数据确实搬到了新的本 rank 行
            torch.testing.assert_close(
                pool.kv_buffer[layer][untouched], before[layer][untouched]
            )  # 未涉及的行保持原样，排除误写

    def test_resolved_pool_takes_ids_as_rows(self):
        pool = _bare_mla_pool()
        pool.write_loc_is_dcp_resolved = True  # 统一池：req_to_token 中已是物理行
        rows = torch.arange(DCP_ROWS, dtype=torch.int64)

        with _dcp():
            kv_cpu = pool.get_cpu_copy(rows)

        for layer in range(pool.layer_num):
            # 双重换算防护：DCP 打开时也要原样透传
            torch.testing.assert_close(torch.cat(kv_cpu[layer]), pool.kv_buffer[layer])

```

# 评论区精华

本 PR 的评审讨论非常轻量：`review_comments_count = 0`，唯一一条 review 来自 ShangmingCai 的 APPROVED ——「This minor fix looks good to me.」。Issue 线程中的两条评论都是流程性的：作者用 `/rerun-test` 指定 5 个测试文件，github-actions[bot] 回报 ubuntu-latest、1-gpu-5090、4-gpu-b200、8-gpu-b200 四个 runner 上的用例全部 ✅。真正有价值的「讨论」发生在 PR body 与 commit message 里：作者主动说明了断言被删的原因（重删前缀是 `seqlen - 1` 的 ragged run）、`_write_loc_dcp_span` 与 `write_loc_is_dcp_resolved` 的契约、以及一个刻意不修的问题——强制重删下约 4/48 的恢复请求会退化成 token 重复，且 `--dcp-size 1` 同样复现，因此判定为独立问题而非 DCP 缺陷。

- 索引换算逻辑该放在哪一层 (design): 确定为 layers/dcp/layout.py 的模块级纯函数 maybe_dcp_kernel_indices(indices, dcp_size, dcp_rank)，dcp_size/rank 显式传参，HiCache host pool 与 MLA 池共用。
- 删除「整加宽页」断言的正确性边界 (correctness): 断言被接受移除，换取重删路径可用；代价是函数对「起始对齐」的假设从强制校验退化为文档契约，未来误用会静默选错槽位。
- 统一池上的双重换算防护 (design): 以池属性作为开关，单一换算点同时服务「加宽 id」与「已物理化 id」两种池，避免二次折叠。
- DCP 回归用例的归属测试文件 (testing): 最终落在 test/registered/unit/mem_cache/test_decode_retraction_backup.py 的 TestDcpRetractionBackup，与既有的 MHA/draft 备份用例共用同一裸池构造风格。
- 评审结论与验证强度 (other): 代码评审轻量通过；深度验证证据主要来自 PR body 中的 4xB200 实测（main 约 12 次重删即 assert，本分支 140-200 次重删无 assert）。

# 风险与影响

- 风险：
 1. **静默错误行风险（最高）**：把加宽 id 当行号使用时，只要 id 落在此 rank 的行范围内就不会崩溃，而是复制 / 写回错误 token 的 KV（Issue#38645 原文明确警告）。本 PR 修掉了该路径，但此类错误缺乏端到端可观测性，主要靠 `test_restore_lands_on_new_owned_rows` 这类单测与 Kimi-Linear-48B 的 logprob 对比守护。
 2. **页对齐断言移除的边界**：新函数删掉了 `indices.numel() % dcp_size == 0` 的 fail-fast 断言，改为口头契约「run starts page-aligned」。若未来有调用方传入起始未对齐的 run（例如带 offset 的 chunk 视图），步长视图会静默选错槽位而不报错；`test_hicache_dcp_host_pool.py` 中原本守护这一点的 `test_ragged_run_is_rejected` 也被删除，回归网出现缺口。
 3. **对池契约的隐式依赖**：`get_cpu_copy` / `load_cpu_copy` 的正确性依赖 `_write_loc_dcp_span`（配合 `write_loc_is_dcp_resolved`）与实际索引空间一致。`MLATokenToKVPool` 的子类（如 `MLATokenToKVPoolFP4`）若未正确维护该属性，会出现误换算或属性缺失；`get_parallel().attn_dcp_rank` 取的是全局并行上下文而非池自身 rank，两者不一致时会算错行。
 4. **未覆盖范围**：Mamba/KDA 的 `mamba_indices` 不做转换，`load_cpu_copy` 中「恢复后重复 token 退化」问题被明确留作后续；`assert kv_cpu.shape[0] == len(chunk_indices)` 是换算对称性的唯一显式检查，一旦备份 / 恢复两侧 span 判断不同会在此处断言失败而非静默错位。
 5. **CI 信号混杂**：PR body 中的 CI 状态槽位显示 PR Test (Base)/(Extra)/(AMD ROCm 10) 三个 job 为 ❌，而针对性 rerun 的 5 个测试全绿，合并前需要确认失败项与本次变更无关。
 - 影响：用户侧：PD 分离 + DCP > 1 + MLA（Kimi 系列）的 decode 节点在容量压力触发重删时不再崩溃，请求可继续完成；作者在 4xB200 上实测 main 分支约 12 次重删即 assert，本分支 140–200 次重删无 assert，per-token logprob 与不重删参考值在噪声范围内。系统侧：`layers/dcp/layout.py` 成为 DCP 索引换算的公共入口，HiCache host pool 与 KV 池的语义一致性问题被收敛到单点；`memory_pool.py` 处于热路径上游，但只在备份 / 恢复时执行，对正常 decode 吞吐无影响（转发路径不受影响）。团队侧：这是「加宽逻辑索引 vs rank 物理行」这一 DCP 核心约定的一次契约补齐，后续任何新增访问 `kv_buffer` 的旁路路径都应先过 `maybe_dcp_kernel_indices`。测试层面新增的 `TestDcpRetractionBackup` 提供了可复用的裸池构造方法（`_bare_mla_pool` + `_dcp`），便于继续补 DCP 场景用例。
 - 风险标记：静默错行风险 , 页对齐断言移除 , 双重换算依赖池契约 , KDA/Mamba 状态未覆盖 , 已知 token 重复退化未修 , PR 主 CI 槽位显示失败

# 关联脉络

- PR #38790 [BugFix][DCP] Localize MLA KV indices for retraction CPU backup/restore: 同一 Issue#38645 的 Draft 验证 PR，本 PR body 明确说明是其正式落地版本；两者修复目标、索引语义与验证矩阵（v0.5.17 A/B、v0.5.19 容量边界扫描）完全对应。
- PR #35233 [AMD] Fix registered HiCache host pointer aliases: 同样触及 HiCache host pool 的指针 / 索引域问题（pool_host/common.py、pool_host/mla.py 等），与本 PR 同属「host pool 与 device/ 逻辑索引空间错配」问题族。
- PR #38941 [Fix] Merge adjacent KV-row frees so a mid-page split under DCP cannot double-free: 同为 DCP 场景下 KV 行索引 / 页内接缝的正确性修复，改动 mem_cache 公共层（base_prefix_cache.py、common.py），与本 PR 一起构成 DCP KV 索引语义的补强序列。