Tune SM90 small-M FP4 schedules

[mgoin]

Summary

Tune the SM90 scheduler for the small-M FP4 paths exercised by MXFP4/NVFP4 A16 and grouped-FP8 activations.

• use a 256-thread N32 warp geometry for grouped-scale A8 when an N128 tile is legal, while keeping block K256 when it is legal
• preserve legal block/warp nesting while fitting N and K
• improve indexed A16/B4 residency with an M8 schedule, up to three resident CTAs, and a same-wave N/K rebalance
• use the existing N128/K64 TMA schedule for small dense A16/B4
• add focused heuristic and exact MXFP4/NVFP4 kernel correctness coverage

This keeps the existing weight layouts and changes only launch scheduling.

Root cause

The grouped-scale SM90 path selected block (M, 128, 256) with warp (M, 16, 128), producing a 512-thread CTA. On the reported Kimi projections this caused very low eligible-warp activity and a large latency cliff.

The A16/B4 paths had separate small-M issues: insufficient resident CTAs, excessive routed-row padding, stale warp K after divisibility fitting, and output grids that did not always justify wide N tiles or Stream-K.

H200 performance

Grouped FP8 + MXFP4, E=12/top-8 balanced routing:

Projection	M	old	new	speedup
W13 N6144 K3584	8	560.1 us	66.9 us	8.4x
W13 N6144 K3584	64	1853.0 us	120.1 us	15.4x
W2 N3584 K3072	8	296.9 us	41.2 us	7.2x
W2 N3584 K3072	64	963.8 us	70.2 us	13.7x

The new geometry is neutral at M=1-4 and about 8x faster at M=256. A block-K128 ablation regressed the small endpoints, so this change does not copy a blanket block-K cap.

Matched BF16 FP4 expert-GEMM sweeps versus Marlin:

• MXFP4 gate/up: worst loss through M64 drops from 36% to 7.9%; Humming is 16-22% faster at M128-256.
• NVFP4 gate/up: one M4 point remains 10.1% behind; the other measured small-M points are 0.4-5.5% behind.
• NVFP4 down: M1 remains 14.2% behind; other measured points are 2.6-8.1% behind.
• Dense NVFP4 gate/up: worst measured loss drops to 6.7%.

The remaining NVFP4 down gap is dominated by E4M3 scale application rather than GEMM scheduling and is intentionally outside this draft.

Validation

• ../.venv/bin/python -m pytest tests/test_sm90_heuristics.py -q: 11 passed
• chg run --gpus 1 --timeout 45m -- ... test_mxfp4.py test_special_paths.py -v: 18 passed, 1 skipped
• chg run --gpus 1 --timeout 45m -- ... humming/tests/kernels/humming -q: 248 passed, 84 skipped
• Ruff checks on all four changed files: passed
• vLLM Humming MoE integration selection: 16 passed
• Qwen3-30B-A3B-MXFP4A16 GSM8K: 0.8976 accuracy, 0 invalid rate
• Qwen3-30B-A3B-NVFP4 GSM8K: 0.8878 accuracy, 0 invalid rate
• full-and-piecewise CUDA graph capture: all 51 configured sizes from 1 through 512 for both models; 128-question replay smoke completed with zero invalid outputs

Scope and known limitations

This draft does not claim that FP16 NVFP4 is uniformly ahead of Marlin; several FP16 points remain 10-14% behind. It also does not add vLLM dispatch or DeepEP plumbing for MXFP4 + block-FP8 group128, and it does not claim NVFP4 grouped block-FP8 or MXFP8 activation support.

The longer-term proposal is to replace the ordered SM90 mutations with bounded schedule candidates, one shared legality/resource analyzer, and sparse offline tuning records. This draft stages the measured schedules and regression fixtures that proposal must preserve.

Duplicate-work check

No open inclusionAI/humming PR matched SM90 A16/FP4 heuristics, H200/Marlin, or indexed A16.

This is not a duplicate of merged Humming #57: that change is already in the base and provides the general indexed-A16 foundation. This draft adds FP4-specific small-M residency/tile selection, dense FP4 scheduling, and grouped-A8 geometry. vLLM #51332 covers application-side block-FP8 dispatch and DeepEP plumbing and is not duplicated here.

Disclosure

AI assistance was used in developing and validating this change.