执行摘要

• 一句话：ROCm sparse MLA 启用 UNIFORM_BATCH CG 模式
• 推荐动作：值得精读，特别是理解 CUDA Graph 模式选择对多 token 预测性能的影响，以及如何通过简单配置对齐不同后端的行为。对于关注 ROCm 性能优化的开发者有借鉴意义。

功能与动机

根据 PR 描述："In order to improve performance especially in CPU bound cases when MTP is enabled UNIFORM_BATCH CG mode needs to be enabled in rocm_aiter_mla_sparse.py. It has not been enabled mostly due to testing reasons and lack of focus for MTP."

实现拆解

1. 修改 CUDA Graph 支持模式：在 ROCMAiterMLASparseMetadataBuilder 类中，将 _cudagraph_support 类属性从 AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE 改为 AttentionCGSupport.UNIFORM_BATCH。这使得 MTP 场景下可以捕获统一批次的 CUDA Graph，减少 CPU 调度开销。
2. 添加 batch 重排阈值初始化：在 __init__ 中新增 self.vllm_config = vllm_config 并调用 self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)。此举确保 ROCm sparse MLA 后端的调度器在推测解码时执行与 CUDA 后端一致的 batch 重排，避免正确性问题。

关键文件：

• vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py（模块 注意力后端；类别 source；类型 core-logic；符号 _cudagraph_support, init）: 唯一修改文件，通过调整 CUDA Graph 支持模式和增加 batch 重排初始化来提升 MTP 性能

关键符号：init

关键源码片段

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

唯一修改文件，通过调整 CUDA Graph 支持模式和增加 batch 重排初始化来提升 MTP 性能

class ROCMAiterMLASparseMetadataBuilder(AttentionMetadataBuilder[ROCMAiterMLASparseMetadata]):
    # 启用 UNIFORM_BATCH CUDA Graph 模式，支持多 token 批处理，提升 CPU 边界性能
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
​
    def __init__(
        self,
        kv_cache_spec: AttentionSpec,
        layer_names: list[str],
        vllm_config: VllmConfig,
        device: torch.device,
    ):
        self.kv_cache_spec = kv_cache_spec
        self.model_config = vllm_config.model_config
        self.model_dtype = vllm_config.model_config.dtype
        parallel_config = vllm_config.parallel_config
        self.device = device
        max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
​
        # 保存 vllm_config 并初始化 batch 重排阈值，对齐 CUDA 后端行为
        self.vllm_config = vllm_config
        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
​
        self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
        self.mla_dims = get_mla_dims(self.model_config)
        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
        # ... 后续初始化 ...

评论区精华

1. 性能基线问题：reviewer tjtanaa 询问 "singlet" 性能数据是否在无此 PR 时收集。作者确认是的，并说明 uniform batch 提升主要在 CPU bound 场景，新增的 reorder 行是为了匹配 CUDA 行为。
2. CI 失败讨论：多个 pre-commit 检查失败，涉及 pip-compile-xpu，但作者和合并者都确认与本 PR 无关，并建议避免频繁 rebase。
3. 依赖 PR #47766：作者表示此 PR 在 #47766 落地方可正常工作，之前因缺少该依赖导致失败。

• 性能基线收集方式 (testing): singlet 数据基于无 PR 版本，reorder 初始化是为了对齐 CUDA，不影响性能基线。
• CI 失败与 pre-commit 问题 (other): CI 失败是基础设施问题，不影响 PR 自身正确性。

风险与影响

• 风险：风险极低。变更仅调整了 CUDA Graph 捕获模式和增加了与 CUDA 后端的对齐初始化，且已在 MI350X/MI355X 上通过吞吐和精度测试（包括 lm_eval）。唯一潜在风险是统一批次模式可能对非 MTP 场景的图捕获行为有影响，但测试表明无精度退化。另外，新增的 batch 重排阈值若与其他调度逻辑不兼容可能引发问题，但由于与 CUDA 后端对齐，风险可控。

• 影响：影响范围：仅限 ROCm 平台、使用 sparse MLA attention 后端、且启用 MTP 的 DeepSeek V3.2 模型。对 CUDA 或其他后端无影响。
  影响程度：中等。对相关场景带来约 3% 的吞吐提升（CPU bound 场景更明显），且无功能退化。团队方面，AMD 的 ROCm 推理用户受益。

• 风险标记：少而精的变更, 依赖其他 PR, 仅 ROCm/DeepSeek 场景

关联脉络

• PR #47766 [ROCm][MRV2] Support MRV2 for ROCm sparse MLA: 本 PR 依赖 #47766 的更改才能正常工作，作者在注释中明确提到 'This PR has been confirmed to be working again now that #47766 has landed.'