The Critical Search: Uncovering SM120 CUTLASS Collective Implementations

In the midst of a deep optimization session for GLM-5-NVFP4 inference on 8× NVIDIA RTX PRO 6000 Blackwell GPUs, the assistant issued a single, seemingly mundane bash command. Yet this command — message [msg 915] — represents a pivotal investigative step in a methodical journey to understand why certain CUTLASS tile configurations fail on SM120 architecture and, ultimately, how to squeeze more throughput from the FP4 GEMM kernels powering the model's MoE layers.

The Message

[assistant] [bash] ssh root@10.1.230.174 'find /root/ml-env/lib/python3.12/site-packages -path "*/cutlass/include/cutlass/gemm/collective/*" -name "*sm120*" -o -path "*/cutlass/include/cutlass/gemm/collective/*" -name "*array*block*" | head -10'
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_mma_array_tma_blockwise_scaling.hpp
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_mma_tma_blockwise_scaling.hpp
/root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp

Context: The Optimization Journey

To understand why this search matters, we must trace the arc of the investigation. The session had already established that the GLM-5-NVFP4 model is compute-bound, not communication-bound — a TP4+PP2 configuration proved 2× slower than TP8, ruling out allreduce latency as the primary bottleneck. The real culprit was FP4 GEMM kernel efficiency on SM120 (Blackwell architecture).

The numbers painted a stark picture. Each GPU draws only ~235W out of its 600W TDP during inference — a mere 39% utilization. The CUTLASS kernels achieve ~1,300 TFLOPS (70% of dense peak) only for very large matrices, but during actual decode, per-expert batch sizes of ~16–64 tokens achieve a paltry 0.8–55 TFLOPS (0.02–3% of peak). The 99KB shared memory limit on SM120 prevents using larger CUTLASS tile configurations — specifically, M128×N256 and M256×N128 fail to initialize with an "Error Internal" from gemm.initialize().

The user's directive in [msg 905] was concise: "Fix tiles, then try 2/3." The assistant had been working through this agenda systematically. Messages [msg 906] through [msg 914] traced the investigation through the CUTLASS initialization code, shared memory measurements, and SMEM requirement calculations. By [msg 913], the assistant had confirmed that the failing tiles require more than 99KB of shared memory when using the finalize epilogue — 155KB for 128×256×128 with a single pipeline stage.

Message [msg 914] began probing the CUTLASS source itself, looking at StageCountAutoCarveout behavior and the collective builder implementations. The assistant started searching for SM120-specific collective headers with a command that found general SM120 files. Message [msg 915] refines this search with a more targeted query.

Why This Search Matters

The command in [msg 915] is not a random file hunt. It is a deliberate, knowledge-driven probe designed to answer a specific question: What CUTLASS collective implementations exist for SM120, and do any of them handle blockwise scaling differently?

The find command uses two parallel path patterns. The first, -name "*sm120*", targets any file with SM120 in its name — the Blackwell architecture identifier. The second, -name "*array*block*", searches for array-based block scaling implementations, which may offer different memory layouts or scheduling strategies that could circumvent the shared memory limitation. The -o (OR) operator between them means the command returns files matching either pattern, and the head -10 limits output to the first 10 results.

The results reveal three key files:

  1. sm120_mma_array_tma_blockwise_scaling.hpp — an SM120-specific array-based collective using TMA (Tensor Memory Accelerator) with blockwise scaling, appearing twice (possibly due to symlinks or multiple copies in the build tree)
  2. sm120_mma_tma_blockwise_scaling.hpp — the non-array variant of the SM120 blockwise scaling collective
  3. sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp — an SM90 (Hopper) array-based collective for FP8 blockwise scaling, included as a reference or fallback

Assumptions and Reasoning

