perf(moe): speed up the nvfp4 moe prefill gemm on rtx gpus

[jason-fxz]

Prefill path of the inline-dequant NVFP4 MoE (fused_experts_nvfp4). Same banks and layout, no repack; decode is unchanged.

Changes

• kernel/nvfp4_moe.py (new): the NVFP4 MoE GEMM dispatch (kernel choice, tiles, switches, launch), moved out of moe/fused_nvfp4.py, which keeps only the layer flow (align -> gate_up -> act -> down -> sum). The kernels are templates; every tile and switch comes from the dispatch.
• Triton prefill kernel (all CUDA GPUs; ROCm unchanged):
  • fast e2m1 decode: lop3/prmt bit tricks and mul.rn.bf16x2 (fma.rn.bf16x2 before sm_90), bitwise equal to the old decode; the decode form is picked per device and tokens per expert (sm_89 keeps the pair decode below 48 tokens per expert);
  • TMA weight loads on sm_120 when K % 256 == 0 (scales by TMA too for BLOCK_SIZE_M <= 32).
• kernel/gluon/ (new): Gluon kernels that replace a Triton kernel on a given GPU when Triton 3.8 is present; FREETOKEN_DISABLE_GLUON=1 turns them off.
  • gluon/nvfp4/moe/sm120.py: warp-specialized prefill GEMM for sm_120 (TMA loader warp, decode warps, mma.sync warps); a PERSISTENT constexpr runs BLOCK_SIZE_M 64 tiles persistently with a third pipeline stage. The dispatch routes BLOCK_SIZE_M 64 / 128 to it on sm_120.
• Tests: the Gluon route against the Triton kernel for both kernel modes, 128- and 256-row tiles, K % 256 == 0 and not; skipped off sm_120 or without Gluon.

Results

