vLLM修复Mamba2核函数在SM90上的NaN输出Bug

大型语言模型推理框架 vLLM 发布了 v0.32.0rc3 版本，其中一项关键 Bugfix 引起了关注：修复了在 NVIDIA SM90 架构（即 Hopper 系列 GPU）上，Mamba2 chunk scan kernel 产生 NaN（Not a Number）输出的问题。

问题背景：Mamba2 与 SM90 的兼容性挑战

Mamba2 是一种基于状态空间模型（SSM）的高效序列建模架构，因其在长序列处理上的线性复杂度优势，被广泛集成到现代推理框架中。vLLM 作为高吞吐量 LLM 推理引擎，对 Mamba2 的支持依赖于高度优化的 CUDA kernel，其中 chunk scan kernel 负责分块并行扫描计算，是性能关键路径。

SM90 架构（Hopper）引入了全新的 Tensor Memory Accelerator（TMA）和线程块集群等特性，为 kernel 优化提供了新可能，但也带来了数值稳定性方面的微妙问题。当 chunk scan kernel 在 SM90 上运行时，特定输入条件下会输出 NaN，导致推理结果完全失效。

状态空间模型（SSM）与 Mamba2

状态空间模型源于经典控制理论，用一组一阶微分方程描述系统状态随时间的演化，形式为 h'(t) = A·h(t) + B·x(t)，输出 y(t) = C·h(t) + D·x(t)。在序列建模中，这一连续形式被离散化为递推关系 h_t = Ā·h_{t-1} + B̄·x_t，使模型能以 O(n) 的复杂度处理长度为 n 的序列，无需像 Transformer 那样计算 O(n²) 的注意力矩阵。

Mamba2 是在初代 Mamba 基础上改进的架构，其核心变化是把 SSM 的状态维度与输入通道绑定，并将原本逐时间步的串行递推重写为矩阵乘法形式。这一改写使 Mamba2 可以像注意力一样利用 GPU 的 Tensor Core 进行大规模并行计算，同时在推理时仍保持常数级的显存占用和每步计算量。正因如此，Mamba2 被视为长上下文场景中 Transformer 的有力替代或补充方案，也是 vLLM 这类高吞吐推理框架需要专门适配的对象。

Chunk Scan（分块扫描）机制

Mamba2 的递推本质上是前缀和（prefix sum）式的线性扫描：每个时间步的状态依赖前一步的状态，天然具有串行依赖。为了在 GPU 上并行化，chunk scan 把长序列切成若干块（chunk），先在每块内部独立完成块内扫描，得到块级的初始状态和局部结果；随后对块间状态做一次小规模的跨块扫描，将各块的起始状态串行传播；最后用跨块得到的起始状态回溯修正每块内部的输出。

这种"块内并行 + 块间串行"的两阶段策略把并行度从序列长度级降低到块数量级，是 Mamba2 在 GPU 上高效运行的关键 kernel。也正因为涉及两阶段的数据依赖与中间状态的跨块传递，chunk scan 对累加器初始化、块边界处理和同步逻辑非常敏感，任何一处边界条件出错都会在整个序列上放大为 NaN 或数值漂移。

NaN 的产生条件

NaN（Not a Number）在 IEEE 754 浮点标准中是表示无效运算结果的特殊值，常见触发条件包括 0/0、∞−∞、∞×0，以及对负数开平方等。在 GPU kernel 中，NaN 还有一类常见来源是读取未初始化或越界的内存：这些位置可能残留任意位模式，被解释为浮点数后恰好奇异，参与后续累加便会污染整条计算链。

关键的传播特性是：NaN 具有极强的"传染性"，任何与 NaN 的算术运算结果都是 NaN。在 chunk scan 这种跨块归约中，一个块产生的 NaN 会通过块间状态传播到后续所有块，最终导致整段输出全部失效，而不会只是局部异常。这也是为什么这类问题往往表现为"推理结果完全错误"而非轻微精度下降。

修复内容：PR #61718 解析

根据提交记录，该修复由开发者 Vadim Gimpelson 提交，并得到 Claude 的协作支持。PR #61718 标题明确为“Fix NaN output of the Mamba2 chunk scan kernel on sm90”，直接针对问题根源。

虽然提交信息未展开技术细节，但从问题性质推断，NaN 通常源于浮点运算中的无效操作（如 0/0、inf-inf）或未初始化内存读取。在 chunk scan 的并行归约过程中，SM90 的特定调度或内存访问模式可能触发了边界条件错误。修复可能涉及对累加器初始化、同步逻辑或数值稳定性的调整。

SM90（Hopper）架构特性

SM90 是 NVIDIA Hopper 架构（H100、H200 等 GPU）的流多处理器代号。相比前代 Ampere（SM80），Hopper 引入了多项影响 kernel 行为的新机制。Tensor Memory Accelerator（TMA）提供硬件级的多维张量异步搬运，能大幅减少地址计算开销和寄存器占用，但要求数据布局严格对齐。线程块集群（Thread Block Cluster）允许同一 GPC 内的多个线程块协同执行并访问彼此的共享内存（分布式共享内存），这改变了传统的同步与内存可见性假设。此外 Hopper 的 warp 调度器和指令发射时序也与前代存在差异。

这些特性使得原本在 SM80 上正确运行的 kernel，在 SM90 上可能因调度顺序、寄存器分配或内存访问时序的改变而暴露隐藏的竞态或边界条件。chunk scan 这类依赖精确同步的两阶段 kernel 尤其容易受影响，本次 NaN 问题很可能就属于此类由架构差异触发的数值稳定性缺陷。

对推理生态的影响

对于使用 vLLM 部署 Mamba2 类模型的用户，尤其是运行在 H100/H200 等 Hopper GPU 上的场景，该修复至关重要。NaN 输出不仅导致生成结果错误，还可能引发下游任务崩溃。此次快速修复体现了 vLLM 社区对硬件兼容性和数值鲁棒性的重视。

说个细节，该版本仍为 release candidate（rc3），建议生产环境在充分测试后升级。

总结与建议

vLLM v0.32.0rc3 通过 #61718 解决了 Mamba2 chunk scan kernel 在 SM90 上的 NaN 问题，提升了框架在新硬件上的可靠性。开发者应关注该版本后续正式发布，并在 Hopper 架构上验证 Mamba2 模型的推理稳定性。