Peering into the CUTLASS Kernel Scheduler: How SM120's Auto-Selection Became a Performance Bottleneck

In the high-stakes world of large language model inference, every microsecond counts. When deploying the GLM-5-NVFP4 model across eight NVIDIA RTX PRO 6000 Blackwell GPUs, the assistant had already established that the model was compute-bound rather than communication-bound — a significant finding that shifted the optimization focus from network topology to kernel efficiency. But the question remained: why were the FP4 GEMM kernels underperforming so dramatically? Message 897 captures the moment the assistant drilled down to the root cause, uncovering a critical detail in CUTLASS's kernel scheduling that explained the performance ceiling.

The Context: A Performance Mystery

The preceding messages had painted a stark picture. The assistant had benchmarked FP4 matrix multiplication kernels on the SM120 architecture (Blackwell) and discovered that only small tile configurations worked reliably. Tiles of size 128×128×128 and 128×128×256 initialized and ran, but the larger, more compute-efficient tiles — 128×256×128 and 256×128×128 — failed during initialization with a "Failed to initialize cutlass TMA WS grouped gemm" error (see [msg 886]). This was not a minor issue: larger tiles mean more data reuse per thread, higher arithmetic intensity, and ultimately better utilization of the GPU's tensor cores. Without them, the kernels were leaving performance on the table.

The assistant had also established that the per-expert batch sizes during decode were tiny — roughly 16 tokens per expert at 512 concurrent requests, and 32 tokens at 1024 concurrency. At these sizes, the FP4 GEMM operations achieved a paltry 0.8 TFLOPS out of a theoretical peak of 1,850 TFLOPS (dense) — a utilization rate of 0.04% ([msg 891]). Even the best-case micro-benchmarks reached only ~1,300 TFLOPS, or 70% of the dense peak. Something fundamental was limiting the kernel's ability to scale.

The Discovery: KernelScheduleAuto

Message 897 opens with a succinct but powerful finding:

SM120 uses KernelScheduleAuto — CUTLASS auto-selects the kernel schedule. SM100 uses explicit 1SM/2SM warp-specialized schedules optimized for that arch. SM120 falls back to the auto scheduler which likely picks suboptimal configurations.

This is the key insight. CUTLASS, NVIDIA's CUDA template library for linear algebra subroutines, supports multiple kernel scheduling strategies. For the SM100 architecture (the previous Blackwell variant used in the GB200 Grace-Hopper systems), CUTLASS provides explicit, hand-tuned schedules — including 1SM (single streaming multiprocessor) and 2SM (dual-SM) warp-specialized variants that are carefully optimized for the hardware's characteristics. These schedules dictate how warps are assigned to different phases of the matrix multiplication: some warps handle the matrix A tile loads, others handle matrix B tile loads, and others handle the compute and epilogue. The specialization allows maximum overlap of memory access and computation.

For SM120, however, no such explicit schedules exist in the FlashInfer build. Instead, CUTLASS falls back to KernelScheduleAuto, which attempts to select a reasonable schedule automatically at compile time. The assistant's diagnosis — that this auto-selection "likely picks suboptimal configurations" — is grounded in a deep understanding of how CUTLASS's auto-scheduler works. Without architecture-specific tuning, the auto-scheduler may choose conservative tile sizes, suboptimal pipeline depths, or inefficient warp assignments that fail to exploit SM120's unique capabilities.

The Follow-Up: StageCountAutoCarveout

The message doesn't stop at the discovery. It immediately follows up with a targeted investigation into a related parameter: StageCountAutoCarveout. This parameter determines how many pipeline stages the kernel can fit given the shared memory budget. The assistant runs a bash command to inspect the relevant section of the launcher template:

ssh root@10.1.230.174 'sed -n "460,520p" /root/ml-env/lib/python3.12/site-packages/flashinfer/data/csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/launchers/moe_gemm_tma_ws_launcher.inl'

The output reveals template code referencing Is2SM and KernelPtrArrayTmaWarpSpecialized2SmBlockScaledMxNvf4UltraVs16Sm103 — confirming that the launcher is designed for warp-specialized, 2SM operation, but the SM120 instantiation uses the auto-scheduler variant. The assistant is methodically connecting the dots: the auto-scheduler, combined with the auto-carveout logic, likely computes that larger tiles require more shared memory than the 99KB available on SM120, causing the initialization failure for M128×N256 and M256×N128 configurations.

