ThunderKittens reaches 22.4 PFLOPS on NVIDIA Rubin

Together AI’s kernels team used early access to NVIDIA Vera Rubin NVL72 to extend ThunderKittens, its embedded CUDA domain-specific language, for Rubin’s instruction set. The resulting hand-tuned matrix-multiplication kernels reached 22.4 PFLOPS with NVFP4 and 12.0 PFLOPS with FP8, roughly matching NVIDIA’s cuBLASLt and CuTe implementations in Together’s tests.

The existing B200 kernels compiled and ran on Rubin without modification because the architecture preserves Blackwell’s programming model. Their performance reached only 42.1% of the NVFP4 roofline and 44.4% of the FP8 roofline. Rubin’s tensor cores consume operands about twice as quickly, so kernels tuned for Blackwell could not supply data fast enough.

NVFP4 is NVIDIA’s block-scaled 4-bit floating-point format, while FP8 uses 8-bit values. Both reduce memory traffic and increase tensor-core throughput compared with wider formats, making them relevant to large-model training and inference.

Rubin widens the data path

Together reports large increases in Rubin’s per-GPU compute capacity, memory bandwidth, multiprocessor count, and peak power:

For kernel authors, Rubin’s smaller architectural and instruction-set changes determine whether code can approach those compute limits:

• A wider MMA K step: Blackwell’s tcgen05.mma instruction consumes 32 bytes along the K, or reduction, dimension per step. Rubin supports 64 bytes in the same cycle count, doubling the work available in each instruction window.
• More tensor memory: The 128-lane tensor-memory space grows from 512 to 576 columns. The additional 32 KiB is accessible through PTX 9.4’s new .exclusive qualifier.
• Larger shared memory: An oversized mode can expand shared memory dynamically to 328 KiB, compared with 228 KiB on Hopper and Blackwell.
• A B-operand collector: Blackwell could retain an A tile in a staging buffer for reuse by the next matrix-multiply-accumulate instruction. Rubin adds the same capability for B, reducing a 2-by-2 MMA block from eight operand reads to five.
• Earlier A release: A new commit variant signals the producer barrier as soon as an MMA finishes reading A. The loader can then begin streaming the next tile before the current instruction retires.

Feeding the faster tensor cores

Enabling the 64-byte K encoding produced little improvement by itself because it increased tensor-core consumption without reducing memory traffic or expanding the pipeline. Together addressed both constraints through a new tile layout and a deeper shared-memory ring.

The original NVFP4 Blackwell kernel used 1-by-1 tiling. Covering a 512-by-256 output region required two independent cooperative thread array pair jobs, each loading its own copy of B. A 2-by-1 layout lets two accumulators along the M dimension share one B chunk, cutting that traffic in half.

Blackwell’s tensor-memory capacity had prevented this layout. Two 256-by-256 accumulators occupied all 512 columns, leaving no space for block-scaling factors. Rubin’s additional 64 columns provide enough room for those factors and make the shared-B layout possible.

A sweep of the shared-memory ring showed consistent gains as the pipeline deepened:

The B-side collector added another 1% to 3% by applying FILL to one MMA and LASTUSE to the next, roughly halving B-operand reads. Early A release produced larger gains on 64K-by-64K and 128K-by-128K NVFP4 matrices, improving throughput by 13.5% and 22.1%, respectively. At those sizes, A tiles were evicted from cache between reuses; smaller matrices already kept A in L2 and saw no benefit.

Source compatibility still needs tuning

The port demonstrates that ThunderKittens can preserve its programming model across GPU generations while exposing new hardware controls through template parameters. Existing B200 kernels ran unchanged, and developers could adopt Rubin’s wider MMA operations, added tensor memory, larger shared-memory allocation, and dual operand collectors without redesigning the DSL.

The results also quantify the limits of source compatibility. A kernel can execute correctly on a newer GPU while leaving most of its added throughput unused. Reaching Rubin’s higher roofline required less B traffic, more in-flight data, architecture-specific memory allocation, and earlier synchronization.

Developers working on FlashAttention-style kernels and other tiled tensor-core workloads can carry over the same decomposition and synchronization model from B200. The main retuning targets are tile shape, shared-memory depth, operand reuse, and barrier timing.

Results come with early-hardware limits

The tests ran on a qualification-sample GPU with CUDA 13.4. Together expects its kernels and NVIDIA’s library baselines to improve as the hardware and software stack mature, so the reported figures should be treated as early measurements rather than final production limits.

These figures measure matrix-multiplication kernel throughput for the tested shapes. They do not establish end-to-end model performance, latency, energy efficiency, or sustained throughput under production serving conditions. The peak-power increase from 1,000 W to 2,300 W also makes power measurements necessary for any efficiency comparison.

Together says LUT-based GEMMs, hardware-native megakernels, and additional engine optimizations remain under development. Its technical account provides the current implementation details and benchmark results.