# PR #29161 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[Fix]: Defer DSA MLA CP KV gather for fp8 trtllm prefill in PD mode
- 合并时间：2026-07-01 15:20
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/29161

---

# 执行摘要

- 一句话：修复 DSA fp8 trtllm CP 下 KV gather 形状不匹配崩溃
- 推荐动作：本 PR 修复了关键的 shape mismatch 崩溃，设计思路清晰（延迟 gather 直到形状正确）。建议精读 `_should_defer_dsa_cp_kv_gather` 的条件判断和 `_all_gather_dsa_trtllm_fp8_kv` 中拼接 - 通信 - 切分的实现。合并后应关注非连续张量的潜在问题。

# 功能与动机

PR body 明确指出：在 PD 分离式 prefill 服务中，DSA prefill + CP + fp8 KV + trtllm backend 是一个常用配置，但 original code 在 `forward_absorb_prepare` 中过早执行了 KV all-gather，而 fp8 trtllm 路径的 RoPE 被推迟到 `_forward_trtllm` 中，导致 `k_rope` 的 batch 维度与局部 `positions` 不匹配，触发 `mla_rope_quantize_fp8: Check failed: k_rope_in.size(0) == nnz` 错误。此修复是必需的。

# 实现拆解

1. 在 `forward_mla.py` 中添加 `_should_defer_dsa_cp_kv_gather` 函数，该函数接收 `dsa_prefill_cp` 和 `fuse_rope_for_trtllm_mla` 布尔参数，仅当两者均为 True 时返回 True，标识需要延迟 CP KV gather。
2. 在 `forward_absorb_prepare` 中缓存 `_fuse_rope_for_trtllm_mla(forward_batch)` 的返回值，避免重复调用。然后根据新条件：如果 CP 启用且非延迟条件，照常执行 `rebuild_cp_kv_cache`；否则跳过，使 KV 保持本地 shard。
3. 在 `dsa_backend.py` 中新增 `_all_gather_dsa_trtllm_fp8_kv` 函数，将 `k` 和 `k_rope` 拼接后以 uint8 类型通过 `cp_all_gather_rerange_output` 进行跨 rank 通信，再按原维度切分返回。
4. 在 `_forward_trtllm` 中，先对 `forward_batch.positions` 进行 CP 重排（`cp_split_and_rebuild_position`），然后调用 `mla_quantize_and_rope_for_fp8` 完成 fused RoPE quantization；若为 CP 模式且需要保存 KV，则调用新函数完成 all-gather。
5. 针对 in-seq-split（zigzag）布局，在传入 `trtllm_batch_decode_with_kv_cache_mla` 之前，根据 `attn_cp_metadata.zigzag_index` 对 `seq_lens` 重新排序，确保与本地 shard 顺序对齐。

关键文件：
- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`（模块 前向逻辑；类别 source；类型 data-contract；符号 _should_defer_dsa_cp_kv_gather）: 添加延迟 CP KV gather 的条件判断和门控逻辑，缓存 fuse_rope_for_trtllm_mla 值
- `python/sglang/srt/layers/attention/dsa_backend.py`（模块 注意力后端；类别 source；类型 core-logic；符号 _all_gather_dsa_trtllm_fp8_kv）: 实现延迟 gather 的核心通信函数，并在 _forward_trtllm 中调用，同时修复 in-seq-split 的 seq_lens 重排

关键符号：_all_gather_dsa_trtllm_fp8_kv, _should_defer_dsa_cp_kv_gather

## 关键源码片段

### `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`

添加延迟 CP KV gather 的条件判断和门控逻辑，缓存 fuse_rope_for_trtllm_mla 值

```python
def _should_defer_dsa_cp_kv_gather(
    *,
    dsa_prefill_cp: bool,
    fuse_rope_for_trtllm_mla: bool,
) -> bool:
    # 当 DSA prefill CP 和 trtllm MLA 融合 RoPE 同时启用时，延迟 KV gather
    return dsa_prefill_cp and fuse_rope_for_trtllm_mla

