执行摘要

此PR修复了vLLM中Mamba2 SSD内核的数值溢出问题，通过钳制dA_cumsum差异为非正，防止因浮点舍入错误导致的Inf和NaN，提升推理稳定性，对齐上游修复。

功能与动机

当Mamba2模型的参数|A|值较大时，dA_cumsum的计算可能因并行前缀扫描（tl.cumsum）的浮点舍入错误，使差异(dA_cs_last - dA_cs_k)略为正数，导致exp()溢出到Inf，进而传播为NaN影响后续解码。PR body指出："When a Mamba2 model has large |A| values, dA_cumsum reaches magnitudes where float32 ULP exceeds safe range... causing exp() to overflow to Inf." 修复旨在消除此风险，确保数值稳定性。

实现拆解

主要修改两个Triton内核文件：

• vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py：在_chunk_scan_fwd_kernel函数中，将cb *= fast_exp(dA_cs_m[:, None] - dA_cs_k[None, :]) 改为 cb *= fast_exp(tl.minimum(dA_cs_m[:, None] - dA_cs_k[None, :], 0.0))。
• vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py：在_chunk_state_fwd_kernel函数中，将scale = fast_exp(dA_cs_last - dA_cs_k) * dt_k 改为 scale = fast_exp(tl.minimum(dA_cs_last - dA_cs_k, 0.0)) * dt_k。

变更确保所有差异在指数化前被钳制到≤0，防止溢出。

评论区精华

review讨论中，gemini-code-assist[bot]强调：

"This pull request addresses a critical numerical stability issue... The changes are correct, minimal, and align with a similar fix in the upstream Mamba implementation."
tdoublep快速批准。Issue评论中，作者kibitzing补充：
"This fix addresses numerical stability in Mamba2 SSD kernels, with an approach consistent with those used in Mamba and NVIDIA/Megatron-LM."
讨论焦点是修复的正确性和上游对齐。

风险与影响

风险：添加tl.minimum可能微增计算开销，但影响可忽略；修复后NaN消除，但未测试所有模型变体，潜在边缘情况风险低。影响：用户端避免推理错误，系统端增强Mamba2内核可靠性，团队端建立数值处理范例。

关联脉络

从近期历史PR看，类似bugfix如#36540（修复TRTLLM ragged MLA预填充数值问题）和#37010（修复FusedMoE权重加载问题），反映团队持续优化内核稳定性趋势。此PR专门针对Mamba2 SSD内核，与上游Mamba PR#713和Megatron-LM实现关联，强调跨项目一致性。