NVFP4 Grouped GEMM on Blackwell

A block-scaled grouped matrix multiply for the NVIDIA B200, written in the CUTLASS CuTe DSL. The kernel reaches a 6.37x geometric-mean speedup over the reference kernel. This repository holds the kernel versions, the measurements, and the analysis notes from a GPU MODE leaderboard problem.

Each group multiplies FP4 operands (float4_e2m1fn_x2) with FP8 scale factors (float8_e4m3fnuz) at a scale-factor vector size of 16, and writes FP16 output.

Result

All times are the median of a benchmark run on a B200, in microseconds, plus or minus the standard deviation. The last column is the geometric mean across the four cases. Lower is better.

Version	G=8 K=7168	G=8 K=2048	G=2 K=4096	G=2 K=1536	Geomean
Reference kernel	407 ± 2.7	404 ± 1.5	187 ± 1.7	170 ± 3.4	268.9
v11 warp-specialized	302 ± 1.4	327 ± 1.1	147 ± 1.1	142 ± 1.0	213.1
v19 persistent	98 ± 0.25	149 ± 0.1	43.8 ± 0.26	34.4 ± 0.21	68.5
v24 final	71.1 ± 0.07	88.3 ± 0.06	30.4 ± 0.03	16.6 ± 0.02	42.2

Speedup of v24 over the reference kernel: 6.37x, derived from the geometric means above. All listed versions pass 10 of 10 correctness tests. notes/experiment-log.md records every version, including the 11 that failed.

What produced the speedup

Three changes account for almost all of the gain. The rest of the versions measured an idea and rejected it.

1. Warp specialization (v11, 1.26x). The reference kernel loads and computes on the same warp. This version gives warp 0 the TMA loads and warp 1 the MMA instructions, so the two overlap.
2. Persistent tile loop (v19, 3.11x). The grid becomes 148 blocks, which is the B200 streaming multiprocessor count. Each block then walks several output tiles in a cutlass.range loop. This reuses the tensor map instead of rebuilding it per tile.
3. Conditional predication with a lookup table (v21b, 1.58x). A precomputed group lookup table replaces the per-tile group search. A fast path skips seven predicated CuTe operations when a tile does not cross a group boundary.

Stage-count tuning (v22 series) supplied a further 2.6 percent. Five stages beat both four and six.

What did not work

These attempts are in kernels/ and carry their measurements in the log.

Attempt	Outcome
CUTLASS C++ through GemmUniversal (v1 to v3)	The evaluation harness blocks cudaLaunchKernelExC.
A 128x256 tile (v5)	Tensor memory holds 512 columns at most. This is a hardware limit.
A 128-bit wide store to global memory (v10)	The MLIR compilation failed.
Six warps at 192 threads (v13)	Rows at index 128 and above returned wrong values. Tensor map setup stayed on warp 0.
The standard persistent tile scheduler (v15)	MLIR rejects tiled_copy_t2r as a loop-carried argument.

Layout

Path	Contents
kernels/	Every kernel version, v0 through v36. v24_best.py is the final one.
kernels/submission_final.py	The version sent to the leaderboard.
notes/experiment-log.md	The full measurement table, one row per version.
notes/first-principles.md	A roofline estimate and the target time it implies.
notes/data-layout.md	The FP4 and scale-factor memory layouts.
notes/reference-analysis.md	What the reference kernel does and where it loses time.
notes/official-kernel-analysis.md	A read of the CUTLASS grouped block-scaled example.
notes/time-breakdown.md	Where the remaining microseconds go.
other-problems/	Submissions for the three earlier problems in the same set.

Running the kernels

The kernels do not run from this repository alone. Each one imports task, utils, and reference, which belong to the GPU MODE problem harness. That harness carries its own license, so this repository does not redistribute it.

1. Get the harness from gpu-mode/reference-kernels.
2. Copy a kernel from kernels/ into the problem directory as submission.py.
3. Submit it with popcorn-cli, which needs B200 access.

The measurements above came from that path. Nothing here reproduces them on other hardware.

Attribution

The problem statement, the reference kernel, and the evaluation harness belong to GPU MODE, under the Researcher Reciprocity License. This repository excludes those files.

notes/official-kernel-analysis.md quotes the CUTLASS example grouped_blockscaled_gemm.py, which NVIDIA licenses under BSD 3-Clause. The file names its source.

The kernel versions and the notes are my own work.