Mamba-2's forward pass moves into Python kernels: 1.14x to 1.82x faster than Triton, with gains growing at longer sequences
Synopsis
Machine learning engineers at Jane Street, working with NVIDIA, rewrote the Mamba-2 forward pass in NVIDIA CUDA Tile (cuTile), a Python-based GPU programming model, and tuned it for Blackwell Tensor Cores. Mamba-2 combines a fixed-size recurrent state with attention-style matrix multiplication, but its forward pass still needs efficient kernels. On a single B300, the kernels run 1.14x to 1.82x faster than a Triton implementation, with speedups growing as sequence length and state size increase; these are kernel-level results, and end-to-end training also depends on the backward pass.
Interpretation
Mamba-2's forward pass now runs as three custom kernels written in Python, rather than hand-written low-level GPU code. Such forward passes have typically been carried by implementations such as Triton, while cuTile is a newer Python tile-level programming model still in early adoption. On a single NVIDIA B300 SXM6 (Blackwell Ultra), against the Triton reference in mamba_ssm 2.3.2.post1, the cuTile kernels are 1.14x to 1.82x faster, and the speedup grows with sequence length and state size.
Three kernels handle intra-chunk cumulative sums, per-chunk initial hidden states, and per-chunk outputs, keeping the sequential dependence inside one program instance. A token-by-token loop cannot be parallelized, while a parallel scan must materialize an enormous hidden-state tensor, making it memory-bound and unable to use Tensor Cores. chunk_cumsum computes cumulative decay sums within each chunk, fwd_pass_state advances the hidden state chunk by chunk with a jump formula, and compute_out fuses the output terms while keeping the interaction matrix in registers; slices over batch, head and state run in parallel across the grid.
The implementation uses cuTile latency hints, operand swapping, memory-layout transposition and autotuning to get close to the hardware. On Blackwell's fifth-generation Tensor Cores the matrix-multiply operands are asymmetric, with A allowed in Tensor Memory but B required in Shared Memory, and default strided access lowers to cache-unfriendly loads. A transpose identity routes the fused operand to slot A to relieve shared-memory contention, decay_cumsum and dt are transposed to a token-contiguous layout to enable TMA, and cuTile exhaustive search picks block size and occupancy per input shape.
Perspective
The results cover kernel-level speedups for the Mamba-2 forward pass only, measured on a single NVIDIA B300 SXM6 at compute capability 10.3 with cuda-tile 1.5.0, Triton 3.7.1 and PyTorch 2.13.0. The implementation is available through TileGym, and the example runs a bfloat16 configuration with batch 4, 16384 tokens, 8 heads, state size 128 and chunk size 128. A follow-up post will describe the backward pass in cuTile.
The speedup figures are kernel-level measurements, and end-to-end training results also depend on the backward-pass implementation. Some exponentials can be positive when computed and are immediately removed by the causal mask so overflow does not poison the result. Bottlenecks shift with input shape, so autotuning is needed rather than one fixed optimization set.