PR 47947 分析报告：AMD MI300X 的 selective_state_update float32 调优配置

执行摘要

此 PR 为 AMD Instinct MI300X GPU 添加了 float32 状态缓存的 selective_state_update kernel 调优配置文件。与已有的 float16 配置 (PR #47945) 互补，使得 MI300X 在两种常见缓存精度下都能使用优化的 Triton launch 参数，获得显著的 kernel 加速。PR 仅添加一个 JSON 文件，风险极低，数据充分，验证完整。

功能与动机

在 Mamba-2 模型中，SSM 状态缓存默认以 float32 存储。此前 MI300X 只有 float16 状态缓存的调优配置 (PR #47945)，当用户使用默认 fp32 状态缓存时，vLLM 会回退到通用启发式 launch 参数 (_get_default_ssm_launch_config())，导致 kernel 执行效率低下。此 PR 增加了对应的 float32 配置，使 MI300X 与 NVIDIA 设备及 MI355 在两种 dtype 下都能享受调优加速。PR 描述指出 "Each NVIDIA device in configs/selective_state_update/ ships both cache_dtype=float16 and cache_dtype=float32，而 MI355 近期也已获得两种配置，此 PR 使 MI300X 达到同等覆盖。"

实现拆解

1. 新增配置 JSON 文件：在路径 vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float32.json 下创建配置文件，包含 13 个 effective_batch 点（128 到 262144）对应的 BLOCK_SIZE_M 和 num_warps 最优组合。
2. 配置生成：使用项目内基准脚本 benchmarks.kernels.benchmark_selective_state_update 在 MI300X 上以 float16 激活、float32 状态缓存运行，搜索每个 effective_batch 下的最优 Triton launch 参数。此脚本与 #47945 和 #47943 使用相同的流程。
3. 验证：
   • --validate：所有 12 个有效点通过 CPU 参考验证，确保正确性。
   • --compare：与通用启发式参数（M=4, w=4）对比，在所有 effective_batch 上均获得加速，最高 1.85× (effective_batch=1024)。
4. 运行时加载：vLLM 在初始化 SSM kernel 时，会根据设备、headim、dstate、cache_dtype 等键在配置目录中查找匹配的 JSON 文件。若找到，则使用其中的参数；否则回退到通用启发式。配置文件仅影响 launch 几何参数，不影响 kernel 计算，因此对模型输出没有影响。

vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI300X,cache_dtype=float32.json

新增的调优配置文件，为 MI300X fp32 状态缓存提供每个 effective_batch 点下的最优 BLOCK_SIZE_M 和 num_warps，是 PR 唯一的变更文件。

{
  // Triton 版本记录，便于未来判断是否需要重新调优
  "triton_version": "3.4.0",
  // 每个 effective_batch 对应的最优 launch 参数
  "128": { "BLOCK_SIZE_M": 8, "num_warps": 4 },
  "256": { "BLOCK_SIZE_M": 8, "num_warps": 4 },
  "1024": { "BLOCK_SIZE_M": 64, "num_warps": 1 },
  "2048": { "BLOCK_SIZE_M": 32, "num_warps": 4 },
  "4096": { "BLOCK_SIZE_M": 32, "num_warps": 4 },
  "8192": { "BLOCK_SIZE_M": 8, "num_warps": 1 },
  "16384": { "BLOCK_SIZE_M": 32, "num_warps": 4 },
  "32768": { "BLOCK_SIZE_M": 8, "num_warps": 1 },
  "65536": { "BLOCK_SIZE_M": 64, "num_warps": 4 },
  "131072": { "BLOCK_SIZE_M": 64, "num_warps": 4 },
  "196608": { "BLOCK_SIZE_M": 64, "num_warps": 1 },
  "262144": { "BLOCK_SIZE_M": 64, "num_warps": 4 }
}

评论区精华

无实质 review 讨论。PR 由 tomeras91 直接 approve，机器人生成了一条自动评论说明 fork PR 的自动化审核被禁用，但核心流程未受影响。PR 作者在 issue 评论中给出了详细的性能数据表格，表示在所有 effective_batch 上均获得加速，最高 1.85× (effective_batch=1024)。

风险与影响

• 风险：极低。仅添加一个 JSON 配置，不修改任何 Python 或 Triton 代码，不会影响模型正确性。配置加载失败时会静默回退到通用启发式，不会导致崩溃。唯一的小风险是配置针对 Triton 3.4.0 生成，未来 Triton 升级后可能需要重新调优。
• 影响：对 MI300X 上使用 fp32 状态缓存的 Mamba-2 模型用户有明显性能提升（kernel 级别加速最高 1.85×），整体端到端延迟降低。由于 Mamba 模型通常包含多个 SSM 层，累加收益可观。对其他 AMD GPU（MI325X、MI308X 等）无影响，但提供了可复用的调优流程。

关联脉络

• #47945：同一设备的 float16 状态缓存配置，此 PR 是其完全对应补充。
• #47943：MI355 的 float32 配置，使用了完全相同的脚本和流程，验证了调优方法的可迁移性。
• #47767：MI355 的 float16 配置，此 PR 系列共同构筑了 AMD GPU 在 selective_state_update kernel 上的全面调优覆盖。
• 跨 PR 趋势：vLLM 项目正在系统性地为所有主流 GPU（NVIDIA、AMD）补全两种状态缓存精度的调优配置，展现了对 Mamba-2 模型推理性能的持续投入。