Together AI 的 ThunderKittens 现已支持 NVIDIA Vera Rubin NVL72

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.

As Rubin preserves the Blackwell programming model, our old GEMMs still work. However, naively running them on Vera Rubin, we observe that our NVFP4 and FP8 kernels only achieve around 42.1% and 44.4% of the roofline — plenty of room for optimization!

The rest of this post is in two parts. First, we cover the new Rubin features that matter for a GEMM and how to use them in ThunderKittens. Then, we slowly integrate these features into our existing Blackwell NVFP4 kernel, taking it to over 22 PFLOPS and competitive with cuBLAS and CuTE DSL.

The core issue we find is that while Rubin enables tensor cores to consume operands twice as fast, our old Blackwell kernel does not feed them fast enough. To reach the compute ceiling, we need tiles to squeeze more reuse out of the data they already have on-chip.

Comparing vendor specifications, we see the following improvements from Blackwell to Vera Rubin.

As it relates to writing performant GEMMs, we specifically take note of the subsequent features.

For review, a tcgen05.mma computes C = A@B + C over a MxNxK tile, consuming a fixed number of bytes along K per step. On Blackwell, this step is 32 bytes, but on Vera Rubin it can be increased to 64 bytes. The MMA itself still takes the same number of cycles, so doubled-K allows us to pack twice as much work into the same instruction window.

In ThunderKittens, we express this through a new template parameter to our existing mma operation.

mma_ABt (...); // Blackwell Default: 32-byte K step mma_ABt (...); // Vera Rubin: 64-byte K step 2. Tensor memory grows to 576 Columns Blackwell introduced the concept of tensor memory, a 128 lane x 512 column x 32-bit space which tensor cores could directly read and write to. On Vera Rubin, this space increases to 576 columns, providing an extra 32 KiB of tensor memory to play with.

Note that these extra columns are reachable only through the .exclusive qualifier, a PTX 9.4 addition that ensures there’s only one live tensor memory allocation on an SM. Non-exclusive allocations are still capped at 512 and must be a power of two.

In ThunderKittens users may request this with a template parameter to our tensor memory allocator, specifying that the allocation is exclusive.

template struct tensor_allocator { .... }

tensor_allocator tm; // Blackwell default: 512 columns tensor_allocator tm; // Rubin: Up to 576 columns 3. Shared memory increases to 328 KiB While Hopper and Blackwell provided 228 KiB of shared memory, Vera Rubin introduces an oversized shared memory mode that can be dynamically increased to 328 KiB. This is a host-side specification that can be called as below.

CUfunction function = nullptr; cudaGetFuncBySymbol(&function,reinterpret_cast (kernel)); cuFuncSetAttribute(function, CU_FUNC_ATTRIBUTE_SHARED_MEMORY_MODE, CU_SHARED_MEMORY_MODE_ALLOW_OVERSIZED_SHARED_MEMORY); 4. B Side Collector Blackwell introduced the concept of a collector buffer, a small MMA staging buffer that could latch onto an A tile so the next instruction takes it from there instead of from shared memory. Vera Rubin extends this functionality to the B tile as well with .collector::b::* .

To utilize this, we annotate each MMA’s operands with one of four labels that describe its action on the collector buffer.

These labels are permission qualifiers for reuse, not guarantees. This means that the Tensor Core could still reload a matrix even if it has permission to reuse.

Now that either operand can be resident in the collector buffer, we can experiment with new patterns. Over a 2x2 block, for instance, collecting on both sides takes four MMAs from eight operand reads to only five.

In ThunderKittens, we can expose this through the following.

mma2_ABt_chunk (C[0][0], a0, b0, ...); mma2_ABt_chunk (C[0][1], a0, b1, ...); mma2_ABt_chunk (C[1][1], a1, b1, ...); mma2_ABt_chunk (C[1][0], a1, b0, ...); 5. Early A release A tcgen05.commit arrives on an mbarrier once the MMA it processes is finished, signaling the producer that a stage slot is free for reuse. PTX 9.4 introduces tcgen05.commit.sync_restrict::shared::read::mma::a , a new instruction that allows us to move the signal earlier for A tiles. Rather than waiting for the MMA to finish, we can fire the barrier as soon as an MMA has finished reading its A operand out of shared memory, enabling us to signal the tma loader to start storing the next stage’s memory.

ThunderKittens introduces a new commit type for a user to express this.

tensor_commit (inputs_finished[stage], mask); // arrives when MMA retires tensor_aread_commit (A_finished[slot], mask); // arrives when MMA finishes reading A Building the GEMM We now have new features and a Blackwell GEMM. The following sections slowly integrate these features into our existing Blackwell kernel, and explain why they’re needed as we move to Vera Rubin.

The most intuitive bottleneck comes from still relying on Blackwell’s 32-byte K step. On Vera Rubin, that encoding's ISA ceiling is ~16.8 PFLOPs, and our NVFP4 Blackwell GEMM already achieves 14.7 PFLOPs out of the box (88% of the ceiling). To push further, we must double the K bytes processed by our MMAs.

However, simply flipping on the wider encoding for the existing Blackwell kernel, we notice that performance only marginally improves, instead of a supposed 2x. While the wider MMA doubles how fast the tensor cores can consume operands, it does not help with how fast we can supply them. For double-K to be useful, we pull two levers to keep our cores happy: moving fewer bytes and deepening the pipeline to ensure these fetches are overlapped.

To fetch fewer bytes, we stack a second output tile on the same CTA pair along the M dimension. As the two accumulators differ only in M, we are able to share the same B chunk for them both. Our original Blackwell NVFP4 kernel utilized a 1x1 tiling format, meaning that covering M512xN256 of output took two pair jobs that each independently transferred their own copy of B. By moving towards a 2x1 format, we can fetch B only once and cover the same output, reducing our operand traffic.

This 2x1 tiling format was not easily achievable in our NVFP4 Blackwell kernel as tensor memory was capped at 256 KiB. Two M256xN256 accumulators already take up 512 columns, meaning that block-scaled MMAs do not have space to store their A and B scales. While programmers could work around this constraint by having the epilogue warps load only a subset of the accumulator columns before signaling the MMA-empty state, allowing the MMA for the next K tile to begin, this introduced a small amount of latency that could not be hidden. Luckily, with Vera Rubin’s 64 extra columns we can store the scaling factors without needing to do this dance.

The changed tiling format reduces operand traffic, but it does not help decrease how long each fetch takes. The next challenge is to keep the tensor cores continuously fed. Vera Rubin’s larger shared memory allows us to create deeper pipelines, staging more tiles ahead of time and giving transfers more time to complete. Sweeping ring depths for our NVFP4 and FP8, 16k square GEMMs, we see the following.

NVFP4 16k Square GEMM:

FP8 (E4M3) 16k Square GEMM:

While the most gain appears from this final peg, we note that these gains are dependent on our earlier optimizations. Below are results from sweeping the K step size, shared memory pipeline, and tiling strategy independently.

To push our kernels even further, we play around with a few final knobs

We note that all measurements above were produced using NVIDIA CUDA 13.4 on a Qualification Sample (QS) GPU. We expect performance of all baselines to continue to improve with Vera Rubin software releases.

We hope you found some of this useful, and we are excited for everyone to start playing with them soon. From LUT GEMMs, hardware native megakernels, and new engine optimizations, there’s still plenty of interesting snippets to share. More on that soon!

The kernels and performance teams at Together AI are actively hiring! If you’d like to learn more about these kernels or work with us on developing the next set of updates, please reach out to Simran or Dan!