Assumptions and Reasoning

The assistant makes several implicit assumptions in this message. First, it assumes that the KernelScheduleAuto fallback is indeed the root cause of the tile initialization failures — a reasonable inference given that SM100's explicit schedules work correctly while SM120's auto-scheduler fails on the same tile configurations. Second, it assumes that the shared memory budget is the limiting factor, connecting the StageCountAutoCarveout mechanism to the observed failures. Third, it assumes that the auto-scheduler's decisions are suboptimal not just for the failing tiles but also for the working ones — meaning even the 128×128 tiles might be running at less than peak efficiency.

These assumptions are well-supported by the evidence gathered in previous messages. The assistant had already confirmed that the larger tiles fail with a specific TensorRT-LLM error message, and that the shared memory available on SM120 is 99KB per block — significantly less than what larger tiles typically require for their pipeline stages. The reasoning is sound, but it remains to be verified whether patching the auto-scheduler or providing explicit SM120 schedules would actually resolve the failures.

Input Knowledge Required

To fully appreciate this message, the reader needs familiarity with several concepts. Understanding CUTLASS's architecture — its template-based approach to kernel generation, the concept of tile sizes (M×N×K), and the role of warp specialization — is essential. Knowledge of GPU shared memory hierarchies, pipeline stages in matrix multiplication kernels, and the distinction between auto-scheduled and explicitly scheduled kernels is also required. Additionally, familiarity with the FlashInfer library's integration of TensorRT-LLM kernels for MoE (Mixture-of-Experts) GEMM operations provides the broader context for why this investigation matters.

The message also builds on the assistant's earlier work benchmarking FP4 kernels ([msg 875]), establishing the compute-bound nature of the workload ([msg 876]), and cataloging the working and failing tile configurations ([msg 886]). Without this foundation, the significance of the KernelScheduleAuto discovery would be lost.

Output Knowledge Created

This message creates several valuable pieces of knowledge. First, it establishes that SM120 lacks explicit CUTLASS kernel schedules and relies on the auto-scheduler — a finding that has direct implications for anyone deploying models on Blackwell GPUs. Second, it identifies StageCountAutoCarveout as the likely mechanism through which the shared memory budget constrains tile size selection, providing a specific target for optimization. Third, it sets up a clear hypothesis: providing explicit SM120 schedules (or overriding the auto-carveout calculation) could unlock larger tile configurations and improve kernel efficiency.

The message also implicitly documents a debugging methodology: when kernel initialization failures occur, trace them through the template instantiation chain to identify which scheduler variant is being used, then compare against architectures where the same tiles work. This approach is generalizable to many CUTLASS-based systems.

The Thinking Process

What's visible in the assistant's reasoning is a systematic narrowing of possibilities. The earlier messages had established what was failing (larger tile configs) and where it was failing (the initializeMoeGroupedGemm function). Message 897 narrows further to why: the auto-scheduler. The assistant doesn't just report the finding — it immediately validates it by examining the relevant source code, confirming that SM120 indeed uses KernelScheduleAuto while SM100 uses explicit schedules.

The thinking also shows a preference for architectural explanations over superficial ones. Rather than attributing the failure to a generic "bug" or "incompatibility," the assistant connects it to a specific design choice in CUTLASS's template metaprogramming — the absence of SM120-specific schedule instantiations. This reflects a deep understanding of the software stack, from the high-level FlashInfer library down to the CUTLASS template machinery.

The Broader Significance

This message represents a turning point in the optimization journey. Having established that the model is compute-bound and that the FP4 kernels are severely underutilized, the assistant now has a concrete hypothesis for why the kernels are underperforming. The auto-scheduler, combined with shared memory constraints, is limiting tile size selection and preventing the kernels from achieving higher arithmetic intensity. This insight will inform the next phase of optimization: either finding a way to provide explicit SM120 schedules, overriding the auto-carveout logic, or working around the tile size limitations through other means (such as increasing batch sizes through higher concurrency).

The discovery also has implications beyond this specific deployment. Anyone using FlashInfer or TensorRT-LLM on Blackwell GPUs for FP4 inference may encounter the same limitations. The assistant's investigation provides a roadmap for diagnosing and potentially resolving similar issues, contributing to the broader understanding of kernel optimization on NVIDIA's latest architecture.