NVIDIA and Jane Street released three GPU kernels, code that performs specific GPU operations, to accelerate Mamba-2's forward pass from inputs to outputs. Developers optimizing models can use the public Python code. According to the company, it runs 1.14~1.82× faster than existing open-source kernels.

1. Turning recurrent computation into chunk-based work

Mamba-2 is a recurrent model that updates a fixed-size internal memory, or hidden state, with previously processed information. During training or prompt processing, it converts chunks of tokens into matrix multiplications rather than computing one token at a time. This enables use of Tensor Cores, the GPU's dedicated matrix computation units.

NVIDIA's cuTile is a tool for writing GPU kernels in Python. Developers define operations on tiles, or groups of data, and leave detailed execution layout to the compiler. Jane Street engineers used it to implement Mamba-2 and optimized the kernels with NVIDIA.

2. Three kernels split state and output computation

The first kernel, chunk_cumsum, accumulates and stores the product of the decay rate and time interval within each chunk. The decay rate represents how much previous information fades. The next two kernels reuse these values.

The second, fwd_pass_state, computes the hidden state at each chunk's start and the state after the last chunk. It preserves chunk order while processing independent state segments in parallel.

The third, compute_out, produces token outputs from the chunk's starting state. Rather than writing intermediate matrices separately to large global memory, it fuses computation and keeps them in registers, small storage spaces near the compute units.

The blog's cumulative-sum kernel shows how tiles are loaded, accumulated within a chunk, and stored:

    dt_tile = ct.load(dt, index=(batch, chunk_idx, head_tile_idx),
                          shape=(1, chunk_size, BLOCK_H),
                          padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
    a_tile = ct.load(a, index=(head_tile_idx,),
                        shape=(BLOCK_H,),
                        padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)

    d_tile = ct.cumsum(a_tile * dt_tile, axis=1) * LOG2E
    ct.store(output, index=(batch, chunk_idx, head_tile_idx), tile=d_tile.astype(output.dtype))

3. Hardware-specific memory and matrix multiplication optimizations

The implementation attaches expected latency information to memory reads and writes so the compiler can adjust instruction order. It reorders multiplication inputs to suit the input layout of Blackwell's 5th-generation Tensor Cores. This reduces shared-memory contention, where multiple tasks compete for storage. It also changes the memory layout of accumulated decay values and time intervals to enable contiguous reads.

Since bottlenecks vary by input shape, it automatically searches block sizes, the amount of data processed at once, and GPU resource utilization settings. This reduces execution time but requires searching configuration combinations for each input shape.

4. Company kernel benchmarks show up to a 1.82× speedup

NVIDIA's blog benchmarks used cuTile 1.5.0 on a single B300 SXM6. The GPU's maximum compute clock was 1920 MHz, with 268 GiB of high-bandwidth memory.

The comparison used open-source kernels written in Triton 3.7.1, the NVIDIA Dynamo-family reference implementation in mamba_ssm 2.3.2.post1.

The benchmarks used PyTorch 2.13.0 and CUDA 13.3.

According to the company, cuTile kernels run 1.14~1.82× faster than Triton, with speedups increasing as sequence length and state size grow. The chart covers 1K~256K tokens and state sizes of 64, 128, 256, 512, but individual bars lack exact numerical labels.

The blog explains that this is a forward-pass kernel comparison. Overall training performance also depends on the implementation of the backward pass, which traces computation from outputs back toward inputs.

5. Explore the implementation in TileGym

The code and benchmarks are available in NVIDIA's TileGym repository. cuTile is the default execution mode. Follow the repository's installation instructions to explore computation examples and performance measurements.

The repository specifies cuTile 1.5.0 or later and the compiler required at runtime. The installation guide lists support for CUDA 13.1 or later on Blackwell GPUs and CUDA 13.2 or later on Ampere GPUs.

This release provides code showing that matrix formulation and hardware-specific tuning of the Mamba-2 forward pass can deliver measured gains.