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

[Missyinlll]

Summary

• Extend the existing SM90 indexed-A16 candidate policy with routed-M-aware BLOCK_M selection.
• Use deterministic per-expert tile estimates instead of relying only on coarse average-M thresholds.
• Add a small padding/tile trade-off guard around BM transitions to avoid overly aggressive increases in BLOCK_M.
• Prefer legal half-K (BK=64) configs for targeted W4A16 indexed cases when BM >= 32.
• Keep behavior unchanged outside the targeted SM90 W4A16 indexed path.

Context

This builds on the SM90 indexed-A16 scheduling and bounded candidate framework introduced in #58 and #59.

The existing candidate flow already owns the indexed-A16 base schedule, legality checks, residency analysis, and schedule transforms. This PR keeps that structure intact and only extends the W4A16 indexed BLOCK_M selection and half-K preference.

It is based on the latest main including #65 and preserves the short-K large-M scheduling guard added there.

Motivation

The existing indexed-A16 policy can switch to a larger BLOCK_M too aggressively near routed-M boundaries.

For MoE workloads, only a small subset of experts may cross the next BM boundary, while increasing BM for all experts can introduce substantially more padding for only a small reduction in M tiles.

The updated selector uses deterministic per-expert tile estimates to choose BM and applies a small guard when a larger BM provides little tile-count reduction but significantly increases padding.

This also avoids the observed BK=128 performance cliff for larger W4A16 indexed tiles by preferring legal half-K configs when BM >= 32.

Validation

Benchmarked on an NVIDIA H100 80GB HBM3 against the latest upstream/main.

Representative workloads:

• GLM-5.2

  • W13: N=4096, K=6144
  • W2: N=6144, K=2048
  • 256 experts, top-k=8

• DeepSeek-V4

  • W13: N=4096, K=4096
  • W2: N=4096, K=2048
  • 256 experts, top-k=6

132 routed-M cases were tested in total.

Performance:

• GLM-5.2 overall geomean: +9.88%
• DeepSeek-V4 overall geomean: +11.85%
• All-case geomean: +10.83%
• Median speedup: +6.34%
• Regressions >1%: 0
• Regressions >2%: 0

Several shapes around BM transition regions show substantially larger gains:

Profile	Workload	M	Upstream BM/BN/BK	New BM/BN/BK	Upstream p50	New p50	Speedup
GLM-5.2	W13	6152	32/256/128	24/256/128	2442.2 us	1504.7 us	+62.31%
GLM-5.2	W2	4920	32/256/128	24/256/128	1079.8 us	759.5 us	+42.17%
DeepSeek-V4	W13	6138	32/256/128	24/256/128	1476.2 us	996.5 us	+48.13%
DeepSeek-V4	W2	7362	32/256/128	32/256/64	778.4 us	526.0 us	+47.97%

The largest improvements occur where the existing policy selects a substantially more expensive configuration near a BM/BK transition, while unchanged regions retain the existing configuration and performance.

Correctness:

• All 132 cases passed
• Independent FP32 matmul reference
• rtol=0.02, atol=2.0
• Baseline and candidate use the same inputs, weights, scales, and deterministic routing

[Missyinlll]

@mgoin This PR is also related to the SM90 indexed A16 path, and I noticed your recent work in #65. I’d really appreciate any feedback on the heuristic approach when you have time. Thanks!