Accelerating Mamba-2 with High-Performance Forward Pass Kernels in NVIDIA CUDA Tile

NVIDIA CUDA Tile (cuTile) is a Python-based DSL and parallel programming model for NVIDIA GPUs that automatically leverages advanced hardware capabilities to deliver peak performance for tile-level GPU programming. By abstracting thread- and warp-level scheduling behind high-level, tile-centric primitives, cuTile enables you to write GPU kernels in Python while compiler optimizations automatically target advanced hardware features like Tensor Cores.

As an early adopter of cuTile, machine learning engineers at Jane Street implemented Mamba-2, a state-of-the-art sequence model, in cuTile and partnered with NVIDIA to optimize the kernels. Jane Street is a deeply technical global trading form, using deep learning to drive innovation in advanced quantitative research. This post walks you through the three custom kernels powering the Mamba-2 forward pass and the cuTile optimizations behind them.

Overview and prerequisites

To use cuTile, ensure your environment meets the requirements outlined in the cuTile Python Quickstart guide. Import the package as cuda.tile and expose it as ct throughout the kernel implementations featured in this post.

Sequence modeling computational trade-off

Sequence modeling architectures face a computational trade-off: traditional softmax attention trains rapidly using matrix multiplication accelerated by Tensor Cores, but at inference, compute and memory requirements increase linearly with sequence length. In contrast, recurrent neural networks (RNNs) suffer from slow, nonparallelizable training but require constant compute and fixed memory footprint regardless of sequence length.

Mamba-2 architecture

The Mamba-2 architecture comes from a family of models aiming for the best of both worlds. Similar to an RNN, it has a fixed size hidden state and a recurrent update that can be computed chunkwise in a computationally efficient manner using Tensor Cores.

The core recurrence of Mamba-2 is:

\(h_t = \exp(A \cdot dt_t) h_{t-1} + dt_t B_t x_t\)

yt=ct⊤ht+Dxty_t = c_t^\top h_t + D x_t

Where:

\(x_t\), \(dt_t\), \(B_t\), and \(c_t\) are the inputs to the Mamba-2 layer (obtained by projecting the inputs to the block, like queries, keys and values in an attention layer) at token \(t\) \(y_t\) is the output at token \(t\) \(h_t\) is the fixed-size hidden state carried across tokens. It’s not a vector as in RNNs, but a matrix (more numbers and more representational capacity) \(A\) and \(D\) are learned model parameters

Intuitively:

\(dt_t\) (one scalar per head) represents how much of the state should be updated. All values in \(dt\) are negative Because the previous state is multiplied \(\exp(Adt_t)\), the state stays the same when \(dt_t = 0\)  \(x_t\) (a vector of size H) plays a role similar to the values in attention \(B_t\) (a vector of size S) plays a role similar to the keys in attention \(c_t\) (a vector of size S) is how the state is mixed to produce the outputs and plays a role similar to the queries in attention

The state \(h_t\) is a matrix of size S,H (\(B_t x_t^\top\) is an outer product). To make computation easier to write, use \(a_t = Adt_t\) and call \(b_t\) the product \(dt_t B_t\). The equation is then simply:

\(h_t = \exp(a_t)h_{t-1} + b_t x_t^{\mathsf{T}}, y_t = c_t^{\mathsf{T}}h_{t-1} + Dx_t\)

This recurrent formulation is ideal for sequential inference. However, training and prefill require something faster than a for-loop. While the linear state update allows for a parallel scan algorithm, this approach requires materializing an enormous hidden state tensor making it primarily memory-bound and unable to utilize Tensor Cores.

An alternative approach leverages Tensor Cores by computing any future hidden state directly through matrix multiplication. Ignoring the \(\exp(a_t)\) term for simplification, going from \(h_t\) to \(h_{t+T}\) is just:

\(h_{t+T} = h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + b_{t+2}x_{t+2}^{\mathsf{T}} + \cdots + b_{t+T}x_{t+T}^{\mathsf{T}}\)

If building a matrix \(B\) of shape (T, S) by stacking all of the \(b_{t+i}\) and a matrix \(X\) of shape (T, H) by stacking all of the \(x_{t+i}\) then this is just \(h_{t+T} = h_t + BX^T\):

With this jump formula, you can efficiently compute the hidden state at various chunk “checkpoints” in the sequence. The next step is to find a way to compute outputs \(y_1\), \(y_2\), \(y_3\), \(y_4\) given only \(h_0\). This would be easy if computing \(h_1\), \(h_2\), \(h_3\), \(h_4\), … but doing so would defeat the point of the jump formula.

Still removing the decay term \(\exp(a_t)\) for now, the goal is to compute:

