Prune SM90 kernel configs whose accumulator does not fit in registers

[hamuzhan]

Prune SM90 kernel configs whose accumulator does not fit in registers

Problem

The SM90 library generator emits kernel configs whose FP32 accumulator cannot
fit in the register file. Each math warp group holds its accumulator fragment
in registers:

Schedule	Accumulator regs/thread
TmaWarpSpecialized, TmaWarpSpecializedPingpong	cta_m * cta_n / 128 (full CTA tile per math warp group)
TmaWarpSpecializedCooperative	cta_m * cta_n / 256 (tile split across two math warp groups)

Above 128 accumulator registers per thread, the accumulator no longer fits in
the 255-register budget together with mainloop state, so ptxas spills it to
local memory on every mainloop iteration.

Evidence (H100, CUDA 13.2, 16-bit dense GEMM, M=N=K=8192, TN)

Throughput by per-warp-group accumulator size. FP16 and BF16 give the same
numbers:

accum regs/thread	example configs	measured
128 (fits)	128x128 and 64x256 pingpong; 128x256 and 256x128 cooperative	532-898 TFLOP/s
256 (spills)	128x256 and 256x128 pingpong and warpspecialized; 256x256 cooperative	43-90 TFLOP/s
512 (spills more)	256x256 pingpong and warpspecialized	18-20 TFLOP/s

Nsight Compute for ..._128x256x64_..._pingpong_epi_tma (BF16 and FP16 report
the same count):

• Local Memory Spilling Requests: 710,525,984 per launch (fitting configs: 0)
• 43% of executed instructions are local loads/stores
• 100% of LDL/STL instructions come from register spilling
• SM throughput 12.4% vs 92.8% for the same tile with the cooperative schedule

Compile-time check with cuobjdump --dump-resource-usage on the built library
(FP16, TN, cluster 1x2, epi_tma; BF16 matches):

tile	schedule	per-WG accum regs	stack frame (spills)
128x128	cooperative	64	0 B
128x128	pingpong	128	0 B
64x256	pingpong	128	0 B
128x256	cooperative	128	0 B
256x128	cooperative	128	0 B
128x256	pingpong	256	1928 B
256x128	pingpong	256	1160 B
256x256	cooperative	256	1904 B
256x256	pingpong	512	6456 B

The stack frame becomes non-zero at exactly 128 registers per thread, matching
the formula above.

This matches limits already present in the C++ layers

The same limit already exists in the kernel and builder code. This PR applies
it in the library generator as well:

• The ping-pong kernel computes the same value and switches to a 24/240
  producer/consumer register split when it reaches 208 registers per thread
  (include/cutlass/gemm/kernel/sm90_gemm_tma_warpspecialized_pingpong.hpp:143-148).
  For the pruned configs the accumulator alone needs 256 or more registers, so
  the 240-register consumer budget still spills.
• The CollectiveBuilder KernelScheduleAuto path selects ping-pong only when
  TileShape_M == 64 and cooperative otherwise
  (include/cutlass/gemm/collective/builders/sm90_gmma_builder.inl:1021-1022),
  so it never picks ping-pong for these tiles.

The builder path already avoids these configs. Only the library generator still
emits them.

Change

Prune these configs in get_valid_schedules at instantiation level 0, using
the existing early-pruning pattern. Higher instantiation levels
(CUTLASS_LIBRARY_INSTANTIATION_LEVEL with a non-zero pruning digit) still emit
them, so exhaustive-search builds are unchanged. ScheduleAuto is gated by the
cooperative (split-tile) budget, since the builder resolves large tiles to a
cooperative schedule.

The rule uses the accumulator size, which is 4 bytes for F16/BF16/TF32/FP8/INT8,
so it applies to all of these paths. The measurements above are from the 16-bit
path.

Impact

• Default 16-bit SM90 manifest: 2417 to 1705 kernels (-712, -29%). Removed: 376
  pingpong, 120 warpspecialized, 176 cooperative (256x256), 40 auto (256x256).
  Library build time drops with the kernel count.
• Other families at the default level (same rule, F32/S32 accumulator): fp8 e4m3
  783 to 697 (-11%), int8 155 to 141 (-9%), tf32 287 to 215 (-25%), grouped f16
  753 to 465 (-38%). All generator paths configure without error.
• A sweep of 22k measurements over LLM-style shapes (M in [16, 16384], N,K from
  Llama-8B/70B layers plus a 128k-vocab LM head, FP16 and BF16, default and
  extended instantiation levels) shows none of the pruned configs reaching the
  top 3 for any (shape, M) point.
• No change for CUTLASS_LIBRARY_INSTANTIATION_LEVEL >= 1xxx builds (checked:
  level 1242 still emits the pruned configs).

Validation

• Configure-time kernel diff on the 16-bit filter shows the predicted 712
  removals and 0 additions.
• Rebuilt cutlass_profiler from the pruned manifest. All 1904 surviving gemm
  operations verify Passed against the device reference (m=640, n=768, k=1024;
  beta=0.5 for source-C kernels, beta=0 for void-C).
• Pruned kernel names are absent from --mode=enumerate on the pruned manifest.
• cutlass_test_unit_gemm_device_tensorop_sm90 is unaffected (unit tests do not
  go through the library generator); spot re-run passes.

Repro:

cmake -S . -B build -GNinja -DCUTLASS_NVCC_ARCHS=90a \
  -DCUTLASS_LIBRARY_KERNELS="cutlass3x_sm90_tensorop_gemm_f16_f16_*,cutlass3x_sm90_tensorop_gemm_bf16_bf16_*"
# before/after: find build/tools/library/generated -name '*.cu' | wc -l

./build/tools/profiler/cutlass_profiler --operation=Gemm \
  --kernels="cutlass3x_sm90_tensorop_gemm_bf16_bf16_f32_void_bf16_128x256x64_1x2x1_0_tnn_align8_warpspecialized_pingpong_epi_tma" \
  --m=8192 --n=8192 --k=8192   # ~44 TFLOP/s on main

[hamuzhan]

@XiaoSongXS

[XiaoSongXS]

i think u likely tag the wrong person?

[hamuzhan]

@XiaoSongXS Sorry to bother you. Since I saw you were a reviewer on a PR related to Blackwell, I figured you were the person in charge of this topic. sorry again

[XiaoSongXS]

• no worry. i'm just concern likely ur question is not routed to the correct PIC.
• for the sm90 reg spill thing, yes reg spill sometimes can happen on hopper kernel as both mma output and epilogue computation use same set of register. And this is why in sm90 kernel, sometimes u see register reconfig and carefully designed epi tile size.
• the purpose of the sm90 generator is simply generate some kernel to show those kernel work, in real production, user commonly pick a subset of kernel they want and generate it, meaning that they will write their own generator script. so this is less a issue on generator script (as it is demo purpose)
• also i highly encourage u to do a profile for fix problems size X various sm90 cta tile X cluster size combine to understand perf relationship between kernel config and final perf. reg spill is indeed a issue, but this is not the real perf limitor of kernel. real limitor of good gemm perf is tensorcore utilization, which commonly require large cta tile and large mma

[hamuzhan]

Thanks, this changes how I’m thinking about the issue. I’ll expand the profiling across fixed problem sizes, CTA tiles, cluster shapes, and schedules before deciding whether generator-side pruning is appropriate