执行摘要

• 一句话：TRITON_MLA 启用 FULL CUDAGraph
• 推荐动作：建议精读。该 PR 展示了一个极简但高效的优化模式：通过覆写 MetadataBuilder 的 _cudagraph_support 即可启用 FULL CUDAGraph，收益显著且风险低。对于其他使用 MLA 或类似自定义 attention backends 的开发者具有参考价值。

功能与动机

MLACommonMetadataBuilder 默认将 _cudagraph_support 设为 NEVER，使得 decode 阶段只能使用 PIECEWISE 模式，unified_mla_attention_with_output 算子无法被 FULL CUDAGraph 捕获，导致每个 decode step 产生不必要的 Python 调度开销。PR body 明确强调此问题，并希望通过声明 UNIFORM_BATCH 支持来启用 FULL 模式捕获。

实现拆解

1. 在 vllm/v1/attention/backends/mla/triton_mla.py 中新增 TritonMLAMetadataBuilder 类，继承自 MLACommonMetadataBuilder[MLACommonMetadata]，将类变量 _cudagraph_support 覆写为 AttentionCGSupport.UNIFORM_BATCH。
2. 在 TritonMLABackend 中新增静态方法 get_builder_cls()，返回 TritonMLAMetadataBuilder，使得 pipeline 能够获取到正确的 MetadataBuilder。
3. 新增导入 MLACommonMetadataBuilder 和 AttentionCGSupport，为上述变更提供类型支持。
4. 该变更仅涉及一个文件，无需配置或部署配套改动。FULL CUDAGraph 捕获使用 worst-case max_seq_len，内核中 inline 的 torch.empty 和数据依赖的 num_kv_splits 均为 replay-safe。

关键文件：

• vllm/v1/attention/backends/mla/triton_mla.py（模块 注意力层；类别 source；类型 core-logic；符号 TritonMLAMetadataBuilder, get_builder_cls）: 核心变更文件，新增 TritonMLAMetadataBuilder 类，并更新 TritonMLABackend 以返回新 builder。全部 10 行新增均在此文件中。

关键符号：TritonMLAMetadataBuilder.init, TritonMLABackend.get_builder_cls

关键源码片段

vllm/v1/attention/backends/mla/triton_mla.py

核心变更文件，新增 TritonMLAMetadataBuilder 类，并更新 TritonMLABackend 以返回新 builder。全部 10 行新增均在此文件中。

# 路径 : vllm/v1/attention/backends/mla/triton_mla.py
# 该 PR 的核心变更：新增 MetadataBuilder 并声明 FULL CUDAGraph 支持
​
from vllm.v1.attention.backend import AttentionCGSupport
from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadataBuilder
​
# ... ( 原有导入和类定义 ) ...
​
class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
    # 声明 CUDA Graph 支持 UNIFORM_BATCH 模式（即 FULL capture）
    # 覆盖基类默认的 NEVER，从而让 decode 阶段能够被 FULL 图捕获
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
​
​
class TritonMLABackend(MLACommonBackend):
    # ... 原有实现 ...
​
    @staticmethod
    def get_builder_cls() -> type["TritonMLAMetadataBuilder"]:
        # 返回自定义的 MetadataBuilder，使上层 pipeline 能获取正确的 builder
        return TritonMLAMetadataBuilder

评论区精华

审核过程非常简洁，ZJY0516、MatthewBonanni、mgoin 均给予 APPROVED，gemini-code-assist[bot] 仅给出了自动回复。无实质性讨论或争议。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低。该 PR 仅自定义了 CUDAGraph 模式，未改动 decode kernel 本身。FULL CUDAGraph 捕获使用 worst-case max_seq_len，内核内联的 torch.empty 和 num_kv_splits 依赖数据但符合 replay-safe 条件。回归风险主要在于 CUDAGraph 图捕获与运行时的兼容性，但类似模式已在 FlashInfer 和 FlashAttn MLA 后端中使用，验证充分。
• 影响：直接影响使用 TRITON_MLA 后端的模型（如 Kimi-K2.6），decode 阶段吞吐提升约 14%，TPOT 中位数降低 10%。无 API 或行为变化，完全向后兼容。对系统资源无额外开销。
• 风险标记：暂无

关联脉络

• 暂无明显关联 PR