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

Engineers at the trading firm Jane Street, working with NVIDIA, rewrote the forward pass of Mamba-2, a sequence model that carries a fixed-size hidden state from token to token like a recurrent network, as GPU kernels in cuTile, NVIDIA's Python language for writing kernels in terms of tiles rather than individual threads (NVIDIA Technical Blog). The core problem is that stepping through tokens one at a time is fine for generation but too slow for training and prefill. The post's approach splits a sequence into chunks and uses a formula that jumps the hidden state from the start of one chunk to the next with a single matrix multiplication, which is the kind of work Tensor Cores accelerate. The per-token decay is handled as cumulative sums in log space, and the authors note that the exponentials they keep are never of positive numbers, which keeps the computation numerically stable.

That becomes three kernels: one computes the cumulative decay sums within each chunk, one walks across chunks to produce each chunk's starting hidden state, and one computes the outputs for every chunk in parallel, fusing the carried-state term, the causally masked interaction matrix and the skip connection so the intermediate matrix stays in registers instead of global memory. Further tuning included compiler latency hints, rearranging operands to suit the Blackwell Tensor Core's memory layout, a transposed memory layout for contiguous loads, and an autotuner that picks settings per input shape. On a single B300 GPU, the authors report the cuTile forward pass is 1.14x to 1.82x faster than an open-source Triton implementation of Mamba-2, with larger gains at longer sequences and bigger state sizes. They stress these are kernel-level numbers; end-to-end training speed also depends on the backward pass, which a follow-up post is meant to cover. The kernels are available through NVIDIA's TileGym.