NVIDIA Accelerates Mamba-2 Forward Kernel with CUDA Tile, Speeding Up State Space Model Inference

NVIDIA's developer blog recently shared a hands-on look at writing a high-performance forward pass kernel for Mamba-2 using CUDA Tile (cuTile). Mamba-2 belongs to the state space model (SSM) family, processing long sequences with linear complexity, and is emerging as an important alternative to Transformers in long-context scenarios. But turning its theoretical advantages into real-world throughput inevitably requires engineering optimization at the kernel level.

What CUDA Tile Is and Why It Suits Writing Mamba-2 Kernels

cuTile is a Python-based domain-specific language (DSL) and parallel programming model introduced by NVIDIA. It lets developers describe tile-level computations in a way that's close to mathematical notation, with the compiler automatically mapping them onto GPU thread hierarchies, shared memory, asynchronous copies, and other hardware capabilities. In other words, work that previously required writing large amounts of CUDA C++ boilerplate is abstracted into higher-level block operations.

This is especially critical for Mamba-2. The forward computation of Mamba-2 involves selective scan, state recurrence, and several element-wise gating operations. It has serial dependencies along the sequence dimension and parallel space along the feature dimension. cuTile's tile semantics happen to express this "intra-block parallelism, inter-block recurrence" structure, allowing the scan process to advance in a pipelined fashion tile by tile while keeping state resident in on-chip memory, reducing round trips to global memory.

Domain-Specific Language (DSL) is a programming language or model designed for a specific application domain, as opposed to general-purpose languages (like C++ or Python). DSLs trade scope limitation for expressiveness and performance: the compiler can assume specific computation patterns (such as tiled matrix operations or tensor contractions), thereby automatically handling thread mapping, memory layout optimization, and instruction selection. In GPU programming, cuTile, Triton, and Halide are all attempts of this kind, aiming to give developers near-handwritten-CUDA performance without manually managing thread blocks, shared memory, and synchronization barriers. The cost is limited flexibility—operations beyond the DSL's expressive range often need to fall back to lower-level implementations.

Selective scan is the core operation of the Mamba family of SSMs. Unlike classical linear time-invariant SSMs, Mamba's discretization parameters Δ, input matrix B, and output matrix C all vary dynamically with each token, enabling the model to "select" what historical state to retain or forget based on content. The forward computation is therefore a recurrence along the sequence dimension: h_t = A_bar_t · h_{t-1} + B_bar_t · x_t, where A_bar and B_bar are determined by Δ. Because the recurrence has strict temporal dependencies, it cannot compute all positions in parallel at once like attention does. Instead, tiled scan or parallel prefix sum algorithms must be used to decompose the serial chain into two phases: intra-block parallelism and inter-block propagation—which is also the main challenge in kernel design.

Key Optimization Focuses of the Mamba-2 Forward Kernel

According to NVIDIA's blog, the core idea behind this forward kernel is to split the sequence into blocks, compute local contributions in parallel within each block, and connect blocks via state propagation. This tiled scan strategy follows the same principle as parallel prefix sum, but is specially adapted to handle the input-dependent parameters of SSMs (such as Δ, B, and C varying with tokens).

The engineering benefits come mainly from three areas: first, tiling increases computational density, allowing matrix multiplications and element-wise operations to more fully feed the tensor cores; second, intermediate state is kept on-chip as much as possible, reducing bandwidth pressure; third, cuTile's automatic scheduling enables overlap between asynchronous data movement and computation. For inference scenarios, the efficiency of the forward kernel directly determines first-token latency and long-sequence throughput, so these optimizations have direct practical value.

Parallel prefix sum (scan) is a classic algorithm for parallelizing recurrence operations that satisfy the associative property. Take accumulation as an example: the naive approach requires O(n) serial steps, while algorithms like Blelloch scan can complete in O(log n) steps at the cost of extra computation. The state recurrence of an SSM is essentially the composition of affine transformations, and the composition of affine transformations satisfies associativity, so scan algorithms can be applied: first split the sequence into blocks, compute local prefix results independently within each block, then perform a small-scale scan on the inter-block aggregation results, and finally broadcast the inter-block results back to each block for correction. This "two-pass scan" structure is the theoretical foundation of the tiled scan kernel, and the choice of block size requires balancing parallelism, on-chip memory usage, and recomputation overhead.

From Python DSL to Production-Grade Kernels

cuTile is not positioned to replace CUDA C++, but to lower the barrier to developing high-performance kernels. Researchers can use Python to rapidly iterate on algorithm variants, validate whether a tiling strategy works, and then push mature solutions down to lower-level implementations. NVIDIA's choice of Mamba-2 as an example is, in a sense, a demonstration of how these new architectures can be well supported on new programming models.

For teams doing model deployment, this means two things: first, the operator library ecosystem for SSM-type models is catching up—the fused operator expertise previously honed around attention mechanisms needs to be re-migrated to scan and recurrence structures; second, the Python-first kernel development workflow is expected to shorten the cycle from "paper reproduction to online serving." Of course, cuTile is still evolving, and its performance ceiling, coverage of the latest hardware, and synergy with other compiler stacks still need more empirical data to verify.

Operator fusion refers to merging multiple element-wise or lightweight operations into a single kernel to avoid repeatedly writing intermediate results to global memory. In the Mamba-2 forward pass, the scan is typically followed by gating, normalization, and output projection. If each is launched as a separate kernel, tensors must be read and written each time, and bandwidth overhead accumulates rapidly. After fusion, intermediate values can remain in registers or shared memory and directly participate in the next computation, reducing both memory accesses and kernel launch latency. The reason why fused operator expertise like FlashAttention, honed around attention mechanisms, is hard to directly transplant is that the fusion boundaries of SSMs are determined by the recurrence structure of the scan, requiring re-partitioning of computation stages and arrangement of state lifetimes on-chip.

Summary

Mamba-2's competitiveness depends not only on architecture design but also on whether the underlying kernels can translate the theoretical advantages of linear complexity into measured throughput. NVIDIA's attempt to write the forward kernel with cuTile demonstrates a path to rapidly approach hardware limits using a high-level DSL. For readers interested in long-context inference costs and the engineering deployment of state space models, this direction is worth continued attention.