Peering Into the CUTLASS Kernel: Tracing Shared Memory Constraints on Blackwell GPUs

In the midst of an intensive optimization campaign for GLM-5-NVFP4 inference on eight NVIDIA RTX PRO 6000 Blackwell GPUs, the assistant issued a deceptively simple command: a grep into a CUTLASS header file to find shared-memory-related symbols. Message [msg 916] captures this moment — a single bash invocation searching for SharedStorage, smem_size, SmemKB, kSmem, MaxShared, and SMEM inside sm120_blockscaled_mma_array_tma.hpp. The output reveals four lines, all referencing the SharedStorage struct and its type aliases. Though brief, this message sits at a critical juncture in a deep investigation into why certain FP4 GEMM tile configurations crash during initialization on the SM120 architecture.

The Context: A Compute-Bound Model Starved for Kernel Efficiency

To understand why this grep matters, one must appreciate the journey that led to it. The assistant had spent dozens of messages diagnosing why the GLM-5-NVFP4 mixture-of-experts model — deployed with SGLang on eight Blackwell GPUs — was underperforming its theoretical potential. Earlier benchmarking had confirmed the model was compute-bound rather than communication-bound: a TP4+PP2 configuration was 2× slower than TP8, ruling out allreduce latency as the primary bottleneck ([msg 899]). The real culprit was poor FP4 GEMM kernel efficiency on SM120.

The numbers were stark. Each GPU drew only ~235W out of its 600W TDP during inference ([msg 903]). The CUTLASS kernels plateaued at roughly 1,300 TFLOPS — about 70% of the dense peak of ~1,850 TFLOPS — but only for very large matrices. During actual decode, per-expert batch sizes of roughly 16–64 tokens achieved a paltry 0.8–55 TFLOPS, or 0.02–3% of peak ([chunk 7.0]). Something was fundamentally limiting the kernel's ability to use the GPU's compute capacity.

The investigation zeroed in on shared memory. Blackwell's SM120 architecture provides 100 KB of shared memory per SM, with a maximum opt-in per-block limit of 99 KB (101,376 bytes). The default per-block limit is just 48 KB ([msg 912]). CUTLASS's warp-specialized FP4 block-scaled GEMM kernels use shared memory extensively — for staging input tiles (FP4 activations and weights with FP8 scale factors), for the mainloop pipeline, and for the epilogue where output tiles are accumulated. Larger tile configurations like M128×N256 and M256×N128 were failing with kErrorInternal during gemm.initialize(), even though can_implement() passed ([msg 910]). The assistant hypothesized that the kernel's shared memory allocation exceeded the 99 KB limit at runtime, causing the initialization to fail.

The Message: A Targeted Probe Into CUTLASS Source

Message [msg 916] is the assistant's attempt to verify this hypothesis by examining the actual CUTLASS collective header. The command is:

grep -n "SharedStorage\|smem_size\|SmemKB\|kSmem\|MaxShared\|SMEM" \
  /root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp \
  | head -20

The path reveals the specific file: sm120_blockscaled_mma_array_tma.hpp — the CUTLASS collective implementation for SM120 block-scaled matrix multiplication using Tensor Memory Accelerator (TMA) with array-based interface. This is the kernel path used for FP4 block-scaled GEMMs, where both operands are in FP4 format with per-block scaling factors in FP8. The "array" variant likely refers to how the warp-specialized threads are organized, while "TMA" indicates it uses CUDA's Tensor Memory Accelerator hardware for efficient data movement between global memory and shared memory.

The grep patterns are carefully chosen. SharedStorage is the CUTLASS struct that holds all shared memory allocations for a kernel invocation — its sizeof directly determines the kernel's SMEM footprint. smem_size and SmemKB are common variable names for computed shared memory sizes. kSmem is a typical constant naming convention. MaxShared would match MaxSharedMemory constants. SMEM is a broad catch-all. The assistant is casting a wide net to find any SMEM-related declarations in this file.

The output is revealing in its sparseness. Only four lines match, all for SharedStorage:

