NVIDIA’s Mamba-2 Kernels Run Up to 1.82 Times Faster in Tests

NVIDIA outlined a cuTile implementation of Mamba-2 that ran 1.14 to 1.82 times faster than a Triton reference in forward-pass tests, offering a measure of kernel performance on one Blackwell Ultra GPU rather than a full training run.

Jane Street engineers worked with NVIDIA to build and optimize the code using cuTile, NVIDIA’s Python interface for GPU programming. The comparison used a single NVIDIA B300 SXM6 GPU. It tested the cuTile kernels against an open-source Mamba-2 implementation written in Triton.

Mamba-2 processes sequences in chunks and carries a fixed-size hidden state from one chunk to the next. The architecture is designed to use matrix operations efficiently during training while keeping memory requirements fixed during inference. Jane Street uses deep learning in quantitative research.

The forward pass was divided among three kernels. chunk_cumsum computes cumulative decay sums within each chunk. fwd_pass_state carries the hidden state across chunks, while compute_out generates the outputs. The design combines operations to limit large intermediate data and keeps matrix calculations suited to NVIDIA Tensor Cores.

The implementation also uses log-space calculations, causal masking and scaling based on powers of two. NVIDIA’s optimization work included adjusting data layouts and memory access, tuning load and store behavior, and searching across cuTile settings. The benchmark used cuTile 1.5.0 and Triton 3.7.1.

The results cover only the forward pass. A complete training run also includes a backward pass, which was outside the benchmark and can affect total performance. NVIDIA plans to detail its backward-pass implementation in a later post.