PR #43083仅为B200、GB200和H100_80GB_HBM3提供了调优配置。在H200和RTX PRO 6000 Blackwell上，加载器回退到Triton内置启发式，导致显著性能损失。本PR为这两种设备提供针对性的调优参数，以释放硬件潜力。
建议合并。变更安全且性能收益明确，可立即改善相关GPU用户的使用体验。后续可关注Triton版本升级后的配置有效性。
审核人tomeras91快速批准（LGTM），无任何争议或问题提出。作者在批准后请求合并。
使用benchmarks.kernels.benchmark_selective_state_update脚本，针对headdim=64, dstate=128、cache_dtype为float16和float32、nheads范围8-256进行调优，生成每个批量大小对应的最佳BLOCK_SIZE_M和num_warps。生成的JSON文件以headdim=64,dstate=128,device_name={设备名},cache_dtype={数据类型}.json命名，存放于vllm/model_executor/layers/mamba/ops/configs/selective_state_update/目录。脚本同时执行验证（atol=0.01），确保所有配置数值正确性。共有4个新文件，分别对应H200 float16/float32和RTX PRO 6000 Blackwell float16/float32。
| 文件 | 模块 | 状态 | 重要度 |
|---|---|---|---|
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float16.json |
内核配置 | added | 5.78 |
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float32.json |
内核配置 | added | 5.38 |
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float16.json |
内核配置 | added | 5.38 |
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_RTX_PRO_6000_Blackwell_Server_Edition,cache_dtype=float32.json |
内核配置 | added | 4.98 |
vllm/model_executor/layers/mamba/ops/configs/selective_state_update/headdim=64,dstate=128,device_name=NVIDIA_H200,cache_dtype=float16.json
infrastructure
H200 float16 配置，加速效果最为显著（最高 2.9x），覆盖主要使用场景
{
// Triton 版本标识，用于向后兼容
"triton_version": "3.6.0",
// 键为有效批量大小 (effective batch)，值为 {BLOCK_SIZE_M, num_warps}
"8": { "BLOCK_SIZE_M": 4, "num_warps": 2 },
"16": { "BLOCK_SIZE_M": 4, "num_warps": 1 },
"32": { "BLOCK_SIZE_M": 4, "num_warps": 1 },
"64": { "BLOCK_SIZE_M": 16, "num_warps": 4 },
"128": { "BLOCK_SIZE_M": 8, "num_warps": 2 },
"256": { "BLOCK_SIZE_M": 8, "num_warps": 2 },
"512": { "BLOCK_SIZE_M": 16, "num_warps": 1 },
"1024": { "BLOCK_SIZE_M": 16, "num_warps": 1 },
"2048": { "BLOCK_SIZE_M": 8, "num_warps": 2 },
"4096": { "BLOCK_SIZE_M": 16, "num_warps": 2 },
"8192": { "BLOCK_SIZE_M": 32, "num_warps": 2 },
// 更大的批量通常需要更大 BLOCK_SIZE 和更多 warp
"196608": { "BLOCK_SIZE_M": 16, "num_warps": 2 },
"262144": { "BLOCK_SIZE_M": 16, "num_warps": 2 }
}
当前评论区没有形成足够清晰的争议点或结论，后续有更多讨论时会体现在这里。
风险极低。所有配置均经过验证（atol=0.01），且仅影响新增设备，现有设备继续使用已有配置或内置启发式。若加载失败，代码回退默认参数。唯一需关注：配置针对Triton 3.6.0生成，若将来Triton版本更新，可能需要重新验证；但现有机制允许按版本承载配置，因此风险可控。
对H200和RTX PRO 6000 Blackwell用户有明确性能提升，尤其是使用Mamba-2结构模型（如Nemotron-Nano-9B）时：H200端到端吞吐量提升约2.6%，TTFT和ITL均下降约2.2-2.6%。对其他GPU（H100、B200等）无影响。变更无API破坏，无需用户侧改动。
当前没有检测到明确关联的 Issue 链接，后续同步到相关引用后会出现在这里。
参与讨论