AI 基礎設施
ThunderKittens Adds Support for New Vera Rubin Instructions, NVFP4 GEMM Exceeds 22 PFLOPS on Sample GPU
Together AI ported its open-source CUDA kernel framework to NVIDIA Vera Rubin, reorganizing data reuse and the memory pipeline. Tests with square 16K matrices reached 22.24 PFLOPS, though the figure comes from a qualification sample and the still-maturing CUDA 13.4.
Together AI announced Vera Rubin support for ThunderKittens on September 10, adding Rubin-specific PTX instructions and NVFP4 and FP8 GEMM support. Although the existing Blackwell kernels could run without modification, they reached only 42.1% and 44.4% of the theoretical roofline for NVFP4 and FP8, respectively. The problem was not insufficient Tensor Core compute, but that the existing data path could not keep pace with the faster compute units.
The team first increased the K-step size of `tcgen05.mma` from 32 bytes to 64 bytes, allowing each instruction to process twice as much data in the same number of cycles. Rubin’s Tensor Memory also grew from 512 to 576 columns. Combined with the `.exclusive` configuration, it can hold two sets of accumulators as well as quantization scales. This allows the kernel to use a 2×1 tile so that two output blocks can reuse the same B data. A new B-side collector can temporarily retain the B tile in the MMA staging buffer, reducing redundant reads; according to the team’s measurements, this change alone delivered an improvement of about 1% to 3%.
Another key feature is the oversized shared-memory mode, which increases the available capacity per SM from 228 KiB to 328 KiB. This lets the kernel preload more tiles and deepen its ring pipeline. For a square 16K NVFP4 GEMM, increasing the pipeline from three stages to five raised performance from 17,054 to 22,239 TFLOPS. The five-stage FP8 implementation reached 11,995 TFLOPS. ThunderKittens has merged Rubin support into its public repository, providing a readable example for kernel development that does not depend on closed-source libraries.
However, these results do not represent end-to-end model throughput and cannot be translated directly into inference costs. They cover only specific GEMM shapes and were obtained using a Qualification Sample GPU and CUDA 13.4. Power consumption, inter-GPU communication, smaller or irregular matrices, and cuBLAS improvements on production hardware could all change the relative performance. Engineering teams should next watch for reproducible results using the production CUDA toolchain, a broader range of kernel types, and real-world MoE and inference workloads.