\(y_{t+i} = c_{t+i}^{\mathsf{T}}h_{t+i} + Dx_{t+i} = c_{t+i}^{\mathsf{T}}\left(h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + \cdots + b_{t+i}x_{t+i}^{\mathsf{T}} \right) + Dx_{t+i}\)

Where the previous formula is used to expand the hidden state. The skip connection term \(Dx_{t+i}\) is easy to reason about. For the others, it looks like this:

Again, grouping a chunk of \(x_t\), \(b_t\), \(c_t\) and \(y_t\) in matrices enables re-expressing this with Tensor Core-friendly matrix products (Figure 3) where the triangle in the \(\mathit{C}\) by \(\mathit{B}\) product is a causal mask (like in causal softmax attention).

When computing the output \(y_{t+i}\) you have some product of \(c_{t+i}\) by all of the \(b_j\) for \(j \leq t + i\) so when computing the full matrix product, you need to mask out the ones in the future (above the diagonal). In math this becomes (adding back the skip connection):

\(Y = Ch_t + \left((CB^{\mathsf{T}})\cdot M\right)X + DX\)

With \(M\) the traditional causal mask.

Note that this looks very similar to attention (without the softmax) because when the decay term is removed, Mamba-2 is just linear attention (with some renaming). Adding back the decay makes things a little bit trickier. The expansion of the recurrence will look like:

\(h_(t+T)=exp⁡(a_(t+1)+a_(t+2)+⋯+a_(t+T) ) h_t+exp⁡(a_(t+2)+⋯+a_(t+T) ) b_(t+1) x_(t+1)+exp⁡(a_(t+3)+⋯+a_(t+T) ) b_(t+2) x_(t+2)+⋯+b_(t+T) x_(t+T)\)

Reformulating in log-space allows for replacing sequential multiplications with cumulative sums. This remains a matmul if the rows of \(B\) (or \(X\)) are pre-multiplied by the exponentials. Keeping track of the cumulative sum of \(a_t\) in a tensor \(G\), this results in the following new equation:

\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)

Plugging the new hidden state jump into the output computation results in the following:

\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)

where the matrix \(\exp(G – G^{\mathsf{T}})\) contains at \((i, j)\) the term \(\exp(g_i – g_j)\). Note that all of the terms where \(i < j\) will be masked by the causal mask, so only keep terms where \(g_t\) is the sum of more of the \(a_t\) than \(g_j\). By construction, all \(a_t\) values are negative, which makes \(g_i – g_j\) negative. Any exponentials of positive numbers are never kept, guaranteeing numerical stability.

This will translate into three kernels:

chunk_cumsum: Compute the cumulative sums of the \(a_t\) inside each chunk. fwd_pass_state: Compute the hidden state at the beginning of each chunk using the jump formula above. compute_out: Compute the outputs for each chunk using the initial hidden state for the chunk computed in the previous step.

Kernel 1: Computing cumulative sums

The cumulative sums of the \(a_t = Adt_t\) (the \(G\) in the previous math) are needed in both Steps 2 and 3. To avoid recomputing several times, this first kernel will compute \(a_t\) once and store it. Note that this is not a global cumulative sum, just the cumulative sum within each chunk.

The inputs to this kernel are a per-head scalar decay a and a per-token timestep dt. The effective decay at token t for head h is a[h] * dt[t, h], and the kernel accumulates that product along the chunk dimension. These cumulative sums are all in log-space and the exponential of them is always taken (or some linear combination of them). Since it’s much cheaper to compute exp2 instead of exp, the result is scaled by log2(e).

This kernel introduces the cuTile primitives that all three kernels use: ct.bid, ct.load / ct.store, and ct.cumsum.

The full kernel is short:

The kernel is launched with a grid of (num_chunks, num_head_tiles, batch_size), the first line determines where you are in the grid using ct.bid. The broadcast product of a_tile (shape [BLOCK_H]) and dt_tile (shape [1, chunk_size, BLOCK_H]) gives the effective per-token, per-head decay over the tile. ct.cumsum(..., axis=1) then scans along the intra-chunk token dimension, and the result is stored in the layout expected by fwd_pass_state and compute_out.

Kernel 2: fwd_pass_state

fwd_pass_state computes the initial hidden state of each chunk and the final hidden state after the last chunk. It does not directly compute the output. Instead, it prepares the carried state that compute_out will later project into outputs. At its core, the kernel is a sequential loop over chunks as described earlier. Each kernel launch loads one (batch, head, head_tile, state_tile) slice of the hidden state, as shown below. Within that slice, hid_tile has a shape [BLOCK_HEAD, BLOCK_STATE] and will contain a tile of the initial hidden state for each chunk.

In each iteration of the loop, first store the hidden state:

