Accelerating Mamba-2 with High-Performance Forward Pass Kernels in NVIDIA CUDA Tile

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. 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. The post describes three custom kernels powering the Mamba-2 forward pass and the cuTile optimizations behind them.