Measured against main at 6b8f200 (the branch now sits on 31c6ef4; #732 only touches GGUF). torch 2.14.1+cu130, Triton 3.8.0.

Kernel: fused_experts_nvfp4 (align + gate_up + act + down + sum) with random banks in 6 checkpoint shapes (Qwen3.6-35B-A3B, Qwen3.5-122B-A10B, Qwen3.8-Flash, Gemma-4-26B-A4B, MiniMax-M2.5, GLM-4.7) x T in {16, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384}; triton.testing.do_bench median, main / PR / main / PR. Speedup over main, geomean (min - max), by tokens per expert:

GPU (driver)	all	< 36	36 - 96	>= 96	output vs main
RTX PRO 6000, sm_120 (595.91)	1.32x (1.09 - 1.66)	1.22x	1.44x	1.48x	<= 1 bf16 ulp (norm rel 2.5e-4)
RTX 5090, sm_120 (580.105)	1.17x (0.94 - 1.33)	1.17x	1.14x	1.17x	<= 1 bf16 ulp (norm rel 2.5e-4)
RTX 4090, sm_89 (580.105)	1.04x (0.99 - 1.15)	1.00x	1.11x	1.10x	bitwise equal
RTX 3080 Ti, sm_86 (595.71)	1.05x (1.00 - 1.13)	1.03x	1.07x	1.08x	bitwise equal

Gluon at T = 8192 (fused_experts_nvfp4, one run with main, the PR with FREETOKEN_DISABLE_GLUON=1 (Triton only) and the PR interleaved, two rounds; Gluon / Triton and PR / main are speedups, i.e. the Triton-only and main times over the PR time):

RTX PRO 6000:

model	main (6b8f200), us	PR, Triton only, us	PR, us	Gluon / Triton	PR / main	PR TFLOP/s	PR MFU
Qwen3.6-35B-A3B	2745	2446	1953	1.25x	1.41x	211	45%
Qwen3.5-122B-A10B	7479	6241	4717	1.32x	1.59x	262	56%
Qwen3.8-Flash	6128	5425	4400	1.23x	1.39x	183	39%
Gemma-4-26B-A4B	4320	3917	3141	1.25x	1.38x	248	53%
MiniMax-M2.5	10801	9072	6671	1.36x	1.62x	278	59%
GLM-4.7	17301	14411	10560	1.36x	1.64x	293	63%
geomean				1.30x	1.50x

RTX 5090:

model	main (6b8f200), us	PR, Triton only, us	PR, us	Gluon / Triton	PR / main	PR TFLOP/s	PR MFU
Qwen3.6-35B-A3B	3009	2834	2612	1.09x	1.15x	158	75%
Qwen3.5-122B-A10B	8686	7841	7169	1.09x	1.21x	173	82%
Qwen3.8-Flash	6207	5579	5224	1.07x	1.19x	154	74%
Gemma-4-26B-A4B	5151	4757	4449	1.07x	1.16x	175	84%
MiniMax-M2.5	13055	11665	10739	1.09x	1.22x	173	82%
GLM-4.7	20671	18435	17002	1.08x	1.22x	182	87%
geomean				1.08x	1.19x

PR MFU: useful GEMM FLOPs (2 x T x top_k x 3 x hidden x moe_intermediate) over the whole layer time, against the bf16 mma.sync peak with fp32 accumulate at the rated boost clock: RTX PRO 6000 188 SMs x 2.43 GHz x 1024 FLOP/clk = 468 TFLOP/s, RTX 5090 170 SMs x 2.41 GHz x 512 FLOP/clk = 210 TFLOP/s (its cuBLAS bf16 GEMM reaches 219 - 227 TFLOP/s, so the card clocks above the rated boost). The RTX PRO 6000 runs BLOCK_SIZE_M 128 tiles here and the RTX 5090 BLOCK_SIZE_M 64 (GeForce cards run fp32-accumulate mma at half rate, so the decode the warp specialization hides is a smaller share).

On sm_120 the output differs only where the Gluon kernel runs (it accumulates in a different order). Decode (fused_experts_decode_nvfp4_marlin, M = 1 - 16, CUDA graph replay) is 1.00x and bitwise equal on all four GPUs.

End to end, RTX PRO 6000, nvidia/Qwen3.6-35B-A3B-NVFP4, time to first token of fresh random-word prompts (max_tokens 1), median over 2 server starts per tree (3 prompts each; 6 each for ~1k / ~2k):

prompt tokens	main (6b8f200)	PR
~1k	54.1 ms	57.8 ms	0.94x
~2k	79.4 ms	69.9 ms	1.14x
~4k	137.2 ms	124.1 ms	1.11x
~8k	267.2 ms	237.8 ms	1.12x
~16k	556.9 ms	499.0 ms	1.12x
~32k	1232 ms	1115 ms	1.11x

ft serve --model <path>/Qwen3.6-35B-A3B-NVFP4 --text-model-only --moe-strategy fused --cache-type naive

A greedy 128-token answer is identical across server starts for each tree; main and PR diverge after ~375 characters (one word choice, both coherent).

Host time: on sm_120 the PR adds host time per MoE layer: ~40 us for the TMA descriptors on the Triton path (BLOCK_SIZE_M 16 / 32) and ~55 us for the Gluon launch. A ~1k-token prefill of Qwen3.6-35B-A3B is host-bound, so it gets slower even though its MoE GPU time drops from 21.4 to 16.5 ms (nsys); with TMA off it is within 1% of main. On a busy host, small layers (Qwen3.6-35B-A3B shapes, T <= 512) also measure host-bound in the kernel A/B on the RTX 5090 now and then.

Tests

pytest tests/moe/test_nvfp4_backends.py tests/kernels/test_e4m3_compat.py tests/moe/test_fused_moe.py, with Gluon on and with FREETOKEN_DISABLE_GLUON=1:

• RTX PRO 6000 / RTX 5090: 56 passed, 5 skipped / 52 passed, 9 skipped
• RTX 3080 Ti: 46 passed, 15 skipped (both)
• RTX 4090: 49 passed, 11 skipped, 1 failed (both): test_e4m3_compat.py::test_forced_emu_matches_native (blk_aq_y) fails the same way on main on this GPU.

Related: #578 also replaces the table decode of the Triton prefill kernel; it predates #723.