The Pulse
NVIDIA’s cuTile Kernels Run Up to 1.82× Faster on Mamba-2
NVIDIA and Jane Street built Mamba-2 forward-pass kernels with NVIDIA CUDA Tile, reporting speedups of 1.14× to 1.82× over a Triton implementation. The comparison ran on a single NVIDIA B300 and measures kernels only, not end-to-end trainin
AI.info Team · Written with AI · fact-checked ·
NVIDIA says a Mamba-2 forward pass written with its CUDA Tile programming model runs between 1.14× and 1.82× faster than an open-source Triton implementation. The company published the results on October 19, describing work by Jane Street machine-learning engineers in partnership with NVIDIA. The comparison measures forward-pass kernels, not complete training runs.
Three kernels carry the Mamba-2 forward pass
The implementation uses cuTile, NVIDIA’s Python-based programming model for writing GPU kernels with tile-oriented operations. Rather than calculate each token’s hidden state in a long serial loop, Mamba-2 can process sequence chunks with matrix multiplications suited to GPU Tensor Cores. That approach makes the kernel implementation—and the way it moves data through the GPU—central to the benchmark.
The authors divide the forward pass among three kernels. chunk_cumsum calculates cumulative decay values within each chunk; fwd_pass_state advances the hidden state between chunks; and compute_out combines that state with the chunk’s inputs to produce outputs. The separation lets the implementation reuse decay calculations while assigning the main matrix operations to Tensor Cores.
The 1.82× result comes from a B300 test
NVIDIA compared its cuTile implementation with an open-source Mamba-2 kernel written in Triton. The tests ran on one NVIDIA B300 SXM6, identified in the post as a Blackwell Ultra GPU, with 268 GiB of HBM. The software setup included CUDA 13.3, PyTorch 2.13.0, cuTile 1.5.0 and Triton 3.7.1.
The reported speedups range from 1.14× to 1.82×, with larger gains as sequence length and state size grow. NVIDIA’s example benchmark uses a sequence length of 16,384 tokens and a state size of 128, alongside a batch size of four and eight heads. Those values describe the example configuration in the post; the headline speedup range summarizes results across tested configurations.
“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.”
Sylvain Gugger, Axel Feldmann, Yinuo Liu and Siddharth Samsi, NVIDIA Technical Blog
Memory layout and tuning shape the gains
The post describes several optimizations alongside the kernel design. The authors use latency hints to guide compiler scheduling for memory operations, rearrange some matrix operands to reduce shared-memory contention, and transpose data layouts to enable more contiguous access. An autotuner searches block-size and occupancy settings for each input shape instead of relying on one configuration for every workload.
That tuning matters because the reported margin varies with sequence length and state size. NVIDIA’s comparison therefore shows the performance of a specific implementation on a specified GPU and software stack, not a universal speed advantage for cuTile over Triton on every Mamba-2 workload.
The benchmark stops at the forward pass
The main limitation is explicit in the post: these are kernel-level measurements. End-to-end training also depends on the backward pass, which the benchmark does not evaluate. A faster forward kernel can reduce one part of a workload, but the reported 1.82× figure does not establish a matching improvement in total training time.
NVIDIA says the implementation is available through TileGym. For now, the concrete result is a forward-pass comparison on one B300: cuTile kernels measure 1.14× to 1.82× faster than the Triton reference across the tested configurations.