# PR #32785 完整报告

- 仓库：`sgl-project/sglang`
- 标题：fix: avoid piecewise prefill graph for trtllm_mla
- 合并时间：2026-08-08 16:00
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/32785

---

# 执行摘要

- 一句话：trtllm_mla 禁用 piecewise prefill 图，修复 TTFT 回归
- 推荐动作：值得精读。PR 虽小，但展示了 " 默认值升级 + 后端例外 + 显式覆盖 " 的三层优先级设计，是分析 CUDA graph 与 MLA 注意力后端兼容性规则的好入口。建议结合 #30889（回归源头）和 #32288（同区域正交修复）一起阅读，能快速理解 sglang 中 prefill CUDA graph backend 选择的演进逻辑。

# 功能与动机

修复 issue #32655：开启 piecewise prefill CUDA graph 后，nvidia/Kimi-K2.6-NVFP4 在 4x GB300/TP4 固定 1K 输入 /1K 输出、concurrency 1 负载下吞吐下降约 3.5%（186.70 vs 180.11 tok/s），TTFT 从 122.68 ms 升至 344.58 ms，TPOT 基本不变（5.24 vs 5.21 ms）。根因是 #30889 将默认 prefill 后端从 breakable 升级为 tc_piecewise 后，Kimi-K2.6 自动选择的 trtllm_mla 后端在捕获时走了 FlashInfer paged-MLA fallback，需要整段转换 FP8 KV cache。PR body 明确指出本 PR 修复的是默认 backend 选择，与 #32288（修复 fallback 状态泄漏到 speculative verify）正交。

# 实现拆解

1. 变更入口：`python/sglang/srt/server_args.py` 中 `ServerArgs._apply_cuda_graph_compatibility()` 的升级分支。原本当 prefill backend 为 BREAKABLE 且模型位于多模态 allowlist 时无条件升级为 TC_PIECEWISE，现在增加 `self._resolved_attention_backends()[0] != "trtllm_mla"` 守卫。
2. 行为分支：trtllm_mla 被排除后保持 BREAKABLE，随后落入 `_disable_breakable_cudagraph_if_incompatible()`，由 `use_mla_backend` 检测把 prefill CUDA graph 置为 DISABLED，恢复 #30889 之前的 eager TRT-LLM ragged-MLA 路径；decode CUDA graph 选择完全不受影响。
3. 显式覆盖优先：`_apply_cuda_graph_compatibility` 开头对 `(Phase.PREFILL, "backend") in self._cuda_graph_config_locked` 直接返回，因此用户显式指定 tc_piecewise（如 `--enforce-piecewise-cuda-graph` 或 JSON 配置）时，即使注意力后端是 trtllm_mla 也保持 TC_PIECEWISE，不改变用户意图。
4. 测试配套：`test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py` 新增两个回归测试（默认 trtllm_mla 保持 breakable 并被禁用、显式 tc_piecewise 覆盖默认规则），并给原有升级测试补上 `_resolved_attention_backends` 的 mock（`("fa3", "fa3")`），避免新守卫导致旧用例误触发。
5. 验证情况：单测 9 个全部通过，pre-commit 通过；PR 未包含完整 4x GB300 性能 A/B，性能结论依赖 issue #32655 的数据。CI Extra 有一次失败记录，通过 `/tag-and-rerun-ci` 重跑后合并。

关键文件：
- `python/sglang/srt/server_args.py`（模块 服务配置；类别 source；类型 core-logic；符号 _apply_cuda_graph_compatibility）: 核心逻辑修改文件：在 `_apply_cuda_graph_compatibility()` 多模态 prefill 图升级条件中新增 trtllm_mla 排除守卫，是本次修复的关键实现。
- `test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py`（模块 预填充图；类别 test；类型 test-coverage；符号 test_trtllm_mla_stays_on_breakable_and_is_disabled_by_compatibility, test_explicit_tc_piecewise_overrides_trtllm_mla_default）: 新增两个回归测试守住 trtllm_mla 的默认行为与显式覆盖语义，并调整原有升级测试的 mock，避免新守卫带来误判。

