perf(moe): pipeline the nvfp4 prefill scale loads on geforce sm_86 / sm_89

[jason-fxz]

The Triton prefill GEMM of the inline-dequant NVFP4 MoE loads its e4m3 block scales one byte per thread on the pointer path (sm_86 / sm_89, and sm_120 without TMA). One byte is too small for cp.async, so each K step issues a synchronous load, waits for it and passes the bytes through shared memory with a CTA barrier before the decode can start. On GeForce sm_86 / sm_89 (half-rate fp32-accumulate mma, 2 CTAs per SM at 235 registers) this was the largest loss after the MMA itself.

Changes

• kernel/triton/nvfp4_fused_moe.py: SCALE_WORDS constexpr (default off). The scales are read as int32 words, 4 per word, and spread over the 8 weight words each word scales; the loads are 4 bytes, so Triton pipelines them with cp.async. Needs FAST_DECODE == 1, pointer loads and BLOCK_SIZE_KB >= 32.
• kernel/nvfp4_moe.py: the Triton prefill tuning becomes one _PrefillTuning record per (rho class, compute capability), looked up as (class, cc) then (class, None). The sm_89 pair decode and the H100 byte decode move into the table. New entries:
  • GeForce sm_86 / sm_89: SCALE_WORDS on BLOCK_SIZE_M 64 GEMMs with K % 64 == 0 and contiguous, 4-byte aligned scales; 8 warps when K >= 2048 (gate_up), 4 otherwise.
  • GeForce sm_86: BLOCK_SIZE_M 64 from 100 tokens per expert (was 150). sm_89 keeps 150: there BLOCK_SIZE_M 64 still loses at 128 tokens per expert.
  • The Gluon steps stay outside the table (_prefill_block_m_steps).

All other GPUs get the same configs as before: prefill_config for 18 faked GPUs (H100, H20, A100, B200, RTX PRO 6000 and RTX 5090 with and without Gluon, RTX 4090, RTX 3080 Ti, RTX 3090, RTX A6000, RTX 6000 Ada, L40S, GB10, Orin, an unknown name, ROCm) x 6 (E, top_k) x 19 token counts differs only on sm_86 GeForce at 102 - 128 tokens per expert (BLOCK_SIZE_M 32 -> 64).

Results

fused_experts_nvfp4 (align + gate_up + act + down + sum) with random banks in 9 checkpoint shapes (Qwen3.6-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash, Gemma-4-26B-A4B, MiniMax-M2.5, GLM-4.7, GLM-5.2, GLM-5.3-Flash, MiniMax-M3) x T in {16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384}; triton.testing.do_bench median, main / PR / main / PR, each tree with its own extensions built in place. torch 2.14.1+cu130, Triton 3.8.0. Speedup over main (ac6a15d):

GPU (driver)	points the PR changes	other points	output vs main
RTX 4090, sm_89 (580.105)	20: 1.069x (1.03 - 1.09)	70: bitwise equal, 0.99 - 1.03	changed points: norm rel <= 3.0e-4
RTX 3080 Ti, sm_86 (595.71)	28: 1.117x (1.03 - 1.19)	62: bitwise equal, 0.95 - 1.04	changed points: norm rel <= 2.8e-4

T = 4096 / 8192 / 16384 (- = BLOCK_SIZE_M 32, unchanged):

model	RTX 4090	RTX 3080 Ti
Qwen3.6-35B-A3B	- / 1.05 / 1.05	1.10 / 1.12 / 1.12
Qwen3.5-122B-A10B	- / 1.08 / 1.09	1.08 / 1.14 / 1.10
Qwen3.8-Flash	- / 1.07 / 1.07	- / 1.11 / 1.10
Gemma-4-26B-A4B	1.04 / 1.03 / 1.03	1.06 / 1.04 / 1.03
MiniMax-M2.5	- / 1.09 / 1.08	1.10 / 1.13 / 1.14
GLM-4.7	1.09 / 1.09 / 1.09	1.13 / 1.14 / 1.13
GLM-5.2	- / 1.09 / 1.07	1.15 / 1.16 / 1.15
GLM-5.3-Flash	- / 1.06 / 1.06	1.19 / 1.13 / 1.14
MiniMax-M3	- / 1.07 / 1.07	1.15 / 1.13 / 1.15

On the changed points the decoded weights reach the MMA through shared memory instead of registers (Triton picks this layout once the scales come from the pipeline), so the accumulation order changes; outputs differ by at most 1 bf16 ulp. The RTX 3080 Ti also gains at T = 2048 on Gemma-4-26B-A4B (1.08x) and GLM-4.7 (1.10x).

Not measured: other GeForce sm_86 / sm_89 models (the tile tables already come from one card per architecture), T > 16384, end-to-end serving on these GPUs.

Tests

pytest tests/moe/test_nvfp4_backends.py tests/kernels/test_e4m3_compat.py tests/moe/test_fused_moe.py:

• RTX PRO 6000: 56 passed, 5 skipped (Gluon on) / 52 passed, 9 skipped (FREETOKEN_DISABLE_GLUON=1), same as main
• RTX 3080 Ti: 46 passed, 15 skipped, same as main
• RTX 4090: 49 passed, 11 skipped, 1 failed, same as main (test_e4m3_compat.py::test_forced_emu_matches_native, blk_aq_y)