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

Jane Street machine learning engineers implemented the Mamba-2 sequence model in NVIDIA's CUDA Tile (cuTile) Python DSL and partnered with NVIDIA to optimize the three custom kernels powering the Mamba-2 forward pass. The post details how cuTile abstracts thread- and warp-level scheduling behind tile-centric primitives so the compiler targets Tensor Cores, and explains the Mamba-2 recurrence h_t = exp(A·dt_t)h_{t-1} + dt_t B_t x_t, in which the hidden state is a fixed-size S×H matrix rather than an RNN-style vector. The work matters because Mamba-2 aims to combine attention's Tensor Core-friendly training with a recurrent model's constant compute and fixed memory footprint at inference.

NVIDIA CUDA Tile https://developer.nvidia.com/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 https://www.janestreet.com/ 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 https://docs.nvidia.com/cuda/cutile-python/quickstart.html 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$ 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: python @ct.kernel def chunk cumsum kernel a: torch.Tensor, num heads dt: torch.Tensor, batch size, num tokens, num heads output: torch.Tensor, batch size, num tokens, num heads chunk size: ct.Constant int , BLOCK H: ct.Constant int , - None: chunk idx, head tile idx, batch = ct.bid 0 , ct.bid 1 , ct.bid 2 dt tile = ct.load dt, index= batch, chunk idx, head tile idx , shape= 1, chunk size, BLOCK H , padding mode=ct.PaddingMode.ZERO .astype ct.float32 a tile = ct.load a, index= head tile idx, , shape= BLOCK H, , padding mode=ct.PaddingMode.ZERO .astype ct.float32 d tile = ct.cumsum a tile dt tile, axis=1 LOG2E ct.store output, index= batch, chunk idx, head tile idx , tile=d tile.astype output.dtype 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 ou t 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. head tile idx, state tile idx, batch head = ct.bid 0 , ct.bid 1 , ct.bid 2 batch = batch head // x.shape 2 head = batch head % x.shape 2 group = head // num heads per group hid tile = ct.load init state, index= batch, head, head tile idx, state tile idx , shape= 1, 1, BLOCK HEAD, BLOCK STATE , .reshape BLOCK HEAD, BLOCK STATE In each iteration of the loop, first store the hidden state: num chunks = ct.cdiv b.shape 1 , chunk size for chunk idx in range num chunks : ct.store hiddens, index= batch, head, chunk idx, head tile idx, state tile idx , tile=hid tile.reshape 1, 1, 1, BLOCK HEAD, BLOCK STATE .astype hiddens.dtype , 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. x tile = ct.load x, index= batch, chunk idx, head, head tile idx , shape= 1, chunk size, 1, BLOCK HEAD , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, BLOCK HEAD b tile = ct.load b, index= batch, chunk idx, group, state tile idx , shape= 1, chunk size, 1, BLOCK STATE , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, BLOCK STATE dt tile = ct.load dt, index= batch, chunk idx, head , shape= 1, chunk size, 1 , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, d tile = ct.load decay cumsum, index= batch, chunk idx, head , shape= 1, chunk size, 1 , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, d last = ct.extract d tile, index= chunk size - 1, , shape= 1, b dt d = b tile dt tile ct.exp2 d last - d tile :, None hid tile = ct.exp2 d last hid tile hid tile = ct.mma x tile.transpose , b dt d.astype b.dtype , hid tile 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. ct.store final state, index= batch, head, head tile idx, state tile idx , tile=hid tile.reshape 1, 1, BLOCK HEAD, BLOCK STATE , 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. chunk idx, head tile idx, batch head = ct.bid 0 , ct.bid 1 , ct.bid 2 batch = batch head // x.shape 2 head = batch head % x.shape 2 group = head // num heads per group cb tile = ct.full chunk size, chunk size , 0.0, dtype=ct.float32 out tile = ct.full chunk size, BLOCK HEAD , 0.0, dtype=ct.float32 Phase 1: build CB = C @ B.T and accumulate the carried-state output for state tile idx in range state size // BLOCK STATE : b tile = ct.load b, index= batch, chunk idx, group, state tile idx , shape= 1, chunk size, 1, BLOCK STATE , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, BLOCK STATE c tile = ct.load c, index= batch, chunk idx, group, state tile idx , shape= 1, chunk size, 1, BLOCK STATE , padding mode=ct.PaddingMode.ZERO, .reshape chunk size, BLOCK STATE h tile = ct.load hiddens, index= batch, head, chunk idx, head tile idx, state tile idx , shape= 1, 1, 1, BLOCK HEAD, BLOCK STATE , .reshape BLOCK HEAD, BLOCK STATE cb tile = ct.mma c tile, b tile.transpose , cb tile out tile = ct.mma c tile, h tile.transpose .astype c tile.dtype , out tile 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. Phase 2: decay + causal mask then final computation d tile = ct.load decay cumsum, index= batch, chunk idx, head , shape= 1, chunk size, 1 , padding mode=ct.PaddingMode.ZERO .reshape chunk size, dt tile = ct.load dt, index= batch, chunk idx, head , shape= 1, chunk size, 1 , padding mode=ct.PaddingMode.ZERO .reshape chunk size, x tile = ct.load x, index= batch, chunk idx, head, head tile idx , shape= 1, chunk size, 1, BLOCK HEAD , padding mode=ct.PaddingMode.ZERO .reshape chunk size, BLOCK HEAD d val = ct.load d, index= head, , shape= 1, cb tile = cb tile ct.exp2 d tile :, None - d tile None, : pairwise decay idx = ct.arange chunk size, dtype=ct.int32 cb tile = ct.where idx :, None = idx None, : , cb tile, 0.0 causal mask out tile = out tile ct.exp2 d tile :, None decay carried state dt x = dt tile :, None x tile out tile = ct.mma cb tile.astype ct.tfloat32 , dt x.astype ct.tfloat32 , out tile out tile = out tile + d val x tile D x skip ct.store out, index= batch, chunk idx, head, head tile idx , tile=out tile.reshape 1, chunk size, 1, BLOCK HEAD .astype out.dtype 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 https://github.com/NVIDIA/TileGym . A complete example is shown below. python import torch import tilegym from tests import common from tests.ops.test mamba2 import Mamba2ProblemSpec tilegym.set backend "cutile" num tokens, state size = 16384, 128 torch.manual seed 0 spec = Mamba2ProblemSpec batch size=4, num tokens=num tokens, num groups=1, num heads=8, state size=state size, head size=64, chunk size=128, dtype=torch.bfloat16, state dtype=torch.float32 fwd inputs = spec.get sample inputs with reference=False, with backward=False fwd = common.benchmark lambda: tilegym.ops.mamba2 chunk forward fwd inputs , use cupti=True, print f"Forward: {fwd 'forward' 'median' 1000:.2f} us" 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 https://developer.nvidia.com/cuda/tile and explore the documentation https://docs.nvidia.com/cuda/cutile-python/index.html , GitHub repos https://github.com/nvidia/cutile-python , and sample code https://github.com/NVIDIA/cutile-python/tree/main/samples .