tune(sm90): fold H200 W4A8 MoE calibrations into the generic grouped-scale policy

[Leoyzen]

Fixes #26

Reworked per review: no separate sm90_h200.py. H200 shares the GH100 die and the 132-SM grid with H100 SXM and H800, so the tuning is SM-count-driven, not device-name-driven. The three calibrated decisions now live in the generic sm90 grouped-scale policy, keyed on num_sms, and H100 SXM / H800 (also 132 SMs) pick them up for free.

Changes

All in sm90_policies.py, gated on problem.device.num_sms >= 128 (so H20 at 78 SMs and H100 PCIe at 114 SMs keep byte-identical behavior with current main; verified on real H20):

1. 2-CTA residency cap for grouped-scale candidates (port of _indexed_a16_ctas_per_sm pattern).
2. Large per-expert token counts (per_expert_m >= 48, N % 256 == 0) select a (64|128, 256, 128) tile with a single resident CTA and stream-K off, inserted as the highest-priority candidate.
3. Small-M W4A8 shapes select the wide (m, 512, 64) tile with 2 CTAs.

Supporting changes: grouped-scale path now carries the device grid size (include_grid_size=True); the vLLM downstream constraints (K-block <= 128, warp_n >= 32 when block_n % 32 == 0) are enforced by construction. sm90_h200.py, name dispatch, and HUMMING_FORCE_HEURISTICS are gone.

Why SM-count keying

Dispatch by device name was already broken in the field: our production H200 node reports an OEM-rebranded name (L20X), so the name-based dispatch in the previous revision never matched and the server ran generic SM90 heuristics. SM count survives rebranding. It is also more robust than hardcoding 132: H800 SM counts are ambiguous across public sources and H100 PCIe has 114 — reading sm_count at runtime picks correctly in all cases.

Results

Kernel bench (real H200, OEM-rebranded L20X, indexed MoE GEMM, 288 experts top-8, fp8e4m3 x int4-128, triton do_bench, vs generic SM90 on the same card):

shape	M	speedup
w13 (N=4096, K=4096)	64	1.05x
w13	512	1.77x
w13	2048	2.06x
w13	16384	1.89x
w2 (N=4096, K=2048)	64	1.07x
w2	512	1.85x
w2	2048	2.17x
w2	16384	2.01x

No regressions on the production packed-K path: configs are byte-identical with stock (verified end-to-end in a serving pod).

Config equivalence against the previous separate-class revision was checked shape-by-shape with scripts/config_diff_h200.py: small-M and large-M selections match exactly; mid-M (4096-8192) differs by one tile step within the same family (threshold table vs measured argmin), expected performance-equivalent.

Accuracy

GSM8K 5-shot (greedy completions), GLM-5.3-W4AFP8 served via vLLM (TP4 + EP), same checkpoint and server config, only heuristics differ:

• this PR: 0.9083
• stock (generic SM90): 0.8976

Within noise (~1.3 sigma). Raw-completion protocol, not comparable to official thinking-mode scores.

Tests

• pytest tests/ -k "sm90 or grouped" — 60 passed, 80 skipped (SM100-gated) on real H20; 66 passed, 2 skipped on CPU with a pinned device profile
• New tests/test_sm90_grouped_scale_sm_count.py (17 tests): calibration trigger at 132 SMs, gate-off byte-identical behavior at 114/None, 2-CTA cap, wide-tile and 256-N selections
• Dispatch verified on real H200 (OEM-rebranded) and real H20

Known limitations

• Calibrated against W4A8 grouped-scale shapes (GLM-5.3 family); the production vLLM packed-K layout routes W4A8 MoE through the legacy seed path and does not hit these candidates — serving deployments need unpacked K layout to pick up the calibrations (follow-up worth its own issue).
• Only benched on H200; H100 SXM shares the same 132-SM grid and die, so the occupancy-driven calibrations should transfer, but a spot-check bench would be good if anyone has one available.
• Dense-path tweaks from the previous revision are intentionally not carried over; H200 dense stays on generic sm90.

Related but non-overlapping: #64 (W4A8 indexed block-M selection, group-32 MXFP4) — this PR does not touch the generic block-M selector.

Assisted-by: opencode (GLM)

[jinzhen-lin]

Thank you. But we should not create a separate sm90_h200.py. The reason for creating sm90_h20.py separately is that its performance parameters are quite unique compared to other sm90 devices like H100/H200/H800, whereas H200 is a typical sm90 device and should be optimized directly in sm90.py.

[Leoyzen]

@jinzhen-lin Agree — H200 shares the 132-SM grid with H100/H800, so the tuning is really SM-count-driven rather than device-specific. Plan to rework: move the calibrated decisions (2-CTA residency cap, large-per-expert-M large-N tile with stream-K off, small-M wide tile) into the generic sm90 policy path keyed on num_sms, and drop the separate class + name dispatch. One question: the current sm90.py routes W4A8 grouped-scale shapes through select_grouped_scale candidates — do you prefer these tweaks inside sm90_policies.py candidate generation, or as a post-selection adjustment in sm90.py?

[Leoyzen]

Reworked as suggested — dropped the separate class and name dispatch, folded the three calibrated decisions into the generic sm90 grouped-scale policy keyed on num_sms (gate at >= 128, so H20 and H100 PCIe are untouched). One note from field testing: our production H200 reports an OEM-rebranded name, which is exactly why name-based dispatch never fired there — SM-count keying is immune to that. Bench numbers and the config-diff harness vs the old revision are in the description.