The rest of the loop loads one chunk of x, b, dt, and decay_cumsum to compute the next chunk’s initial hidden state. As shown previously, the math is:

\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)

ct.extract enables taking the last token from the decay tile loaded (to compute \(G_last\)). Use exp2 instead of exp for performance reasons, as shown previously. ct.mma is the API that will execute the matrix-multiply add while leveraging the Tensor Cores.

After the final chunk, hid_tile contains the final hidden state, which is stored. During inference this value can be fed back as the next call’s init_state; during training it is usually discarded.

The sequential dependence stays inside a single grid point (program instance), while the independent (batch, head, head_tile, state_tile) slices run in parallel across the grid.

Kernel 3: compute_out

compute_out is the final kernel in the forward pass. It takes the per-chunk initial hidden states produced by fwd_pass_state and computes the token outputs, implementing the output equation derived earlier:

\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)

A direct PyTorch implementation would compute these three terms separately, materializing the full [B, n_chunks, chunk_size, chunk_size] interaction matrix C @ B.T in global memory. compute_out fuses them into one tiled kernel and keeps that matrix in registers.

Each kernel launch owns one (batch, chunk, head, head_tile) output tile and works in two phases. Phase 1 loops over state tiles, using ct.mma to build the interaction matrix cb_tile = C @ B.T and accumulate the carried-state term C @ H.T at the same time.

Phase 2 applies the decay and causal mask to cb_tile (ct.arange and ct.where build the lower-triangular mask M, so no mask tensor is ever materialized) and then adds all of the terms to get the final output tile. Both accumulators start as ct.full register tiles, and the final ct.mma casts to ct.tfloat32 to run through the Tensor Cores at near-FP32 accuracy.

As mentioned, while some numerical unstable exponentials are computed in the ct.exp2(d_tile[:, None] - d_tile[None, :]) line (some of those are positive numbers), immediately mask them out so potential overflows don’t end up poisoning the computation.

Mamba-2 performance optimization

With the Mamba-2 workload defined in the cuTile-Python DSL, several key performance optimizations are introduced:

Latency hint: cuTile allows you to annotate ct.load and ct.store operations with latency hints to model DRAM traffic. The compiler leverages these hints to optimize instruction scheduling and pipeline buffer depth. Swap A/B: On the NVIDIA Blackwell fifth-generation Tensor Core, MMA operands are asymmetric: operand A can live in Tensor Memory, but operand B must live in Shared Memory. For an input-fusion pattern like the SIMT-op + Dot-op in fwd_pass_state, the identity \(B^{\mathsf{T}}A^{\mathsf{T}} = (AB)^{\mathsf{T}}\) is used to route the fused operand (b_dt_d) to operand A, relieving shared-memory contention. GMEM layout: The default strided memory access for decay_cumsum and dt lowers to cache-unfriendly LDG instructions. Transposing to a [batch, head, token] layout ensures contiguous access, enabling the use of TMA. cuTile-Python autotuner: Since performance bottlenecks shift with input shape, no single optimization strategy fits every model configuration. These optimizations are integrated alongside cuTile parameters (block_size, occupancy) into a unified search space, leveraging cuTile exhaustive_search to automatically find the optimal setup for each input shape.

Mamba-2 forward pass performance benchmarks

This implementation of the Mamba-2 forward pass kernel is available through TileGym. A complete example is shown below.

The performance of the implemented Mamba-2 forward pass was evaluated against an open source implementation of the Mamba-2 kernel written in Triton, under the following software and hardware configurations:

Software

mamba_ssm 2.3.2.post1 (NVIDIA Dynamo-Triton reference), Triton 3.7.1 cuTile: cuda-tile 1.5.0 TileIR assembler (tileiras) 13.4 PyTorch 2.13.0, CUDA 13.3 (Driver and Runtime API, NVCC V13.3.33), Driver 595.58.03 cuda-python / cuda-bindings 13.3.1 Python 3.12.3, Ubuntu 24.04.4 LTS, kernel 6.8.0-107

Hardware

Single NVIDIA B300 SXM6, compute capability 10.3 (NVIDIA Blackwell Ultra) 1920 MHz max SM clock, 1100 W power limit, 268 GiB HBM

The results show that the optimized cuTile forward-pass kernels are 1.14x to 1.82x faster, with speedups growing as sequence lengths and state size increases. Note that these are kernel-level speedups only and end-to-end training results will also depend on the backward pass implementations.

Get started with cuTile

cuTile enables you to write clear, readable, and highly performant kernels in Python. Stay tuned for a follow-up post describing how the backward pass is implemented in cuTile.

Get started with cuTile and explore the documentation, GitHub repos, and sample code.