tune(sm90): improve W4A8 indexed MoE block-M selection

[Missyinlll]

Summary

• Add a routed-M-aware BLOCK_M selector for SM90 indexed MoE GEMMs with FP8 E4M3 activations and FP4 E2M1 weights using weight scale group size 32.
• Generate legal BM candidates dynamically from the active SM90 max_block_m.
• Use deterministic per-expert estimates with a local wave/padding correction.
• Keep the existing configuration policy unchanged outside the targeted path.

Motivation

The existing SM90 indexed MoE heuristic selects BLOCK_M by minimizing an estimated block count based on a sampled routing distribution. For medium and large routed-M workloads, this can select overly large BM values.

For a given minimum per-expert M-tile count, a smaller BM that preserves the same tile count can reduce padding without increasing the primary tile count. The new selector uses the expected routed rows per expert to determine the primary BM regime, then considers nearby legal candidates when the estimated last CTA wave is poorly utilized.

Implementation

The new selector is enabled only when all of the following are true:

• SM90
• num_experts > 0
• GemmType.INDEXED
• activation dtype is FP8 E4M3
• weight dtype is FP4 E2M1
• weight scale group size is 32

Legal BM candidates are generated dynamically as 8-row-aligned values from 8 through the active max_block_m, while retaining the existing dtype-specific legality constraints.

The primary BM target is derived from:

expected_m_per_expert = routed_m / num_experts

For workloads above max_block_m, the selector estimates the minimum number of M tiles required per expert and chooses a compact BM that preserves that tile count. A local candidate window is then evaluated using estimated CTA-wave utilization and padding.

For the validated W4A8 get_config2 path:

• BN=128 is unchanged.
• BK is intentionally derived from the BM selected by the existing heuristic, so this change does not introduce a new K-reduction configuration.
• num_stages=4 is unchanged.
• Stream-K policy is unchanged.
• Warp-N and warp-K are unchanged; warp-M follows the selected BM.
• No kernel, TMA layout, multicast, quantization, or epilogue implementation is changed.

The selector is integrated into the current SM90 candidate-based tuning structure and uses the existing candidate legality/resource checks.

The experimental runtime switch used during development has been removed. The new selector is the default behavior for the eligible path.

Performance

Measured on an NVIDIA H100 80GB (SM90) using GLM-5.2-style indexed MoE shapes:

• W13: K=6144, N=4096
• W2: K=2048, N=6144
• E=256, top_k=8
• FP4 E2M1 weights
• FP8 E4M3 activations
• weight scale group size 32

Routed M denotes post-top-k routed rows. With top_k=8, the corresponding input-token count is routed_m / 8.

Routing is deterministic and evenly distributed across 256 experts for the reported cases.

The standalone kernel benchmark uses GEMM-only CUDA-event timing. Weight preparation, activation quantization, setup, JIT compilation, and warmup are outside the measured interval.

The A/B comparison uses separate source checkouts and subprocesses:

• baseline: upstream main / Refactor SM90 tuning around bounded schedule candidates #59 at 8611853cbf60cab5abf1f1c01fd96c8f85f9315b
• final: current PR head at 660444371df022fe8e153bb53b4449eb3c1ade09

No forced-BM oracle or production environment switch is used.

GEMM	Routed M	Baseline config	Final config	Baseline us	Final us	Speedup
W13	2048	BM24/BN128/BK256	BM8/BN128/BK256	1234.725	1203.659	+2.58%
W13	4096	BM32/BN128/BK256	BM16/BN128/BK256	1335.060	1246.423	+7.11%
W13	8192	BM56/BN128/BK128	BM32/BN128/BK128	1831.129	1537.968	+19.06%
W13	16384	BM96/BN128/BK128	BM64/BN128/BK128	3731.747	2055.866	+81.52%
W2	2048	BM24/BN128/BK256	BM8/BN128/BK256	712.178	665.857	+6.96%
W2	4096	BM32/BN128/BK256	BM16/BN128/BK256	752.906	699.421	+7.65%
W2	8192	BM56/BN128/BK128	BM32/BN128/BK128	1100.789	920.162	+19.63%
W2	16384	BM96/BN128/BK128	BM64/BN128/BK128	2169.567	1265.137	+71.49%

Across the 8 latest-main standalone GEMM cases:

