To Infinity and Beyond: ThunderKittens Now on NVIDIA Vera Rubin NVL72!

Dylan Lim, Xinyi Li, Sonny Li, Peter Wu, Dan Fu, Simran Arora

40+ Models Chosen for Production...40+ Models Chosen for Production...40+ Models Chosen for Production...

The kernels team at Together recently received access to the NVIDIA Vera Rubin NVL72 platform. We spent the past few days digging through the new ISA and poking the chip with micros. There are many fun new features! We finished adding some functionality into ThunderKittens to write NVFP4 and FP8 GEMMs on Vera Rubin, as well as help fellow kittens explore the stars.

Before diving into what Vera Rubin delivers, we start our journey with a quick refresher of a Blackwell GPU GEMM.

NVIDIA Blackwell architecture’s fifth-generation tensor cores fundamentally changed the GEMM programming model. While NVIDIA Hopper architecture’s wgmma instruction was issued collectively by a warpgroup, Blackwell’s tcgen05 instruction is issued by a single thread, enabling one small producer warp to drive the tensor cores. The accumulator also 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.

To reach competitive performance on Blackwell, our GEMM:

Through these efforts, we realized the following results.

Read more about these kernels and their optimizations in our earlier Together blog post or the ThunderKittens 2.0 release!