NVIDIA cuTile Mamba-2 Kernels: Up to 1.82x Faster Than Triton on B300

NVIDIA has published a walkthrough of Mamba-2 forward-pass kernels written in CUDA Tile, or cuTile, its Python-based GPU programming model. On one Blackwell Ultra B300, the company reports its optimized kernels run 1.14x to 1.82x faster than the open-source Triton implementation shipped with mamba_ssm. The post, on the NVIDIA Developer blog, is as much a tutorial on how to write a state space model kernel as it is a benchmark, and the code lives in the open-source TileGym repository. NVIDIA credits Jane Street ML engineers with implementing Mamba-2 in cuTile and working with NVIDIA on the kernels.

The figures cover the forward pass only, on one GPU, against one baseline. That is still a useful signal: Mamba-style layers now appear inside hybrid models from several labs, including NVIDIA's own, and kernel quality decides whether those layers are fast in practice. We covered the same family of ideas in NVIDIA ReplaySSM and in the Nemotron Labs TwoTower diffusion LLM.

TL;DR: what readers will ask

Why Mamba-2 needs custom kernels at all

Softmax attention trains well on Tensor Cores but its inference cost grows with sequence length. Classic recurrent networks have constant inference cost but are hard to parallelize during training. Mamba-2 sits between the two: it keeps a fixed-size matrix hidden state with a recurrent update, where each step decays the previous state and adds a new outer-product term, and computes outputs from that state plus a skip term.

Run naively as a parallel scan, the model would have to materialize a very large hidden-state tensor. That makes the computation memory-bound and keeps it off the Tensor Cores. The standard answer, and the one NVIDIA follows, is a chunked formulation. A "jump formula" gets the state at each chunk boundary with matrix multiplies, and outputs inside a chunk come from the chunk's starting state plus a causal-masked interaction matrix. Working in log space turns products of decays into cumulative sums, and masking keeps only exponentials of non-positive values, which keeps the math numerically stable.

Chip and compute graphic illustrating how a Mamba-2 state space model carries state across sequence chunks on a GPU

The three kernels, in plain language

NVIDIA splits the forward pass into three kernels.

1. chunk_cumsum computes per-head cumulative sums of the decay term within each chunk. It scales by log2(e) so the cheaper exp2 instruction can be used later, and stores the results for the next two kernels.
2. fwd_pass_state loops sequentially over chunks inside each batch, head, head-tile and state-tile slice to produce each chunk's initial hidden state and the final state. It uses ct.mma for Tensor Core matmuls, and independent slices run in parallel across the grid. The final state can be fed back as init_state for inference.
3. compute_out produces the output tile for each batch, chunk, head and head-tile. It fuses the carried-state term, the decayed and causally masked interaction between C and B transposed, and the D times x skip connection. The interaction matrix stays in registers instead of round-tripping through global memory, and masking uses ct.arange and ct.where, so no mask tensor is ever materialized. The last matmul casts to ct.tfloat32 to use Tensor Cores at near-FP32 accuracy, and exponentials of positive values are masked out before they can overflow.

If you are new to why fusing steps like this matters, our explainer on AI chip architectures covers the memory hierarchy that makes keeping data in registers so valuable.

The optimizations behind the speedup

NVIDIA lists four tuning steps beyond the basic port.

• Latency hints. Annotating ct.load and ct.store with latency hints helps the compiler model DRAM traffic, improving instruction scheduling and pipeline buffer depth.
• Swapping the A and B operands. On Blackwell's fifth-generation Tensor Cores, operand A can come from Tensor Memory but operand B must come from Shared Memory. Using the identity that the transpose of a product equals the product of the transposes in reverse order, the kernel moves the fused operand into the A slot, which reduces shared-memory contention.
• Global-memory layout. The decay cumulative sum and dt are transposed to a batch, head, token layout, so accesses are contiguous and can use TMA instead of cache-unfriendly strided loads.
• Autotuning. Block size and occupancy are searched together with the options above using exhaustive_search, picking the best configuration for each input shape.

A grid of chip squares with one lit row, representing GPU tile-level work in the cuTile Mamba-2 kernels

What was measured, and what was not