• Geometric mean speedup: +24.05%
• Median speedup: +13.35%
• Regression cases: 0 / 8

The gains increase with routed M, reaching about 19% at routed M=8192 and 71-82% at routed M=16384 for the tested W13/W2 shapes.

A reverse-order check on the two smallest W13 cases reproduced the gains at approximately:

• routed M=2048: +3.19%
• routed M=4096: +6.00%

This check was run with the PR checkout first and the upstream checkout second, and supports that the small-M gains are not explained by fixed benchmark execution order.

A standalone fused MoE wrapper benchmark covering W13, activation, W2, and combine reported:

Routed M	Speedup
2,048	+3.56%
4,096	+5.60%
8,192	+16.54%
16,384	+69.73%

Geometric mean fused-MoE speedup: +21.27%, with no regressions across the four tested routed-M values.

The fused wrapper includes activation quantization inside the timed path, unlike the standalone GEMM-only benchmark.

This is an operator-level benchmark, not a model-level end-to-end benchmark.

Correctness

Current-head correctness validation used:

• baseline: upstream Refactor SM90 tuning around bounded schedule candidates #59 at 8611853cbf60cab5abf1f1c01fd96c8f85f9315b
• final: current PR head at 660444371df022fe8e153bb53b4449eb3c1ade09
• FP4 E2M1 weights with group size 32
• FP8 E4M3 activations
• GemmType.INDEXED
• deterministic indexed routing
• rtol=0.02
• atol=2.0

The same 8 W13/W2 shapes used in the standalone performance table were validated:

GEMM	Routed M	Baseline BM/BK	Final BM/BK	Max abs	Max rel	Mean abs	Status
W13	2048	BM24/BK256	BM8/BK256	0	0	0	PASS
W13	4096	BM32/BK256	BM16/BK256	0	0	0	PASS
W13	8192	BM56/BK128	BM32/BK128	0	0	0	PASS
W13	16384	BM96/BK128	BM64/BK128	0	0	0	PASS
W2	2048	BM24/BK256	BM8/BK256	0	0	0	PASS
W2	4096	BM32/BK256	BM16/BK256	0	0	0	PASS
W2	8192	BM56/BK128	BM32/BK128	0	0	0	PASS
W2	16384	BM96/BK128	BM64/BK128	0	0	0	PASS

Summary:

• Cases: 8
• PASS: 8
• FAIL: 0

The baseline logical output is used as the reference. BM-specific padded routing metadata is regenerated independently for each checkout, and only the final logical routed output is compared.

The fused MoE wrapper validation also passed at routed-M values:

2048, 4096, 8192, 16384

Repository tests additionally cover:

• selector dispatch scope
• dynamic BM candidate legality
• representative W13/W2 config selection
• FP4 E2M1 / FP8 E4M3 indexed MoE correctness
• preservation of the existing BK decision

Validation

Targeted tests:

pytest -q \
  tests/test_sm90_w4a8_heuristics.py \
  tests/kernels/humming/test_sm90_w4a8_indexed.py

The targeted tests passed on an H100/SM90 environment.

Additional checks:

python -m py_compile \
  tests/test_sm90_w4a8_heuristics.py \
  tests/kernels/humming/test_sm90_w4a8_indexed.py

git diff --check

Scope

This change is limited to the SM90 indexed MoE path with:

• FP8 E4M3 activations
• FP4 E2M1 weights
• weight scale group size 32
• num_experts > 0

The change does not modify:

• Dense GEMM dispatch
• Grouped-contiguous MoE dispatch
• Grouped-masked MoE dispatch
• other activation or weight dtype combinations
• other weight scale group sizes
• other GPU architectures
• kernel implementation
• quantization implementation
• TMA or multicast implementation
• epilogue logic

[mgoin]

@Missyinlll Please migrate to the new structure and retest now that my PRs landed, thanks!

[mgoin]

Can you also update the performance numbers compared to the latest main since other optimizations have landed? Just curious to understand directly what this affects

[Missyinlll]

Can you also update the performance numbers compared to the latest main since other optimizations have landed? Just curious to understand directly what this affects

Thanks! This PR is already based on the latest main after #58/#59 landed and integrated into the new tuning structure.

I’ve also rerun the benchmarks against latest main (8611853) using the current PR head (6604443) and updated the performance/correctness results accordingly.