Mamba-2 offers a different scaling path
NVIDIA and Jane Street engineers have published an implementation study showing how CUDA Tile can accelerate key parts of the Mamba-2 state-space model. Unlike attention, which compares tokens across a sequence, Mamba-style models update a fixed-size hidden state. Their chunkwise formulation can support long sequences with different memory and computation characteristics from a conventional Transformer.
The post focuses on engineering rather than announcing a new model. The team rewrote three forward-pass operations: chunk_cumsum, fwd_pass_state and compute_out. Together they cover cumulative parameters inside chunks, state propagation and output construction. Optimising these stages matters because an efficient high-level architecture can still underperform if its kernels move data poorly or fail to use specialised GPU hardware.
CUDA Tile lifts programming to tile operations
CUDA Tile and its cuTile Python interface let developers express computation over multidimensional tiles rather than scheduling individual threads directly. The compiler maps those operations to GPU resources and can target Tensor Cores for matrix work. This raises the abstraction level while preserving access to layout, memory and latency decisions that determine real performance.
A higher-level model does not mean performance appears automatically. Engineers still need to understand the shapes, reuse patterns and hardware limits of their workload. The advantage is that they can express intent in composable tile operations and explore alternatives without hand-writing every synchronisation step. That can make specialised kernels easier to maintain as model dimensions and GPUs evolve.
Memory layout is part of the algorithm
The implementation uses transposed global layouts so the Tensor Memory Accelerator can move data in a form better suited to the next operation. It also swaps operands in selected matrix multiplications to improve placement between Tensor Memory and shared memory. These changes do not alter Mamba-2’s mathematical result, but they can reduce transfers, conversions and stalls.
This is a recurring lesson in accelerator work: two programs with identical arithmetic can perform very differently because of where data lives and how it travels. Teams evaluating a custom kernel should inspect memory traffic and occupancy alongside raw operation counts. They should also compare numerical results across sequence lengths and data types, because a faster layout is useful only if it preserves acceptable accuracy.
Hints and autotuning search the hardware space
CUDA Tile supports latency hints that help the compiler overlap work, while an autotuner can test configurations such as tile sizes and pipeline choices. The authors use these facilities rather than assuming one launch setup is optimal for every shape. Automated search is particularly valuable for Mamba-2 because batch size, sequence length, head dimensions and chunking change the balance of compute and memory.
Autotuning results should be tied to the actual GPU, software stack and representative shapes. A configuration selected on one device can regress elsewhere, and a benchmark dominated by a single convenient dimension may not reflect production. Teams need a fallback configuration, reproducible tuning records and correctness tests that run for every selected variant.
The forward pass is progress, not the whole training story
The published work covers the forward pass, and the authors identify backward kernels as future work. That boundary is important for organisations considering training or fine-tuning: end-to-end improvement depends on gradients, optimiser work and the rest of the model pipeline. Inference-heavy uses can benefit sooner, but even there overall latency includes data preparation, surrounding layers and framework overhead.
The collaboration is a useful example of domain engineers and platform specialists refining an emerging architecture together. CUDA Tile aims to retain enough hardware control for serious optimisation while reducing the burden of thread-level code. The strongest evidence will come from reproducible end-to-end benchmarks and maintained backward support. For now, the three kernels show that tile-level programming can translate Mamba-2’s theoretical scaling advantages into more practical GPU execution.
What teams should test next
Portability will be another measure of the programming model. CUDA Tile is valuable if engineers can carry the intent of an optimisation across GPU generations without repeatedly rebuilding thread-level machinery. That does not remove the need for architecture-specific tuning, but it can narrow the portion of code that must change. Open benchmarks should report compiler and driver versions, warm-up, input distributions and numerical tolerance. Those details let other teams distinguish an improvement in the kernel from a favourable benchmark setup and help the ecosystem identify where the abstraction still needs lower-level escape hatches.
Procurement and engineering teams should record the source date, product version, supported region and evaluation conditions before relying on the result. They should then repeat the same representative tests after a model, compiler, connector or policy changes. Published performance and capability claims are useful starting points, but production confidence comes from evidence collected in the organisation’s own environment. A named owner, measurable acceptance threshold and documented fallback turn a promising technical release into a service that can be governed over time.