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.
重要度重大証拠E2 未複製執筆簡易
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.