关键符号：_apply_cuda_graph_compatibility, _disable_breakable_cudagraph_if_incompatible, _resolved_attention_backends, test_trtllm_mla_stays_on_breakable_and_is_disabled_by_compatibility, test_explicit_tc_piecewise_overrides_trtllm_mla_default, test_supported_multimodal_model_upgrades_default_to_tc_piecewise

## 关键源码片段

### `python/sglang/srt/server_args.py`

核心逻辑修改文件：在 `_apply_cuda_graph_compatibility()` 多模态 prefill 图升级条件中新增 trtllm_mla 排除守卫，是本次修复的关键实现。

```python
def _apply_cuda_graph_compatibility(self):
    """根据运行配置自动关闭不兼容的 prefill CUDA graph。

    规则按 backend 拆分：`TcPiecewise` 与 `Breakable` 约束不同；
    仅当用户没有显式指定 prefill backend 时才生效。
    """
    if (Phase.PREFILL, "backend") in self._cuda_graph_config_locked:
        # 用户显式锁定 backend（如 `--enforce-piecewise-cuda-graph`），直接返回
        return

    # Breakable 是通用 CUDA 默认值，但与多模态 prefill 不兼容；
    # 白名单中的模型已验证可在 tc_piecewise 下运行 decoder prefill。
    if (
        self.cuda_graph_config.prefill.backend == Backend.BREAKABLE
        and self.get_model_config().is_multimodal_piecewise_cuda_graph_supported
        # 关键修复：trtllm_mla（Blackwell 上 Kimi-K2.6 自动选择）继续保持
        # breakable 路径。当前 MLA 兼容性规则会禁用 prefill CUDA graph，
        # 从而避免 tc_piecewise 落入 FlashInfer paged-MLA fallback 时
        # 整段转换 FP8 KV cache 带来的 TTFT 回归。
        and self._resolved_attention_backends()[0] != "trtllm_mla"
    ):
        logger.info(
            "Using tc_piecewise CUDA graph for validated multimodal "
            "decoder prefill."
        )
        self.cuda_graph_config.prefill.backend = Backend.TC_PIECEWISE

    # 按最终选定的 backend 分别做兼容性检查
    if self.cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE:
        # torch.compile + piecewise 受 dynamo 限制，规则较多
        self._disable_tc_piecewise_cudagraph_if_incompatible()
    elif self.cuda_graph_config.prefill.backend == Backend.BREAKABLE:
        # trtllm_mla 的 MLA 检测在这里触发，最终将 prefill 图置为 DISABLED
        self._disable_breakable_cudagraph_if_incompatible()
    elif self.cuda_graph_config.prefill.backend == Backend.FULL:
        self._disable_full_prefill_cudagraph_if_incompatible()

```

### `test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py`

新增两个回归测试守住 trtllm_mla 的默认行为与显式覆盖语义，并调整原有升级测试的 mock，避免新守卫带来误判。

```python
def test_trtllm_mla_stays_on_breakable_and_is_disabled_by_compatibility(self):
    args = ServerArgs(model_path="dummy")
    args.model_config = SimpleNamespace(
        is_multimodal_piecewise_cuda_graph_supported=True
    )
    # 默认 prefill backend 为 `BREAKABLE`
    args.cuda_graph_config = CudaGraphConfig(
        prefill=PhaseConfig(backend=Backend.BREAKABLE)
    )
    args._cuda_graph_config_locked = set()

    # 模拟 Blackwell 上 Kimi-K2.6 自动解析出 trtllm_mla 的场景
    with (
        patch.object(
            args,
            "_resolved_attention_backends",
            return_value=("trtllm_mla", "trtllm_mla"),
        ),
        patch.object(args, "use_mla_backend", return_value=True),
    ):
        args._apply_cuda_graph_compatibility()

    # 期望：不升级到 tc_piecewise，而是留在 breakable 分支后由 MLA
    # 兼容性规则把 prefill 图置为 DISABLED，走 eager TRT-LLM 路径
    self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED)


def test_explicit_tc_piecewise_overrides_trtllm_mla_default(self):
    args = ServerArgs(model_path="dummy")
    # 用户显式指定 `TC_PIECEWISE` 并锁定 backend 键
    args.cuda_graph_config = CudaGraphConfig(
        prefill=PhaseConfig(backend=Backend.TC_PIECEWISE)
    )
    args._cuda_graph_config_locked = {(Phase.PREFILL, "backend")}

    with patch.object(
        args,
        "_resolved_attention_backends",
        return_value=("trtllm_mla", "trtllm_mla"),
    ):
        args._apply_cuda_graph_compatibility()

    # 显式选择优先于默认规则，仍保持 `TC_PIECEWISE`
    self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.TC_PIECEWISE)

```

