Jane Street and NVIDIA Release Mamba-2 cuTile Kernels
Jane Street engineers and NVIDIA implemented three custom kernels for the Mamba-2 forward pass in CUDA Tile (cuTile), a Python-based DSL and parallel programming model for tile-level NVIDIA GPU programming. NVIDIA reports that the kernels were 1.14x to 1.82x faster than an open-source Triton implementation in tests on a single NVIDIA B300 SXM6. The results measure kernel performance only; end-to-end training performance will also depend on backward-pass implementations.
The work provides a concrete cuTile implementation of Mamba-2 forward-pass kernels and reports measurable speedups over a Triton implementation on NVIDIA Blackwell Ultra hardware, while the results remain limited to kernel-level benchmarks.
NVIDIA CUDA Tile (cuTile) is a Python-based DSL and parallel programming model for NVIDIA GPUs that automatically leverages advanced hardware capabilities to deliver peak performance for tile-level GPU programming. By abstracting thread- and warp-level scheduling behind high-level, tile-centric primitives, cuTile enables you to write GPU kernels in Python while compiler optimizations automatically target advanced hardware features like Tensor Cores.
As an early adopter of cuTile, machine learning engineers at Jane Street implemented Mamba-2, a state-of-the-art sequence model, in cuTile and partnered with NVIDIA to optimize the kernels. Jane Street is a deeply technical global trading form, using deep learning to drive innovation in advanced quantitative research. This post walks you through the three custom kernels powering the Mamba-2 forward pass and the cuTile optimizations behind them.
Overview and prerequisites
To use cuTile, ensure your environment meets the requirements outlined in the cuTile Python Quickstart guide. Import the package as cuda.tile and expose it as ct throughout the kernel implementations featured in this post.
Sequence modeling computational trade-off
Sequence modeling architectures face a computational trade-off: traditional softmax attention trains rapidly using matrix multiplication accelerated by Tensor Cores, but at inference, compute and memory requirements increase linearly with sequence length. In contrast, recurrent neural networks (RNNs) suffer from slow, nonparallelizable training but require constant compute and fixed memory footprint regardless of sequence length.