vLLM Fixes NaN Output Bug in Mamba2 Kernel on SM90

vLLM Fixes NaN Output Bug in Mamba2 Kernel on SM90

The LLM inference framework vLLM has released version v0.32.0rc3, and one key bugfix has drawn attention: it resolves an issue where the Mamba2 chunk scan kernel produces NaN (Not a Number) output on NVIDIA SM90 architecture (i.e., the Hopper series GPUs).

Background: Mamba2 and SM90 Compatibility Challenges

Mamba2 is an efficient sequence modeling architecture based on state space models (SSM). Thanks to its linear complexity advantage in long-sequence processing, it has been widely integrated into modern inference frameworks. As a high-throughput LLM inference engine, vLLM's support for Mamba2 relies on highly optimized CUDA kernels, among which the chunk scan kernel handles block-wise parallel scan computation and sits on the performance-critical path.

The SM90 architecture (Hopper) introduced entirely new features such as the Tensor Memory Accelerator (TMA) and thread block clusters, opening up new possibilities for kernel optimization—but also bringing subtle issues in numerical stability. When the chunk scan kernel runs on SM90, under specific input conditions it outputs NaN, causing inference results to fail completely.

State Space Models (SSM) and Mamba2

State space models originate from classical control theory, describing the evolution of a system's state over time with a set of first-order differential equations, in the form h'(t) = A·h(t) + B·x(t), with output y(t) = C·h(t) + D·x(t). In sequence modeling, this continuous form is discretized into the recurrence relation h_t = Ā·h_{t-1} + B̄·x_t, enabling the model to process a sequence of length n with O(n) complexity, without needing to compute the O(n²) attention matrix as Transformers do.

Mamba2 is an improved architecture built on the original Mamba. Its core change is to bind the SSM's state dimension to the input channels and rewrite the originally per-timestep serial recurrence into matrix multiplication form. This rewrite allows Mamba2 to leverage GPU Tensor Cores for large-scale parallel computation just like attention, while still maintaining constant-level memory footprint and per-step computation during inference. For this reason, Mamba2 is seen as a strong alternative or complement to Transformers in long-context scenarios, and is also a target that high-throughput inference frameworks like vLLM need to specifically adapt to.

Chunk Scan Mechanism

Mamba2's recurrence is essentially a prefix-sum-style linear scan: the state at each timestep depends on the previous step's state, giving it an inherently serial dependency. To parallelize on GPUs, chunk scan cuts a long sequence into several chunks, first independently completing the intra-chunk scan within each chunk to obtain chunk-level initial states and local results; then performing a small-scale inter-chunk scan over chunk states to serially propagate each chunk's starting state; finally using the starting states obtained from the inter-chunk pass to retrospectively correct the output within each chunk.

This two-phase strategy of "intra-chunk parallelism + inter-chunk serialization" reduces the degree of parallelism from sequence-length level down to chunk-count level, and is the key kernel enabling Mamba2 to run efficiently on GPUs. Precisely because it involves two-phase data dependencies and cross-chunk transfer of intermediate states, chunk scan is extremely sensitive to accumulator initialization, chunk boundary handling, and synchronization logic—any error in one boundary condition will be amplified across the entire sequence into NaN or numerical drift.

Conditions That Produce NaN

NaN (Not a Number) is a special value in the IEEE 754 floating-point standard representing an invalid operation result. Common triggers include 0/0, ∞−∞, ∞×0, and taking the square root of a negative number. In GPU kernels, another common source of NaN is reading uninitialized or out-of-bounds memory: such locations may hold arbitrary bit patterns that, when interpreted as floats, happen to be anomalous, and once they enter subsequent accumulation they contaminate the entire computation chain.

The key propagation property is that NaN is extremely "contagious": the result of any arithmetic operation involving NaN is NaN. In a cross-chunk reduction like chunk scan, a NaN produced by one chunk propagates through inter-chunk states to all subsequent chunks, ultimately invalidating the entire output rather than merely causing a local anomaly. This is why such problems typically manifest as "completely wrong inference results" rather than a slight loss of precision.

The Fix: Analyzing PR #61718

According to the commit record, the fix was submitted by developer Vadim Gimpelson, with collaborative support from Claude. The title of PR #61718 is explicitly "Fix NaN output of the Mamba2 chunk scan kernel on sm90", directly targeting the root cause.

Although the commit message does not elaborate on technical details, judging from the nature of the problem, NaN usually stems from invalid operations in floating-point arithmetic (such as 0/0, inf-inf) or uninitialized memory reads. During the parallel reduction of chunk scan, SM90's specific scheduling or memory access patterns may have triggered a boundary condition error. The fix likely involves adjustments to accumulator initialization, synchronization logic, or numerical stability.

SM90 (Hopper) Architecture Features

SM90 is the streaming multiprocessor codename for NVIDIA's Hopper architecture (GPUs such as H100, H200). Compared with the previous generation Ampere (SM80), Hopper introduced several new mechanisms that affect kernel behavior. The Tensor Memory Accelerator (TMA) provides hardware-level asynchronous multidimensional tensor movement, which can greatly reduce address computation overhead and register usage, but requires strictly aligned data layouts. Thread Block Clusters allow multiple thread blocks within the same GPC to execute cooperatively and access each other's shared memory (distributed shared memory), which changes traditional assumptions about synchronization and memory visibility. In addition, Hopper's warp scheduler and instruction issue timing also differ from previous generations.

These features mean that kernels originally running correctly on SM80 may, on SM90, expose hidden races or boundary conditions due to changes in scheduling order, register allocation, or memory access timing. Two-phase kernels like chunk scan that depend on precise synchronization are especially susceptible, and this NaN issue very likely belongs to this category of numerical stability defects triggered by architectural differences.

Impact on the Inference Ecosystem

For users deploying Mamba2-class models with vLLM, especially in scenarios running on Hopper GPUs such as H100/H200, this fix is crucial. NaN output not only causes incorrect generation results but can also crash downstream tasks. This rapid fix reflects the vLLM community's emphasis on hardware compatibility and numerical robustness.

One detail worth noting: this version is still a release candidate (rc3), so production environments are advised to upgrade only after thorough testing.

Summary and Recommendations

vLLM v0.32.0rc3 resolves the NaN issue of the Mamba2 chunk scan kernel on SM90 through #61718, improving the framework's reliability on new hardware. Developers should watch for the subsequent official release of this version and verify the inference stability of Mamba2 models on Hopper architecture.

One-line Summary

vLLM v0.32.0rc3 fixes the issue where the Mamba2 chunk scan kernel outputs NaN on SM90.

Paragraph Summary

vLLM released version v0.32.0rc3, fixing a bug where the Mamba2 chunk scan kernel produced NaN output on NVIDIA SM90 (Hopper) architecture. Mamba2 is based on state space models and can process long sequences with linear complexity; its efficient implementation in vLLM relies on a chunk-wise parallel scan CUDA kernel. New features introduced by Hopper architecture, such as TMA and thread block clusters, changed scheduling and memory access timing, exposing the kernel's numerical stability issue under specific inputs and causing inference results to fail completely. The fix was completed in PR #61718, submitted by developer Vadim Gimpelson. Since the version is still rc3, the article recommends upgrading in production only after thorough testing and watching for the subsequent official release.