Jane Street ML Engineers Map Mamba-2 Recurrent State Computation to Tensor Cores with cuTile
Jane Street machine learning engineers Sylvain Gugger and Axel Feldmann, together with NVIDIA's Yinuo Liu and Siddharth Samsi, explain in detail on the NVIDIA Developer Blog the three-kernel decomposition of the Mamba-2 forward pass: chunk_cumsum computes cumulative sums of the decay term in log-space, fwd_pass_state advances the hidden state across chunks via the jump formula using matrix multiplication, and compute_out fuses the carried-state term, the decayed and causally masked interaction matrix, and the skip connection into a single tiled kernel. cuTile's Python DSL abstracts thread-level scheduling behind tile-centric primitives, enabling this mathematical structure to map directly onto Tensor Cores. Self-reported kernel-level speedups range from 1.14x to 1.82x, but apply only to the forward pass and depend on a backward pass implementation not yet published.
BedeutungWesentlichBeweisE2 nicht repliziertAufbereitungSchnell
Sequence modeling architectures face a fundamental computational trade-off: standard softmax attention accelerates training via matrix multiplication on Tensor Cores, but at inference compute and memory requirements grow linearly with sequence length; recurrent neural networks have constant compute and fixed memory at inference, but training cannot be parallelized. Mamba-2 belongs to a family of state-space models aiming to achieve both — it has a fixed-size hidden state like an RNN, but the state update can be computed chunkwise in parallel.
Jane Street machine learning engineers Sylvain Gugger and Axel Feldmann, collaborating with NVIDIA's Yinuo Liu and Siddharth Samsi, implemented the Mamba-2 forward pass as cuTile (CUDA Tile) Python DSL kernels and presented the full derivation on the NVIDIA Developer Blog. cuTile abstracts thread- and warp-level scheduling behind tile-centric primitives; the compiler automatically maps operations onto advanced hardware features such as Tensor Cores.
The core mechanism is the jump formula. Ignoring the decay term, advancing the hidden state from time t to t+T simply requires accumulating the outer products of intermediate tokens: h(t+T) = h(t) + BX^T, where B and X are matrices formed by stacking the b and x vectors for each token within a chunk. This is directly a matrix multiplication, so it can execute on Tensor Cores. When decay is reintroduced, cumulative sums in log-space replace sequential multiplications with additions; pre-multiplying the rows by exponentials preserves the matrix-multiplication form.
This derivation decomposes into three kernels. chunk_cumsum uses ct.cumsum to compute per-token, per-head cumulative sums of the decay term within each chunk, storing results for reuse by the subsequent two kernels. fwd_pass_state is a sequential loop over chunks; each iteration uses ct.mma to execute the jump formula and advance the hidden state, multiplying the decay-pre-multiplied b operand with the transpose of x. compute_out is the final output computation, fusing the carried-state term, the decayed and causally masked interaction matrix, and the skip connection into a single tiled kernel that keeps intermediate matrices in registers and casts to TF32 for Tensor Core execution.
The compute_out design deserves attention. A direct PyTorch implementation would compute the three terms separately and write the full [B, n_chunks, chunk_size, chunk_size] interaction matrix to global memory. The cuTile version fuses all three terms; the interaction matrix cb_tile = C @ B^T is built incrementally via ct.mma in Phase 1, while Phase 2 applies decay and the causal mask. The causal mask is constructed using ct.arange and ct.where as a lower-triangular condition, so no mask tensor is ever materialized. In the decay matrix exp(G - G^T), positions where i < j produce positive exponents, but the causal mask zeroes them immediately, guaranteeing numerical stability.
Hardware optimizations targeting the Blackwell fifth-generation Tensor Core include: operand swapping — Tensor Core MMA operands are asymmetric, with operand A able to reside in Tensor Memory while operand B must reside in Shared Memory; the identity B^T A^T = (AB)^T is used to route the fused operand to the A side, relieving shared-memory contention. Memory layout transposition — decay_cumsum and dt are transposed from the default strided layout to [batch, head, token] to ensure contiguous access and enable TMA. And the cuTile autotuner, which searches block size and occupancy parameters for each input shape.
Self-reported performance results show the optimized cuTile forward-pass kernels are 1.14x to 1.82x faster than the open-source Triton implementation, with speedups growing as sequence length and state size increase. The test environment is a single NVIDIA B300 SXM6 (Blackwell Ultra, 1920 MHz max SM clock, 1100 W power limit, 268 GiB HBM), benchmarked against the mamba_ssm 2.3.2.post1 Triton reference implementation.
These speedups are kernel-level only. End-to-end training performance also depends on the backward pass implementation; the authors preview a follow-up post describing how the backward pass is implemented in cuTile. cuTile currently targets NVIDIA GPUs exclusively, and the forward-pass implementation is available through the TileGym library.
From a broader perspective, the value of this article lies not only in the performance numbers but in its complete demonstration of how Mamba-2's mathematical structure maps onto GPU hardware primitives: log-space cumulative sums correspond to ct.cumsum, the jump formula corresponds to ct.mma, and the causal mask corresponds to register-level conditional selection. This provides an inspectable concrete mechanism for understanding why state-space models can maintain RNN-style constant inference memory while achieving attention-like training parallelism.