#48159 [ROCm] Add tuned selective_state_update config for AMD MI350

原始 PR 作者 giuseppegrossi 合并时间 2026-07-15 15:18 文件变更 2 提交数 1 评论 2 代码增减 +102 / -0

执行摘要

vLLM 已为 NVIDIA GPU 及 AMD MI300X/MI355 提供 SSU 调优配置，但 MI350X 使用默认启发式配置，性能不佳。此 PR 填补了这一空白，使 MI350X 用户能获得堪比 NVIDIA 硬件的 Mamba 解码性能。

值得合并，支持新硬件且性能收益明显。应鼓励类似模式为其他未覆盖的 GPU 添加调优配置。

讨论亮点

无实质性 review 讨论，维护者 Tomer Asra 已批准。

步骤：

在 MI350X 上运行内部基准脚本 benchmark_selective_state_update，扫描 effective_batch 128–262144 范围，对每个缓存 dtype 选择最优的 BLOCK_SIZE_M 和 num_warps。
将结果保存为 JSON 文件，命名包含设备名、headdim、dstate 和 cache_dtype，确保运行时自动加载。
使用 --validate 对比 CPU 参考验证每个配置点的数值正确性（12/12 通过）。
使用 --compare 确认速度提升相对于默认启发式。
端到端使用 lm_eval lambada_openai 确认模型输出无变化。

文件	模块	状态	重要度
`vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI350_OAM,cache_dtype=float16.json`	SSU 配置	added	5.19
`vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=AMD_Instinct_MI350_OAM,cache_dtype=float32.json`	SSU 配置	added	5.19

分析完成后，这里会展示 LLM 生成的相对完整源码片段和详细注释。

当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。

风险极低：配置仅调整 Triton 启动几何参数，不影响内核计算逻辑。已验证每个 effective_batch 点与 CPU 参考一致，lm_eval 精度和困惑度在配置有无情况下完全相同。但若未来硬件或驱动变化，配置可能不再最优；但框架设计允许通过删除文件回退到默认启发式。

直接影响：使用 AMD MI350X 且运行 Mamba2 模型（headdim=64, dstate=128）的用户，Mamba 解码阶段自动加载调优配置，获得 1.2×–2.3× 性能提升。影响范围限于该特定设备/模型组合，其他用户无影响。团队维护负担小：配置无需代码维护。

特定硬件依赖配置需维护

当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。