Mamba-2 in NVIDIA cuTile: Faster Forward Kernels and What the Benchmark Means

NVIDIA and Jane Street published a Mamba-2 optimization walkthrough on October 19, reporting forward-pass kernels that run 1.14 to 1.82 times faster than their Triton reference. The useful development is an inspectable implementation, not a new model or a promise that training becomes 82% faster.

The technical walkthrough tests a single B300 SXM6 GPU. It uses cuTile 1.5.0, PyTorch 2.13.0 and a mamba_ssm 2.3.2.post1 reference with Triton 3.7.1. These are the authors' measurements, not independently reproduced results.

Why Mamba-2 is a useful kernel example

Mamba-2's original research dates to May 2024. Its state-space duality framework connects state-space models with forms of attention, making efficient matrix-based algorithms possible. The architecture and this week's implementation work are different contributions.

The forward implementation splits the sequence into chunks. One kernel prepares cumulative decay values, another calculates the state entering each chunk, and a third produces the outputs. Intermediate interactions stay in registers rather than being written as a large matrix to GPU memory.

That distinction matters when evaluating alternative architectures. A mathematically efficient operation can still be slow if its implementation moves too much data or makes poor use of the available compute. Counting operations is a starting point, not a deployment benchmark.

What the speedup does and does not establish

The reported improvement is at kernel level. NVIDIA explicitly separates it from end-to-end training, which also depends on the backward pass. The optimization includes memory-layout changes, hardware-aware matrix multiplication and tuning for different input shapes.

For a team evaluating the result, the first question should be how much total runtime these kernels currently occupy. Consider an illustrative workload where half the time is spent in the part being optimized. Even doubling that part would reduce total time by only a quarter, assuming everything else stays unchanged.

The same reasoning applies to serving. Request preparation, other model layers, transfers and scheduling can dilute a kernel gain. A useful evaluation should therefore record both the operation's timing and the latency or throughput of the complete job. Neither number replaces the other.

How to evaluate the implementation

The implementation is available through NVIDIA's TileGym repository, which contains kernel examples, operation tests and benchmark utilities. Its documented setup calls for CUDA 13.1 or later on Blackwell; Ampere support requires CUDA 13.2 or later. It lists PyTorch 2.9.1 or a compatible version and cuda-tile 1.5.0 or newer.

Those are repository requirements, not a reproduction of the B300 benchmark environment. Keep that separation explicit when comparing results on a different GPU.

TileGym's installation option with the tileiras extra includes the runtime compiler dependencies. For a first evaluation, use an isolated environment and record the repository revision, GPU, compiler and framework versions. Run correctness tests before performance tests, then compare the sequence lengths, state sizes and batch sizes your application actually uses.

I would also keep compilation and tuning time separate from steady-state execution. A configuration that wins after extensive tuning may still be unattractive for a workload that changes shape frequently.

Do not confuse the benchmark with cuTile 1.7

Separately, cuTile 1.7.0 was released on October 19. It adds Windows arm64 wheels, more Python language support and new diagnostics. Text-format TileIR dumping requires CUDA Toolkit 13.4 or later.

The release also adds a cache-flushing option to exhaustive-search autotuning, enabled by default. That is relevant to measurement: developers should document benchmark settings rather than compare results produced under different cache conditions. The Mamba-2 figures above used cuTile 1.5.0, so they should not be presented as a 1.7 upgrade benefit.

For developers working on state-space models, this is a useful implementation to inspect and test. The next meaningful result is whether the gain survives integration into a complete workload at the required numerical accuracy. That is where a faster kernel becomes a reason to change a production stack.

I'm Paras. I work at the intersection of AI, engineering and distribution.

I help technical teams turn implementation work and product launches into clear stories that developers and founders can understand and share.

Book a 30-minute call.