# 在 forward_absorb_prepare 中，缓存并判断
dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
mla_prefill_cp = mla_use_prefill_cp(forward_batch)
fuse_rope_for_trtllm_mla = self._fuse_rope_for_trtllm_mla(forward_batch)
defer_kv_gather_until_after_rope = _should_defer_dsa_cp_kv_gather(
    dsa_prefill_cp=dsa_prefill_cp,
    fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
)
if (dsa_prefill_cp or mla_prefill_cp) and not defer_kv_gather_until_after_rope:
    # 只有不需要延迟时才执行 CP KV 重建（bf16 路径保持不变）
    k_nope, k_pe = self.rebuild_cp_kv_cache(
        latent_cache, forward_batch, k_nope, k_pe
    )

```

### `python/sglang/srt/layers/attention/dsa_backend.py`

实现延迟 gather 的核心通信函数，并在 _forward_trtllm 中调用，同时修复 in-seq-split 的 seq_lens 重排

```python
def _all_gather_dsa_trtllm_fp8_kv(
    forward_batch: ForwardBatch,
    k: torch.Tensor,
    k_rope: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    # 拼接 k 和 k_rope，以 uint8 进行 all-gather，再恢复 dtype
    kv_lora_rank = k.shape[-1]
    qk_rope_head_dim = k_rope.shape[-1]
    kv_dtype = k.dtype
    kv = torch.cat((k, k_rope), dim=-1).view(torch.uint8)
    kv = cp_all_gather_rerange_output(
        kv,
        get_parallel().attn_cp_size,
        forward_batch,
        torch.cuda.current_stream(),
    ).view(kv_dtype)
    # 按原始维度拆分并返回（返回的是视图，非连续，下游需注意）
    return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1)

# 在 _forward_trtllm 中，RoPE 量化后执行
if save_kv_cache and dsa_use_prefill_cp(forward_batch):
    k, k_rope = _all_gather_dsa_trtllm_fp8_kv(forward_batch, k, k_rope)

```

# 评论区精华

- gemini-code-assist 指出 `_all_gather_dsa_trtllm_fp8_kv` 返回的 split 切片是非连续的，下游 CUDA/Triton 内核可能因 `.view()` 失败或静默数据损坏，建议添加 `.contiguous()`。
- JustinTong0323 在 4×B300 上验证：round-robin-split 修复后 GSM8K 准确率从 30% 恢复至 94%，停止率 96%；但 in-seq-split 仍为 0% 准确率，需要额外对齐 zigzag 布局。
- Dovis01 随后提交了 seq_lens 重排修复，Justin 重新验证后 in-seq-split 准确率达到 94%，确认所有组合正常工作。

- Non-contiguous split tensors (correctness): 未在 PR 中直接修改，但开发者已知悉；可能需要后续追加 .contiguous()
- round-robin-split 数值验证通过 (correctness): 修复有效，可合并
- in-seq-split 数值错误与修复 (correctness): 修复后 in-seq-split 准确率恢复至 94%，与 round-robin-split 一致

# 风险与影响

- 风险：
 1. 非连续张量风险：`split` 返回的视图未被 `.contiguous()` 保护，若下游内核隐式使用 `.view()` 或假设连续内存，可能导致静默数据损坏（评论中已指出，但本 PR 未修改）。
 2. 仅影响 fp8+trtllm+CP 组合，bf16 CP 路径不受影响。
 3. in-seq-split 的 seq_lens 重排依赖 `attn_cp_metadata`，若该元数据格式变更或为空，可能导致索引越界。
- 影响：
 - 用户：修复了 PD 分离 prefill 部署中 fp8 trtllm 配置下的崩溃，GSM8K 任务准确率从 30% 恢复至 94%，显著提升稳定性。
 - 系统：延迟 gather 使通信发生在 RoPE 之后，减少了形状不匹配风险，但每个 rank 需保留本地 KV 直到 gather，可能小幅增加内存峰值。
 - 团队：设计模式（条件门控延迟执行）值得推广；需后续跟进非连续张量问题。
 - 风险标记：非连续张量风险 , in-seq-split 曾需要额外修复 , 核心路径变更

# 关联脉络

- 暂无明显关联 PR