258:  struct SharedStorage {
273:    using PipelineStorage = typename MainloopPipeline::SharedStorage;
277:  using TensorStorage = typename SharedStorage::TensorStorage;
278:  using PipelineStorage = typename SharedStorage::PipelineStorage;
279:  using TensorMapStorage = typename SharedStorage::TensorMapStorage;

None of the other patterns — smem_size, SmemKB, kSmem, MaxShared, SMEM — appear in this file. This is itself informative: the file does not explicitly compute or store SMEM sizes as named constants or variables. The shared memory layout is entirely defined through the SharedStorage struct's member types, which are themselves aliased from other components (the mainloop pipeline, tensor storage, and tensor map storage). The actual byte count is determined at compile time via sizeof(SharedStorage), as the assistant would discover in the very next message ([msg 918]), where it finds the SharedStorageSize static assert.## The Reasoning Behind the Grep

Why did the assistant reach for a source-code grep rather than, say, running a CUDA occupancy calculator or instrumenting the kernel at runtime? The answer lies in the nature of the failure mode. The error — kErrorInternal during initialize() — occurs at kernel launch time, inside CUTLASS's own initialization logic. This is a compile-time or launch-time validation failure, not a runtime crash. The kernel binary is already compiled; the issue is that when CUTLASS attempts to set up the kernel's shared memory configuration and launch parameters, the requested SMEM exceeds what the device can provide.

The assistant had already done the math in [msg 913], writing a Python script to estimate SMEM requirements for various tile configurations. That analysis showed that the 128×256×128 tile with a single pipeline stage required approximately 91 KB for the mainloop plus default epilogue, or 155 KB with the finalize epilogue — both dangerously close to or exceeding the 99 KB opt-in limit. But those were estimates based on a simplified model of CUTLASS's memory layout. The grep was an attempt to validate those estimates against the actual source code.

The assistant made a specific assumption here: that the SMEM allocation is primarily determined by the SharedStorage struct in the collective header, and that the struct's size would be computed from its member types. This assumption is correct for CUTLASS, which uses a template metaprogramming approach where shared memory layouts are composed from smaller building blocks (pipeline storage, tensor storage, tensor map storage) at compile time. The SharedStorage struct's sizeof is indeed the authoritative SMEM footprint.

However, the assistant also implicitly assumed that the grep patterns would find explicit SMEM size calculations or constants in this particular file. The sparse output — only SharedStorage matches — disproves that assumption. The SMEM size is not computed in this file; it emerges from the composition of the member types, which are themselves defined in other headers (the mainloop pipeline, the tensor storage layouts). This is a common pattern in CUTLASS: the collective defines the structure, but the actual byte counts are determined by the template instantiations of its members.

The Input Knowledge Required

To fully understand this message, one needs substantial background knowledge spanning multiple domains:

CUTLASS architecture: Understanding that CUTLASS uses a hierarchical template structure where GemmKernel composes a CollectiveMainloop (which handles the main GEMM computation) and a CollectiveEpilogue (which handles output conversion and writeback). The SharedStorage struct is typically defined in the collective mainloop and contains storage for the pipeline, tensor slices, and tensor map descriptors.
SM120 architecture specifics: Blackwell GPUs have 100 KB of shared memory per SM, with a 99 KB opt-in per-block limit. The SM120 also introduces new TMA hardware and warp-specialized scheduling that CUTLASS exploits. The sm120_blockscaled_mma_array_tma.hpp file is specifically for SM120's block-scaled MMA (matrix multiply-accumulate) with TMA and array-based warp specialization.
FP4 block-scaled format: The model uses FP4 weights with per-block FP8 scale factors. This means each GEMM operand has two components: the packed FP4 data (0.5 bytes per element) and the scale factors (1 byte per 16 elements). The shared memory budget must accommodate both.
The broader optimization context: This grep is one step in a multi-pronged investigation that includes benchmarking at various concurrency levels, measuring GPU power draw, testing different tensor parallelism configurations, and exploring alternative kernel paths like cuBLASLt.

The Output Knowledge Created

The immediate output of this message is modest: four lines of code showing the SharedStorage struct definition and its type aliases. But the knowledge created extends beyond these lines:

Confirmation that SMEM is defined structurally, not numerically: The absence of explicit SMEM size constants tells the assistant that the size is determined by template composition. This guides the next step: finding where sizeof(SharedStorage) is computed and asserted.
Identification of the specific file and line numbers: Line 258 for the struct definition, lines 273-279 for the type aliases. These become targets for deeper inspection.
A narrowing of the search space: Since the collective header doesn't compute SMEM sizes directly, the assistant must look at the kernel-level header (which wraps the collective) or the builder logic. This leads directly to [msg 918], where the assistant finds SharedStorageSize = sizeof(SharedStorage) and the critical static_assert against sm120_smem_capacity_bytes. The grep also implicitly creates negative knowledge: the assistant learns that smem_size, SmemKB, kSmem, MaxShared, and SMEM are not used as identifiers in this file, ruling out certain code patterns and narrowing where the SMEM budget is actually enforced.

Assumptions and Potential Mistakes

The assistant operated under several assumptions in this message:

That the failing tiles fail due to SMEM overflow: This was the leading hypothesis, strongly supported by the SMEM estimation calculations in [msg 913]. The 128×256×128 and 256×128×128 tiles were estimated to require 91 KB (default epilogue) or 155 KB (finalize), both at or above the 99 KB limit. This assumption proved correct — the next message ([msg 918]) reveals the static_assert that enforces SharedStorageSize <= sm120_smem_capacity_bytes.
That the SMEM information would be findable via grep in the collective header: This was partially correct — the SharedStorage struct is indeed in this file — but the actual size enforcement is in the kernel header (sm120_blockscaled_mma_array_tma.hpp vs the kernel wrapper). The assistant needed to search more broadly.
That the grep patterns would be sufficient to find all relevant SMEM code: The patterns missed SharedStorageSize (which would be found in the kernel header) and sm120_smem_capacity_bytes (defined in cutlass/arch/arch.h). The assistant compensated by following up with additional searches in subsequent messages.
That the head -20 limit was sufficient: The file might have contained relevant SMEM code beyond line 20 of matches. However, given the sparse output (only 4 matches total), this was a safe assumption. A subtle mistake in the grep itself: the pattern SmemKB uses a capital K, but CUTLASS code typically uses lowercase in variable names (e.g., smem_kb). Similarly, kSmem might match C++ constant naming conventions but CUTLASS doesn't consistently use that style. These are minor pattern-matching issues that didn't affect the outcome since the broader patterns (SharedStorage, SMEM) caught the relevant code.

The Broader Significance

This message exemplifies the kind of deep systems investigation that characterizes high-performance ML engineering. The assistant is not content to observe that tiles fail — it traces the failure through the CUTLASS source code, verifying hypotheses against the actual implementation. The grep is a reconnaissance tool, mapping the terrain before committing to a more detailed analysis.

The investigation that follows this message is productive. In [msg 918], the assistant finds the SharedStorageSize static assert and the sm120_smem_capacity_bytes constant (101,376 bytes = 99 KB). In [msg 919], it confirms the constant definition in cutlass/arch/arch.h. This chain of discovery — from runtime error, to SMEM estimation, to source code grep, to the actual capacity constant — demonstrates a methodical approach to debugging kernel initialization failures.

The practical outcome is significant: the assistant confirms that the larger tile configurations are fundamentally impossible on SM120 without reducing the epilogue footprint or restructuring the shared memory layout. This leads to a strategic shift — rather than fixing the tiles, the assistant explores alternative optimization paths including cuBLASLt FP4, piecewise CUDA graphs, expert parallelism, and continuous decode steps. The 28% throughput improvement achieved later in the segment ([chunk 7.0]) comes from these alternative approaches, not from fixing the tiles.

In the end, message [msg 916] is a small but crucial step in a larger journey — a moment of focused investigation that produces negative knowledge ("these tiles cannot work") and redirects effort toward more fruitful optimizations. It is a reminder that in systems engineering, understanding why something cannot work is often as valuable as finding something that does.