Optimizing an NVFP4 Blockscaled GEMM on RTX PRO 6000 Blackwell GPU (SM120) - Colfax Research

Skip to content

2805 Bowers Ave, Santa Clara, CA 95051 | 408-730-2275<br>research@colfax-intl.com

Search

Optimizing an NVFP4 Blockscaled GEMM on RTX PRO 6000 Blackwell GPU (SM120)

This article is a continuation of our series on NVFP4 blockscaling on SM12x GPUs. In Part 1, we covered relevant PTX instructions, scale-factor layout details, and implementation details in CuTe DSL, including how to convert a CUTLASS dense GEMM example into an NVFP4 blockscaled GEMM. In this article, we optimize the NVFP4 GEMM from Part 1 for the NVIDIA RTX Pro 6000 Blackwell Server Edition GPU. We iteratively apply a series of optimizations, outlining the logic behind each one as well as the implementation steps.

We note at the outset that the version from the previous post is already fairly performant at mid-sized problem shapes (e.g., 8k square). Broadly speaking, the optimizations covered in this article fall into two categories:

Optimizations targeting well-known issues in small and large problem shape regimes — wave quantization and L2 cache thrashing, respectively.

Micro-optimizations, whose cumulative effect will be to lift compute throughput by a few percentage points overall.

At the end of our optimization ladder, we achieve compute throughput gains of 29% at 2k, 6% at 4k, 4% at 8k, 16% at 16k, and 40% at 32k, maxing out at 1666 TFLOP/s for 16k for a utilization rate of 83%.

All benchmarks were produced using Python version 3.13.13, PyTorch version 2.12.1, and nvidia-cutlass-dsl version 4.6.0.

Code for all the optimizations discussed in this article is included in the Colfax Research github repo at https://github.com/ColfaxResearch/cfx-article-src/tree/master/sm120_nvfp4_gemms.

RTX Pro 6000 Server Edition Specs

RTX PRO 6000 Blackwell streaming multiprocessor

Four processing partitions, each with a warp scheduler and dispatch unit,<br>register file, FP32 and INT32 execution lanes, fifth-generation Tensor Core,<br>load-store units, and a special-function unit. A shared 128 KB L1 data cache<br>and shared-memory block and four texture units appear below. The RT Core is omitted.

Figure 1. RTX Pro 6000 Streaming Multiprocessor (SM) diagram.

The RTX Pro 6000 has the following specifications:

96 GB of GDDR7 memory with ~1.6 TB/s of memory bandwidth

24,064 CUDA cores

188 Streaming Multiprocessors (SMs)

12 Graphics Processing Clusters (GPCs)

752 fifth-generation Tensor Cores (4 per SM)

L1 cache size: 128 KB/SM

L2 cache size: 128 MB

Peak FP4 Tensor TFLOP/s with FP32 Accumulate: 2015.2

Max SM Clock Rate: 2.43 Ghz

Version 1: The Baseline Kernel

We begin with a quick review of the structure of our kernel from the previous post. The kernel is warp-specialized with a producer-consumer pipeline in which each CTA consists of 1 TMA load warp and 8 MMA warps. The dedicated load warp issues TMA copies into SMEM for the A, B, SFA, and SFB operands. The eight MMA warps wait on those copies before performing SMEM-to-RMEM copies and issuing the appropriate warp-level mma.sync instruction. More specifically, the eight MMA warps form a tiled MMA with 4 warps along M, 2 along N, and 1 along K. The warp-level MMA atom has shape 16 x 8 x 64, so the MMA warps together span a 64 x 16 x 64 tile that is then repeated to cover the 128 x 128 x 128 CTA tile. Since the load warp requires fewer registers, register reallocation is performed with the load warp calling setmaxregister_decrease(40) and the MMA warps calling setmaxregister_increase(232).

After the mainloop, the MMA warps perform the epilogue, writing the output to SMEM, before warp 0 issues the TMA store from SMEM to GMEM. The kernel uses a static persistent tile scheduler in which a single CTA remains in residence on each SM and is repeatedly assigned work tiles.

We now proceed to evaluate the kernel’s performance. GEMM is a compute-bound problem, so we will evaluate the kernel in terms of measured TFLOP/s, both in absolute terms and as a percentage of the device maximum of 2015.2 TFLOP/s. As a rule of thumb, an optimized GEMM kernel at large problem shapes should achieve a utilization rate of 80% or more.

Figure 2 contains compute throughput numbers for Version 1 derived from the mean runtime of 20 iterations executed after 3 warmup iterations. For an 8k square GEMM, we see 1476 TFLOP/s, or about 73% utilization. Figure 2 also contains performance numbers for the NVFP4 GEMM kernels shipped with the two most recent versions of the cuBLAS library (13.5 and 13.6). For cuBLAS, version 13.5 exclusively dispatched to cutlass backend kernels, while 13.6 changed to nvjet for problem shapes 2k and 32k as well as switching to a different cutlass kernel for 16k.

Figure 2. Version 1 vs cuBLAS compute throughput.

Averaged across the five problem shapes, Version 1 achieves approximately 93% of cuBLAS 13.6’s performance, with a notable collapse at 32k. In addition, for the larger shapes we observe...