执行摘要

• 一句话：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 排除守卫，是本次修复的关键实现。

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，避免新守卫带来误判。

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 所在区域持续演进，需保持"默认值升级 + 后端例外 + 显式覆盖"语义。