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.