The benchmark machine was a single B300 SXM6 (compute capability 10.3) with a 1920 MHz maximum SM clock, a 1100 W power limit and 268 GiB of HBM. Software was cuda-tile 1.5.0, tileiras 13.4, PyTorch 2.13.0, CUDA 13.3 with driver 595.58.03, and Python 3.12.3 on Ubuntu 24.04.4. The baseline was the open-source Triton implementation from mamba_ssm 2.3.2.post1 with Triton 3.7.1. The example workload used batch size 4, 16,384 tokens, one group, eight heads, state size 128, head size 64 and chunk size 128, with bf16 inputs and fp32 state.

The headline range is therefore a kernel-level result on one machine. Things the post does not establish, at least in the text we could read:

• End-to-end training speed. The backward pass is not covered yet.
• Absolute latencies. The article text gives ratios; the raw timings are in a chart image.
• Other GPUs. The optimizations lean on Blackwell features, and no results on earlier generations are described.
• Fairness of the baseline. The comparison is against one Triton version. A newer or hand-tuned Triton kernel could narrow the gap, and the post is NVIDIA comparing its own language to a competing approach.

The authors say speedups grow with sequence length and state size, which is the regime where long-context hybrid models live.

How to try it

The implementation is in TileGym. The example in the post sets the backend with tilegym.set_backend("cutile"), calls tilegym.ops.mamba2_chunk_forward, and benchmarks with common.benchmark(..., use_cupti=True). For the language itself, NVIDIA points to the cuTile overview, the Python quickstart for setup requirements, the API docs, the open-source repo and its sample kernels.

A reasonable first experiment is to run the TileGym example on your own GPU, then change the sequence length and state size to see whether the reported trend holds for your shapes. Compare against the mamba_ssm Triton path with the same CUPTI-based timing, so you are measuring kernel time and not Python overhead.

Why it matters for builders

Three reasons stand out. First, state space and hybrid architectures are no longer curiosities. We have seen them in small-model experiments, in NVIDIA's Nemotron 3.5 Lightning MoE, and in decoding research like ReplaySSM. Faster kernels make them cheaper to train and serve.

Second, the language matters as much as the kernel. NVIDIA is pushing several ways to write GPU code beyond hand-tuned CUDA C++; our piece on CUDA kernels in Rust covers another. A Python tile model that competes with Triton on a hard kernel is NVIDIA's argument that you can stay in Python without leaving speed on the table.

Third, it shows how much of the AI performance story is now software. Hardware peaks are only reachable when kernels are written for the memory system, which is why the operand swap and layout changes above produce real gains. The same theme runs through Amazon's TorchNeuron and vLLM's KV transfer work.

A stack of memory wafers feeding a small chip, illustrating the memory-bound nature of the Mamba-2 forward pass kernels

Caveats before you rewrite anything

Do not rip out a working Triton kernel on the strength of a vendor chart. The gains were measured in one configuration on the newest NVIDIA data-center part, the code is tied to a young language whose APIs are still moving (note the version numbers above), and an unmeasured backward pass can erase forward-pass wins in a training loop. If you serve Mamba-2 or hybrid models on Blackwell Ultra, the forward-pass result is directly relevant. If you train on older hardware, wait for the follow-up and for independent numbers.

Which readers should care now

If you maintain inference infrastructure for hybrid or state space models on Blackwell Ultra, put the TileGym example on your benchmark list this week and record kernel time at your real sequence lengths. If you write kernels for a living, read the post for the technique rather than the number: the operand swap, the transposed layout for TMA, and keeping the interaction matrix in registers are ideas that transfer to other chunked recurrent layers. If you only consume models through an API, nothing changes for you today; the effect, if it materializes, shows up later as lower serving cost for architectures that use these layers. Either way, wait for the backward-pass follow-up and independent reproductions before treating the headline as settled.

Related reading

• NVIDIA ReplaySSM: cache SSM inputs, not state
• NVIDIA Nemotron Labs TwoTower diffusion LLM
• Nvidia's push to write GPU kernels in Rust
• AI chip architectures explained
• BitNet on a 1975 6502 with Mamba
• vLLM push-based and bi-directional KV transfer

Versions, benchmark figures and repository details are as reported in NVIDIA's post as of October 19, 2026.