Peering into the CUTLASS SharedStorage: A Micro-Investigation into FP4 GEMM Tile Failures on Blackwell

The Message

ssh root@10.1.230.174 'sed -n "250,300p" /root/ml-env/lib/python3.12/site-packages/flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp'
    cutlass::bits_to_bytes(size(take<0,2>(SmemLayoutA{})) * sizeof_bits<ElementA>::value));

  static constexpr uint32_t TmaTransactionBytesNK = static_cast<uint32_t>(
    cutlass::bits_to_bytes(cosize(take<0,2>(SmemLayoutSFB{})) * cute::sizeof_bits_v<ElementSF>) +
    cutlass::bits_to_bytes(size(take<0,2>(SmemLayoutB{})) * sizeof_bits<ElementB>::value));

  static constexpr uint32_t TmaTransactionBytes = TmaTransactionBytesMK + TmaTransactionBytesNK;

  struct SharedStorage {
    struct TensorS...

At first glance, this message appears trivial: a single sed command that extracts 50 lines from a C++ header file, with the output truncated mid-struct. But this seemingly minor action sits at a critical juncture in a much larger investigation into GPU kernel efficiency. To understand why this message matters, we must examine the reasoning chain that led to it, the assumptions it tests, and the knowledge it seeks to create.

The Context: A Performance Investigation at the Kernel Level

The broader session (Segment 7 of the conversation) is a deep-dive into optimizing inference throughput for the GLM-5-NVFP4 model running on 8× NVIDIA RTX PRO 6000 Blackwell GPUs (SM120 architecture). The assistant had already established that the model is compute-bound rather than communication-bound — a TP4+PP2 configuration was 2× slower than TP8, ruling out allreduce latency as the primary bottleneck. This shifted focus squarely onto the FP4 GEMM kernel efficiency.

Earlier measurements revealed a troubling picture: during inference, the GPUs drew only ~235W out of their 600W TDP — a mere 39% utilization. The CUTLASS kernels achieved respectable ~1,300 TFLOPS (70% of dense peak) for large matrices, but during actual decode, per-expert batch sizes of ~16–64 tokens meant the GEMMs were operating at a tiny fraction of peak — as low as 0.02–3% of theoretical maximum. The assistant had identified that larger CUTLASS tile configurations (M128×N256 and M256×N128) could potentially improve small-batch efficiency, but these tiles were crashing during initialization with an "Error Internal" from initializeMoeGroupedGemm.

Why This Message Was Written

Message 917 is the culmination of a careful diagnostic chain. The assistant had already:

  1. Traced the crash source (msg 907–909): Located the failing code path in moe_gemm_tma_ws_launcher.inl, where gemm.initialize(args, ...) returned kErrorInternal even though can_implement had passed.
  2. Measured shared memory limits (msg 910–912): Discovered that SM120 has 99KB opt-in shared memory per block (101,376 bytes) and 100KB per SM (102,400 bytes). This was done through a series of Python queries, with attribute name corrections along the way (e.g., shared_memory_per_multiprocessor instead of max_shared_memory_per_multiprocessor).
  3. Estimated SMEM requirements (msg 913): Built a Python script to calculate shared memory consumption for different tile configurations, revealing that the failing tiles (128×256 and 256×128) would need ~91KB for just the default epilogue, leaving only ~34KB for the mainloop — insufficient for even one pipeline stage.
  4. Formulated a hypothesis (msg 914): Concluded that the epilogue output tile (M×N×sizeof) is the real problem, but then immediately expressed doubt: "Actually wait — my epilogue estimate might be wrong. CUTLASS epilogues use different strategies." This last point is crucial. The assistant realized that their SMEM estimation was based on assumptions about the epilogue's memory layout — specifically, that the output staging required M×N×2 bytes (BF16) or M×N×4 bytes (FP32 accumulator + BF16 output). But CUTLASS might use more sophisticated epilogue strategies (e.g., register-based accumulation, warp-level reductions) that change the memory footprint. The Python estimate was a useful approximation, but to truly understand why the tiles fail, the assistant needed to see the actual CUTLASS source code for the SM120 block-scaled collective. Message 917 is the direct result of this realization. The assistant reaches for the actual header file — sm120_blockscaled_mma_array_tma.hpp — to examine the SharedStorage struct and the TMA transaction byte calculations. This is a shift from estimation to verification, from modeling to direct inspection.## The Assumptions Embedded in This Message This message carries several implicit assumptions worth examining: Assumption 1: The SharedStorage struct is the key to understanding the failure. The assistant assumes that by examining the SharedStorage layout in the CUTLASS header, they can determine exactly how much shared memory each tile configuration requires. This is a reasonable assumption — CUTLASS computes SMEM usage at compile time via sizeof(SharedStorage), and a static_assert in the kernel checks this against sm120_smem_capacity_bytes (101,376 bytes, as discovered in msg 919). However, the assistant is looking at the collective (mainloop) shared storage, not the full kernel's shared storage which includes epilogue memory. The truncated output shows struct TensorS... — the struct definition is cut off, and the epilogue's contribution to shared memory is defined elsewhere. Assumption 2: The TMA transaction byte calculations are relevant to the SMEM overflow. The assistant prints TmaTransactionBytesMK and TmaTransactionBytesNK — these are the bytes transferred via TMA (Tensor Memory Accelerator) per pipeline stage, not the shared memory allocation size. While related (TMA loads data into shared memory buffers), these are not directly the SMEM footprint. The assistant is gathering clues from multiple angles. Assumption 3: The failing tiles fail for SMEM reasons, not for other reasons. The assistant has already ruled out can_implement failing (it passes), and the "Error Internal" from initialize is consistent with runtime SMEM allocation failure. But there could be other causes — unscheduled instruction sequences, register pressure, or hardware limitations specific to the TMA warp-specialized kernel on SM120. The assistant is narrowing down the hypothesis space but hasn't definitively proven the SMEM theory. Assumption 4: The Python SMEM estimation model is directionally correct. The assistant's earlier Python script assumed specific data types (FP4 packed at 0.5 bytes/element, FP8 scales at 1 byte per 16 elements) and a simple epilogue model (BF16 output staging, FP32 accumulator). These are educated guesses based on the block-scaled FP4 format, but the actual CUTLASS implementation might use different layouts — for instance, the scales might be stored in a different format, or the epilogue might use register accumulation to avoid shared memory entirely.

Input Knowledge Required

To fully understand this message, one needs:

  1. CUTLASS architecture knowledge: Understanding that CUTLASS kernels are composed of a mainloop (which loads tiles via TMA and computes MMA operations) and an epilogue (which finalizes the output, potentially applying activation functions or scaling). The SharedStorage struct is a union or struct of all temporary buffers needed by both phases.
  2. SM120 Blackwell architecture: Knowing that SM120 has 100KB of shared memory per SM (configurable up to 99KB per block with opt-in), and that the TMA engine handles data movement between global memory and shared memory. The 99KB limit is a hard constraint that tile configurations must respect.
  3. FP4 block-scaled format: Understanding that FP4 weights are stored as 4-bit values packed into bytes (2 values per byte), with FP8 scale factors shared across blocks of 16 elements. This affects the memory footprint calculations.
  4. FlashInfer's integration of CUTLASS: The file path (flashinfer/data/cutlass/include/cutlass/gemm/collective/sm120_blockscaled_mma_array_tma.hpp) reveals that FlashInfer vendors its own copy of CUTLASS headers. The assistant is probing a specific version of CUTLASS that may have SM120-specific collective implementations not yet merged upstream.
  5. The MoE grouped GEMM context: The failing tiles are used in a Mixture-of-Experts grouped GEMM kernel, where multiple small matrix multiplications (one per expert) are batched together. The tile dimensions (M, N, K) correspond to the batch size per expert, the hidden dimension, and the intermediate dimension respectively.

Output Knowledge Created

This message creates several forms of knowledge:

Direct output: The contents of lines 250–300 of sm120_blockscaled_mma_array_tma.hpp, showing:

The Thinking Process Visible in the Reasoning

The assistant's reasoning in the messages leading up to and following msg 917 reveals a methodical, hypothesis-driven investigation:

  1. Observe symptom: Larger CUTLASS tiles crash during initialization.
  2. Form hypothesis: SMEM overflow — the tiles require more shared memory than available.
  3. Gather evidence: Measure SMEM limits on SM120 (99KB opt-in).
  4. Build model: Estimate SMEM requirements for each tile configuration using a Python script.
  5. Test model: The model predicts the failing tiles need ~91KB+ — close to the 99KB limit, suggesting SMEM pressure.
  6. Identify uncertainty: The epilogue memory model might be wrong.
  7. Seek ground truth: Examine the actual CUTLASS source code to understand the real SharedStorage layout. This is classic scientific debugging: move from symptom to hypothesis to prediction to verification. The assistant doesn't stop at the Python model — they recognize its limitations and seek primary sources. Message 917 is the "seek ground truth" step. The truncated output is itself informative. The sed command captures lines 250–300, but the SharedStorage struct begins at line 258 and the output cuts off mid-struct. This truncation is a natural consequence of the terminal output being captured, but it means the assistant doesn't get the full picture from this single command. The very next message (msg 918) acknowledges this: "This is the mainloop only. The total SMEM is mainloop + epilogue." The assistant then pivots to examining GemmKernel::SharedStorageSize and the static_assert that checks against sm120_smem_capacity_bytes.

Why This Matters

This message is a microcosm of the entire optimization effort. It represents the moment when the investigation shifts from high-level benchmarking ("the model is compute-bound") to low-level kernel analysis ("why does this specific tile configuration crash?"). The assistant is operating at the boundary between application code and hardware — reading CUTLASS collective implementations to understand how shared memory is allocated for FP4 block-scaled GEMMs on Blackwell.

The ultimate goal — fixing the failing tiles to enable larger CUTLASS configurations — could yield significant throughput improvements. If the M128×N256 or M256×N128 tiles can be made to work, they would allow the MoE grouped GEMM kernel to process more tokens per expert per iteration, improving arithmetic intensity and bringing GPU utilization closer to the 600W TDP. But first, the assistant must understand why they fail. Message 917 is a step toward that understanding.

In the broader narrative of the session, this investigation into tile failures represents one of several parallel optimization paths being explored. The assistant has already achieved a 28% throughput improvement through server parameter tuning (raising --max-running-requests to 2048 and --num-continuous-decode-steps to 8), and is documenting optimization approaches as glb5improvement-xx.md files. The tile-fixing effort, if successful, could unlock further gains by enabling more efficient GEMM computation on the Blackwell architecture's unique shared memory constraints.