GPU MODE NVFP4 Blackwell Challenge
The unfiltered worklog
I participated in GPU MODE’s NVIDIA Blackwell challenge series, here’s my story of working through three of them: nvfp4_gemm, nvfp4_dual_gemm, and nvfp4_group_gemm.
Challenge 1: nvfp4_gemm — Where It All Began
This was my very firs CuTeDSL kernel, and it humbled me quickly. What I thought would be a straightforward port turned into a crash course in MLIR debugging and the subtle constraints of CuTeDSL’s compilation pipeline.
I went through five attempts, each 🥹failing:
Attempt 1 — I started with a basic grid-stride loop, dynamically creating tensors inside a while loop. MLIR serialization error. The issue was that constructing make_layout((m, n, l)) with runtime values inside the loop was something the compiler simply couldn’t serialize.
Attempt 2 — I tried fixing it by changing accessing shape of C to a fixed (128, 128, 1) for the identity tensor. Same MLIR serialization error.
Attempt 3 — I moved on to a different approach, but hit a TMEM reallocation issue — “total allocated columns must be <= 512.” Turns out you can’t reallocate TMEM inside a loop. The hardware just doesn’t support it.
Attempt 4 — Removed the TMEM reallocation, kept a single allocation. Back to MLIR serialization errors. Dynamic tensor creation was still lurking somewhere in the loop body.
Attempt 5 — I tried multiple T2R copy attempts for the epilogue, only to get rank mismatch errors. The T2R copy is designed for a very specific TMEM layout, and I was feeding it something it didn’t expect.
😥 I never managed to submit this one within the competition timeline, but when I finally got it working and benchmarked it locally, I hit 33.210 μs which was like 38th place if I would have submitted.
Challenge 2: nvfp4_dual_gemm — SwiGLU in PTX
Final result: 95th place — 63.210 μs
The dual GEMM challenge adds a second matrix multiplication and fuses it with a SwiGLU activation. The reference CuTeDSL kernel came in at:
⚡ 71.6 μs 🐌 71.9 μs | ⚡ 71.1 μs 🐌 274 μs | ⚡ 60.1 μs 🐌 60.2 μs | ⚡ 76.4 μs 🐌 76.4 μs
My starting approach was straightforward: take the kernel I’d been building for nvfp4_gemm, pass in the new B tensors, make a second scaled_mm call, and add the SwiGLU epilogue:
temp1 = accumulator1 * (1 / (1 + tl.exp(-accumulator1)))
accumulator2 = temp1 * accumulator2This particular part was tricky to get the TMEM column offsets right when you have dual accumulators working simultaneously the accumulators share TMEM space and you need to be careful about where each one lives.
I didn’t had time to proceed further so, I submitted this but later discovered some resource on Inline PTX by the discussions in discord and also from Arseni ivanov's Blog. So I tried to write the SwiGLU epilogue in PTX by myself:
// SiLU: x * sigmoid(x) = x * (1 / (1 + exp(-x)))
mov.f32 %t0, %g0; // copy x
neg.f32 %t1, %t0; // -x
mul.f32 %t2, %t1, 1.44269504; // -x * log2(e)
ex2.approx.ftz.f32 %t3, %t2; // exp(-x) = 2^(-x * log2(e))
mov.f32 %t4, 1.0; // load 1.0
add.f32 %t5, %t4, %t3; // 1 + exp(-x)
rcp.approx.ftz.f32 %t6, %t5; // 1 / (1 + exp(-x))
mul.f32 %t7, %g0, %t6; // x * sigmoid(x)
mul.f32 %r0, %t7, %v0; // result * v0This PTX rewrite gave me a ~20% speed bump on the epilogue, which was satisfying at last. but couldn’t submit in the timeline still a very good lesson learnt. Lezz GOO!
Challenge 3: nvfp4_group_gemm — The Big One
Final result: 42nd place - 38.460 μs
This is where I locked in.
Grouped GEMM is one of the most practically relevant kernels in modern ML infrastructure. You have multiple matrix multiplications of different sizes that need to execute efficiently on the same GPU. The canonical use case is Mixture-of-Experts models, where you have a dynamic number of tokens routed to each expert, potentially with different MMA-shapes per group. I didn’t document the metrics but documented the techniques used
The reference kernel was very basic a single TMA-MMA warp with lots of hard-coded copy ops:
⚡ 334 μs 🐌 391 μs | ⚡ 334 μs 🐌 368 μs | ⚡ 148 μs 🐌 203 μs | ⚡ 132 μs 🐌 162 μs.
Rather than write everything from scratch, I turned to the grouped GEMM implementation in the CUTLASS GitHub repository and started adapting it.
The competition problem requires you to materialize tensors inside the kernel call, while the CUTLASS reference passes pointers all the way until the kernel launch. The result matrix also doesn’t have its own tensor map in the reference which a structural mismatch that required rethinking how data flows through the epilogue.
I needed tensor maps for the scale factors (SFA and SFB), and the whole pipeline had to use 2-SM cooperative MMA, where two CTAs share TMEM and collaborate on the same tile. Getting this wired up correctly through CuTe’s layout algebra was the first real challenge.
The 1-CTA Approach
My first working version used a single CTA (Cooperative Thread Array) per tile. A CTA is CUDA's name for a threadblock where a group of threads that can synchronize and share memory with each other. It's the basic unit of work the GPU scheduler assigns to an SM. Simple: each CTA handles one complete GEMM tile, and the hardware scheduler distributes CTAs across SMs. I focused on getting TMA loads pipelined correctly and the MMA accumulator → FP16 epilogue fused.
Going 2-CTA Cooperative
Next, I tried 2-CTA cooperative kernels, where two CTAs collaborate on a single GEMM via shared TMEM. The idea: for larger tiles, split the work and have two CTAs each handle half, synchronizing through the cluster barrier mechanism.
The speedup was... marginal. A few microseconds here and there. The overhead of inter-CTA synchronization ate into the gains, especially for the smaller groups in the batch.
The TMEM Experiment
Inside each CTA, I used warp specialization dedicating different warps to different jobs rather than having all threads do the same thing. The structure was one TMA warp that handles all global memory loads through TMA descriptors, one MMA warp drives the tensor core math, and four epilogue warps convert the FP32 accumulator to FP16 and write the result back. The TMA warp and MMA warp communicate through an async pipeline with multiple stages, so loads for the next K-tile overlap with compute on the current one. The MMA warp signals the epilogue warps through a separate accumulator pipeline once the full K-reduction is done.
Blackwell’s tensor cores operate on TMEM (Tensor Memory) a dedicated 256KB-per-SM scratchpad that only the tensor core can access directly. The epilogue warps need to explicitly copy the accumulator out of TMEM, convert it, stage through shared memory, and then TMA-store to global memory. It’s a lot of moving parts, but each warp only worries about its own stage.
It worked, and it was reasonably fast. But with smaller M values in some groups (as low as 40), I was leaving SM occupancy on the table. A group with M=40 and a 128-wide M-tile means exactly one tile along M not much parallelism to exploit.
The Persistent Kernel Detour
Persistent kernels are the hot technique for Blackwell that launches exactly as many threadblocks as you have SMs and have them loop over tiles. I tried this with CLC (Cluster Launch Control) scheduling, but it’s slower than the static schedule. For this problem with only 2–8 groups and relatively small total tile counts, the overhead of the persistent scheduling machinery wasn’t worth it. So I rolled back to the older one.
In this challenge, I think maybe some of my implementation or the exploration of data flow might be wrong, but still it laid pretty good ground work for me.
Built with CuTeDSL, CUTLASS, PTX, and a lot of MLIR error messages. Thanks to GPU MODE community for actively pushing and helping me push my learning deeper. Let me cook something with TTS soon.