Together AI Ports ThunderKittens GEMM Kernels to NVIDIA Vera Rubin NVL72

Together AI's kernels team received access to the NVIDIA Vera Rubin NVL72 platform and spent several days examining the new instruction set architecture with microbenchmarks. The team added functionality to ThunderKittens for writing NVFP4 and FP8 GEMMs on Vera Rubin, and published results on September 10, 2026.

The starting point was an existing Blackwell GEMM. NVIDIA's Blackwell architecture introduced the tcgen05 instruction, issued by a single thread, which lets one small producer warp drive the tensor cores. The accumulator moved out of registers into Tensor Memory, and operands are read directly from shared memory, enabling a single MMA to span two CTAs across two SMs. The Blackwell kernel launched threadblock clusters so each CTA pair shares operands via TMA multicast, cutting HBM traffic in half. Warps specialized within clusters: loaders bring A and B into shared memory, one MMA warp drives the tensor cores, and a consumer warpgroup moves finished accumulators to HBM. The kernel ran persistently, with one tile's inputs streaming in while the previous tile's outputs drained.

Because Vera Rubin preserves the Blackwell programming model, the old GEMMs still run. But naively running them on Vera Rubin, the NVFP4 and FP8 kernels achieved only about 42.1% and 44.4% of the roofline. The core issue: Rubin enables tensor cores to consume operands twice as fast, but the old kernel does not feed them fast enough.

What changed on Vera Rubin

Comparing vendor specifications, Vera Rubin NVL72 offers 35 PFLOPS per GPU for NVFP4 tensor cores versus 9 PFLOPS on HGX B200, 17.5 PFLOPS for FP8 versus 4.5 PFLOPS, and 4 PFLOPS for FP16/BF16 versus 2.25 PFLOPS. Memory bandwidth rises to 22 TB/s per GPU from 8 TB/s, SMs increase to 224 from 148, and peak power goes to 2300W from 1000W.

Several features matter for writing GEMMs. Tensor cores take twice the K: a tcgen05.mma step consumes 64 bytes along K on Vera Rubin instead of 32 bytes on Blackwell, while taking the same number of cycles. ThunderKittens expresses this with a new template parameter to the mma operation. Tensor memory grows to 576 columns from 512, adding 32 KiB, but the extra columns are reachable only through the .exclusive qualifier, a PTX 9.4 addition ensuring only one live tensor memory allocation per SM. Non-exclusive allocations stay capped at 512 columns and must be a power of two. Shared memory increases to 328 KiB from 228 KiB via an oversized shared memory mode set with a host-side attribute. A B side collector extends Blackwell's collector buffer to latch B tiles as well, with four labels: FILL reads from shared memory and latches, USE reads from the buffer, LASTUSE reads and releases, and DISCARD skips the latch. These are permission qualifiers, not guarantees. Early A release uses a new PTX 9.4 instruction, tcgen05.commit.sync_restrict::shared::read::mma::a, which signals the mbarrier as soon as an MMA finishes reading its A operand, letting the TMA loader start storing the next stage earlier.

Building the GEMM

The most intuitive bottleneck was the 32-byte K step. On Vera Rubin, that encoding's ISA ceiling is about 16.8 PFLOPS, and the NVFP4 Blackwell GEMM already achieved 14.7 PFLOPS out of the box, 88% of the ceiling. Flipping on the wider encoding alone barely improved performance. Double-K required two levers: moving fewer bytes and deepening the pipeline.

To fetch fewer bytes, the team stacked a second output tile on the same CTA pair along the M dimension. The two accumulators differ only in M, so they share the same B chunk. The original kernel used a 1x1 tiling format, meaning covering M512xN256 of output took two pair jobs that each transferred their own copy of B. A 2x1 format fetches B only once. This was not easily achievable on Blackwell because tensor memory was capped at 256 KiB, and two M256xN256 accumulators already take up 512 columns, leaving no space for block-scaled MMA scales. Vera Rubin's 64 extra columns store the scaling factors without extra work.

Deeper pipelines came from the larger shared memory, staging more tiles ahead of time. Sweeping ring depths for 16k square GEMMs, NVFP4 achieved 17,054 TFLOPS with 3 stages and 202 KiB, 20,595 TFLOPS with 4 stages and 258 KiB, and 22,239 TFLOPS with 5 stages and 314 KiB. FP8 E4M3 achieved 10,895 TFLOPS with 4 stages and 209 KiB, 11,995 TFLOPS with 5 stages and 257 KiB, and 11,288 TFLOPS with 6 stages and 305 KiB. These gains depended on the earlier optimizations.

Final tuning included kernel config tuning across workloads. The original 1x1 tile format was most performant for smaller square workloads, while larger shapes used 2x1 CTA pair tiling with cluster sizes of 2, 4, or 8 CTAs. Tile rasterization order was permuted to improve memory locality. The B side collector, using FILL on one MMA and LASTUSE on the next, reduced B reads to once instead of twice, measured at 1-3% improvement. Early A release gave 13.5% and 22.1% speedups for 64k and 128k square NVFP4 GEMMs respectively, useful at larger sizes where A tiles contend for residency and get evicted between reuses. At smaller sizes, A tiles never leave L2, so reads are already fast. Using early A release required decoupling A and B tiles from the same ring so A loads act independently.