The assistant makes several assumptions in this search. First, that the CUTLASS source bundled within flashinfer's data directory contains the relevant SM120 implementations — a reasonable assumption given that flashinfer ships with a vendored CUTLASS distribution. Second, that the file naming convention accurately reflects architectural targeting (sm120 = Blackwell, sm90 = Hopper). Third, that examining these header files will reveal why the larger tiles fail and how to work around the shared memory constraint.

A notable assumption is that an "array" variant might offer different behavior. In CUTLASS terminology, "array" collectives typically use a different warp specialization strategy or data movement pattern compared to the standard TMA collectives. The assistant is implicitly hypothesizing that the array-based variant might handle shared memory allocation differently, potentially fitting the larger tiles within the 99KB limit.

Input Knowledge Required

To fully grasp this message, one needs:

Output Knowledge Created

The search yields critical intelligence. The presence of two SM120-specific collective implementations — an array variant and a standard TMA variant — confirms that CUTLASS has been extended for Blackwell's block-scaled FP4 capabilities. The array variant (sm120_mma_array_tma_blockwise_scaling.hpp) is particularly interesting because it may use a different warp arrangement or shared memory partitioning strategy.

The discovery of the SM90 FP8 array collective (sm90_mma_array_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp) provides a potential reference implementation. If the SM120 array variant inherits from or parallels the SM90 design, understanding the Hopper implementation could illuminate the Blackwell one.

Most importantly, the search establishes the file-level targets for the next phase of investigation. The assistant now knows exactly which source files to examine to understand why M128×N256 and M256×N128 tiles fail, and whether the array-based collective offers a viable path forward. This transforms the investigation from black-box debugging (observing failures at the API level) to white-box analysis (understanding the kernel generation logic).

The Thinking Process Visible in the Reasoning

The assistant's thinking, visible across messages [msg 906] through [msg 915], follows a classic diagnostic pattern:

  1. Observe symptom: Larger CUTLASS tiles fail with "Error Internal" during initializeMoeGroupedGemm
  2. Measure constraints: Query GPU properties to find 99KB opt-in shared memory limit
  3. Estimate requirements: Calculate SMEM needs for each tile configuration, confirming the failing tiles exceed the limit
  4. Trace the code: Examine the CUTLASS initialization path to understand where and why the failure occurs
  5. Search for alternatives: Look for different collective implementations that might handle memory differently Message [msg 915] is step 5 in this progression. The refinement from the previous search (which found general SM120 files) to this more targeted query (adding the array/block pattern) shows the assistant adapting its search strategy based on partial results. The -o operator and dual-path search pattern reveal a sophisticated understanding of the codebase organization — the assistant knows that relevant implementations might be named either by architecture (sm120) or by algorithmic feature (array, block).

Broader Significance

This message, while small, exemplifies the kind of systematic investigation that characterizes effective performance optimization work. Rather than randomly tweaking parameters or trying brute-force configuration sweeps, the assistant is building a mental model of the system's constraints and capabilities. The shared memory limit is a physical constraint (99KB per block), but the software implementation (CUTLASS collective builder) determines how that constraint maps to tile configurations. By examining the source code, the assistant hopes to find either a configuration workaround or an alternative collective that better utilizes the available shared memory.

The search also reveals the layered nature of the optimization problem. The assistant is not just tuning server parameters (like --max-running-requests or --num-continuous-decode-steps) but diving into the kernel generation layer. This is deep optimization work — understanding how the fundamental computational primitives interact with the hardware's physical constraints.

Conclusion

Message [msg 915] is a small but critical pivot point in a larger optimization narrative. It transitions the investigation from measurement and estimation to source code analysis. The discovered files — sm120_mma_array_tma_blockwise_scaling.hpp and sm120_mma_tma_blockwise_scaling.hpp — become the targets for the next phase of work: understanding why the standard collective fails for larger tiles and whether the array variant offers a path forward. In the broader context of segment 7, this search contributes to the deep analysis of FP4 GEMM kernel efficiency on SM120, which ultimately leads to a 28% throughput improvement through server parameter tuning and a ranked optimization plan documented across multiple glb5improvement-xx.md files.