# 评论区精华

仓库中没有针对 patch 的 formal review 评论。有效的讨论集中在 PR 评论和 body：nvpohanh 询问是否合入（"should we merge this? thanks"），mickqian 以 `/tag-and-rerun-ci` 重跑 CI 后合入。PR body 中作者强调了与 #32288 的边界：本 PR 只改默认 backend 选择，不触碰 fallback 状态逻辑，避免两处修复相互干扰。gemini-code-assist 的提示仅为工具下线声明，不构成技术讨论。

- 与 #32288 的修复边界 (design): 不触碰 fallback 状态逻辑，只改 backend 默认值选择；后续 BCG 支持原生 MLA 后可从 breakable 侧直接放开。
- 合并确认与 CI 状态 (question): CI 重跑后由作者合入，无进一步技术争议。

# 风险与影响

- 风险：时序依赖：`_apply_cuda_graph_compatibility()` 新增调用 `_resolved_attention_backends()`，要求注意力后端在此时已完成解析；单测用 mock 覆盖，真实 B300 启动路径已验证，但后续若重构解析时序需同步关注。行为回退：trtllm_mla 多模态模型继续禁用 prefill CUDA graph，prefill 回到 eager 路径，这是有意的回退，后续 BCG 原生 MLA 支持落地后可放开。性能验证缺口：PR 自身未附 4x GB300 完整 A/B 数据，仅依赖 issue 数据，修复收益需用户自行复测。影响面窄：仅影响 allowlist 中自动解析为 trtllm_mla 的模型（主要为 Blackwell 上的 Kimi-K2.6），其他多模态模型仍升级到 tc_piecewise，显式指定不改变。
- 影响：用户侧：Kimi-K2.6/trtllm_mla 用户恢复 #30889 之前的 eager TRT-LLM prefill 路径，TTFT 回到约 122 ms 量级，吞吐回升约 3.5%。系统侧：无新增配置项、无 schema 变更，decode CUDA graph 与显式 tc_piecewise 行为不变，CLI 兼容性保持。团队侧：在代码注释中明确了 trtllm_mla 的过渡状态（"preferred future path: breakable"），为后续 BCG 原生 MLA 支持提供清晰的放开点，并与 server_args 配置体系的重构方向衔接。
- 风险标记：默认路径行为回退 , 缺 4x GB300 完整 A/B 验证 , 依赖 attention backend 解析时序 , 仅单测覆盖，无端到端回归测试

# 关联脉络

- PR #30889 （源 PR，标题未在材料中提供）默认 prefill CUDA graph 升级为 tc_piecewise: 本 PR 修复的回归源头：#30889 将多模态 allowlist 模型默认 prefill 后端从 breakable 升级为 tc_piecewise，导致 trtllm_mla 落入 paged-MLA fallback。
- PR #32288 （PR body 提及）修复 stale fallback 状态泄漏到 speculative verify: 与本 PR 正交，同属 prefill CUDA graph fallback 路径的修复，PR body 明确划清两者边界。
- PR #33887 config: retire ServerArgs.derive; per-runner values are constructor arguments: 后续对 server_args.py 配置体系的重构，与本 PR 修改的 `_apply_cuda_graph_compatibility` 所在区域持续演进，需保持 " 默认值升级 + 后端例外 + 显式覆盖 " 语义。