[Kernel] SM120: stop routing misaligned-M blockwise FP8 GEMMs to the small-M swapAB config

[lucifer1004]

Purpose

The SM120 blockwise FP8 GEMM dispatch routes every call with M % 4 != 0 to the
swapAB config (TILE_N=32), which was added in #38325 as a small-M/decode optimization:

bool swap_ab = (M <= 64) || (M % 4 != 0);

Under V1 chunked prefill, mixed prefill+decode batches routinely make M misaligned at
prefill scale (M ~ 8k). In an Nsight Systems capture of Qwen3.8-27B-FP8 (a dense 27B
hybrid linear-attention model: 64 layers, 3:1 gated-delta-net to full attention, FP8
e4m3 blockwise) serving 8192/1024-token requests at concurrency 32 on an NVIDIA RTX PRO
6000 Blackwell Server Edition (SM120), 82% of prefill-scale blockwise GEMM
instances (14400/17523) went through the swapAB path, where they averaged 4.84 ms vs
3.10 ms on the default 128x128x128 config. The misroute roughly doubled prefill GEMM
time end-to-end.

Force-routing both paths over a dense M x N x alignment matrix (M = 48..8192 including
M%4 in {0,1,2}, N in {5120, 15360, 34816}, K = 5120) shows:

• swapAB only wins at M <= ~80 (its intended decode regime);
• from M ~ 96-128 up, the non-swap configs win everywhere, by 1.3x-2.0x at prefill scale;
• outputs are bit-identical between the swapAB and non-swap paths at every measured
  misaligned point, so the M % 4 != 0 clause has no correctness role on SM120.

The fix restricts swapAB to its design regime:

bool swap_ab = (M <= 64);

Note: scaled_mm_blockwise_sm90_fp8_dispatch.cuh ((a.size(0) % 4) != 0) and
scaled_mm_blockwise_sm100_fp8_dispatch.cuh ((m < 16) || (m % 4 != 0)) carry the same
misaligned-M clause. They are untouched here because the measurements above are
SM120-only; the same recheck may be worthwhile on those targets.

Test Plan

• Microbenchmark driving torch.ops._C.cutlass_scaled_mm across an M x N x alignment
  matrix with both kernel paths force-routed (temporary env-var instrumentation, not
  included in this PR), comparing per-path latency and bitwise outputs.
• End-to-end serving benchmark before/after: 8192-token prompts, 1024-token outputs,
  concurrency 32, NVIDIA RTX PRO 6000 Blackwell Server Edition, Qwen3.8-27B-FP8.

Test Result

Microbenchmark (N=34816, K=5120):

M	before (ms)	after (ms)
64	0.13	0.13 (unchanged, still swapAB)
8192 (M%4==0)	3.86	3.89 (unchanged)
8193	7.89	4.12
8225	7.64	3.92

Outputs at M=8193/8225 are bit-identical before and after.

End-to-end (concurrency 32, 8K/1K):

metric	before	after
output throughput	424.6 tok/s	499.8 tok/s (+17.7%)
mean TTFT	9660 ms	6872 ms
mean TPOT	65.9 ms	57.3 ms

The patched build reaches parity with an SGLang baseline of the same model, shapes, and
hardware (501.2 tok/s), closing what was previously an ~18% gap that traced back to this
misroute.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @lucifer1004, CI is now available for this PR.

• /ci run starts a CI build.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /ci cancel cancels scheduled or running CI builds for this PR branch.

[ZJY0516]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84406 for commit 16aa5116499c.