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

NVIDIA published a deep technical walkthrough of three custom GPU kernels that power the forward pass of Mamba-2, a state-space sequence model, written in cuTile, its Python-based language for tile-level GPU programming. The work came from machine learning engineers at Jane Street, the quantitative trading firm, who were early adopters of cuTile and then optimized the kernels together with NVIDIA.

On a single B300 Blackwell Ultra GPU, NVIDIA reports the cuTile kernels ran 1.14x to 1.82x faster than an open-source Triton implementation of Mamba-2, with the gap widening at longer sequence lengths and larger state sizes. The code is available through NVIDIA's TileGym project.

This matters beyond one architecture. Mamba-style models promise constant memory at inference no matter how long the input gets, which is exactly the property long-context agents and streaming workloads want. Their weak spot has been fast, hardware-friendly training kernels. Every credible speedup moves them closer to being a practical alternative to pure attention.

We like this story because it is the unglamorous plumbing that decides which model designs actually win. Attention dominates partly because it maps cleanly onto matrix-multiply hardware. Mamba-2 gets around that by chunking the sequence, jumping the hidden state forward with matrix products, and handling the decay terms in log-space so the math stays on Tensor Cores. The post shows how that idea turns into three kernels: one for per-chunk cumulative sums, one that carries state across chunks, and one that fuses the output terms so a large interaction matrix never leaves registers.

The more strategic point is the language. cuTile lets engineers describe work in tiles and leaves thread and warp scheduling to the compiler. The kernels shown here are short and readable Python. If that holds up across more workloads, writing serious GPU code stops being a craft for a handful of CUDA specialists, and more teams can experiment with architectures that are not attention.

Now the caveats, and they are real. These are kernel-level speedups only, and NVIDIA says so plainly: end-to-end training gains depend on a backward pass that has not been published yet. The comparison is against one Triton reference on NVIDIA's newest hardware, measured by NVIDIA. And several of the tricks, like swapping operands so a fused input lands in Blackwell's Tensor Memory, are tied to NVIDIA's specific chip design.

That last point is where this leads if it scales. Triton has been a relatively portable layer that other chipmakers can target. A Python DSL that wins by exploiting Blackwell-specific features deepens the moat around NVIDIA hardware, even as it lowers the bar for developers. Both things can be true. Our read is that this is good for model diversity and good for NVIDIA's lock-in at the same time, and buyers should notice the second half.

I'm also watching who is doing this work. A trading firm co-authoring kernel research tells us the frontier of efficient sequence modeling is not just inside the big labs. Finance has always cared about long, fast streams of data, and state-space models fit that shape.

Mamba-2 belongs to a family of models that keep a fixed-size hidden state, like a recurrent network, but can be trained in parallel chunks. Without its decay term, the post notes, Mamba-2 reduces to a form of linear attention. cuTile is NVIDIA's tile-centric Python DSL; the benchmarks used cuda-tile 1.5.0 against mamba_ssm 2.3.2.post1 on Triton 3.7.1.