[jason-fxz]

Prefill path of the NVFP4 MoE on sm_120 for checkpoints whose experts carry static activation scales (ModelOpt NVFP4, compressed-tensors NVFP4 with dynamic: local inputs): both GEMMs run on NVFP4 activations (W4A4). Same banks as the Triton kernels, no repack; decode, small prefills and W4A16 checkpoints are unchanged.

Changes

• New NVFP4 expert kernel w4a4 (layers/quantization/moe/nvfp4.py), first in the table on sm_120 with Gluon, for checkpoints whose experts carry input_scale. It keeps the Triton kernel's banks, decode and CPU format; prefill runs W4A4 from 4 tokens per expert and main's W4A16 below. --quant-backend moe.nvfp4=triton keeps W4A16.
• kernel/gluon/nvfp4/moe/sm120_w4a4.py: two warp-specialized Gluon GEMMs on mma.sync kind::mxf4nvf4. gate_up applies the gated activation and quantizes its output to NVFP4 in the epilogue; down applies the routing weight and sums the top-k routes with bf16 atomics (add order varies run to run). No bf16 intermediates, no act / moe_sum kernels. Below I = 1024 the down GEMM runs persistently with separate epilogue warps.
• Static activation scales: the NVFP4 expert reader reads each expert's input_scale (compressed-tensors: input_global_scale), kept as GPU-resident alphas; each GEMM input uses the max over the layer's experts, as vLLM does.
• Dispatch (W4A4 vs W4A16, tiles) in kernel/nvfp4_moe.py; a test of the w4a4 kernel against a torch emulation of the NVFP4 activation rounding.

Results

Measured against main at ac6a15d. torch 2.14.1+cu130, Triton 3.8.0.

Kernel: the whole MoE layer, main fused_experts_nvfp4 vs PR fused_experts_nvfp4_w4a4, 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, 32, 64, 128, 256, 512, 768, 1024, 1536, 2048, 3072, 4096}, static input scales from a calibration batch, each model's activation; triton.testing.do_bench median. Speedup over main, geomean (min - max), by tokens per expert:

GPU (driver)	all	< 4	4 - 32	32 - 128	>= 128
RTX PRO 6000, sm_120 (595.91)	1.25x (1.00 - 2.96)	1.00x	1.05x	1.45x	2.54x
RTX 5090, sm_120 (580.105)	1.55x (0.70 - 4.54)	1.00x	1.11x	2.37x	3.72x

Below 4 tokens per expert the PR runs main's path. The RTX 5090 minimum is a host-bound point of a small layer (Qwen3.6-35B-A3B / Gemma-4-26B-A4B shapes, T <= 256; the busy host adds ~130 us now and then); the other shapes there are 1.00 - 1.04x.

W4A4 at T = 8192 (main / PR is the speedup):

RTX PRO 6000 (the kernels of this PR before the port into the tree; they match the PR output to 1e-3):

model	main, us	PR, us	PR / main	PR TFLOP/s	PR MFU
Qwen3.6-35B-A3B	1965	677	2.90x	609	33%
Qwen3.5-122B-A10B	4875	1623	3.00x	762	41%
Qwen3.8-Flash	4551	1478	3.08x	545	29%
Gemma-4-26B-A4B	3320	1080	3.07x	722	39%
MiniMax-M2.5	7116	2250	3.16x	825	44%
GLM-4.7	10877	3341	3.26x	926	49%
GLM-5.2	17263	5584	3.09x	886	47%
GLM-5.3-Flash	11215	3750	2.99x	880	47%
MiniMax-M3	13314	4055	3.28x	915	49%
geomean			3.09x

RTX 5090:

model	main, us	PR, us	PR / main	PR TFLOP/s	PR MFU
Qwen3.6-35B-A3B	2615	679	3.85x	607	36%
Qwen3.5-122B-A10B	7067	1680	4.21x	736	44%
Qwen3.8-Flash	5188	1527	3.40x	527	31%
Gemma-4-26B-A4B	4465	1097	4.07x	711	42%
MiniMax-M2.5	10628	2361	4.50x	786	47%
GLM-4.7	17049	3648	4.67x	848	51%
GLM-5.2	27674	5971	4.63x	829	49%
GLM-5.3-Flash	18955	3814	4.97x	865	52%
MiniMax-M3	20457	4261	4.80x	871	52%
geomean			4.32x

At T = 32768 the geomean is 3.06x on the RTX PRO 6000 (MFU 38 - 57%) and 4.63x on the RTX 5090 (42 - 57%).

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 dense FP4 mma.sync peak at the rated boost clock: RTX PRO 6000 188 SMs x 2.43 GHz x 4096 FLOP/clk = 1871 TFLOP/s, RTX 5090 170 SMs x 2.41 GHz x 4096 FLOP/clk = 1676 TFLOP/s. An MMA-only loop reaches 1781 / 2007 TFLOP/s: under this load the RTX PRO 6000 runs ~2.31 GHz and the RTX 5090 ~2.88 GHz, above its rated boost. FP4 runs at full rate with fp32 accumulate on GeForce too, while bf16 does not, so main's W4A16 GEMM loses more on the RTX 5090.

vs FlashInfer b12x W4A4 (flashinfer 0.6.18 b12x_fused_moe(quant_mode="nvfp4"), weights and swizzled scales prepared once, the same static input scales and activations, whole layer): b12x time / ours, the kernels of this PR before the port (they match the PR output to 1e-3):

model	PRO T=8192	16384	32768	5090 T=8192	16384	32768
Qwen3.6-35B-A3B	1.04x	1.03x	1.05x	0.97x	0.96x	1.03x
Qwen3.5-122B-A10B	1.67x	1.68x	1.63x	1.58x	1.63x	1.68x
Qwen3.8-Flash	1.57x	1.67x	1.56x	1.49x	1.60x	1.57x
Gemma-4-26B-A4B	1.69x	1.71x	1.64x	1.58x	1.62x	1.63x
MiniMax-M2.5	1.71x	1.77x	1.71x	1.62x	1.70x	1.67x
GLM-4.7	1.80x	1.80x	1.75x	1.69x	1.78x	1.67x
GLM-5.2	1.73x	1.79x	1.71x	1.65x	1.66x	1.64x
GLM-5.3-Flash	1.74x	1.85x	1.79x	1.83x	1.74x	1.64x
MiniMax-M3	1.77x	1.86x	1.84x	1.76x	1.72x	1.67x
geomean	1.63x			1.57x

b12x's gated kernel (I <= 512) keeps Qwen3.6-35B-A3B level; on the other shapes our GEMMs run at 1.5 - 1.9x its rate.

Decode (fused_experts_decode_nvfp4_marlin) and W4A16 prefill are untouched.

Accuracy: W4A4 rounds the activations to 4 bits, the precision the checkpoint was calibrated for, so outputs differ from W4A16 by more than rounding noise. RadixArk Qwen3.8-Flash-Next-NVFP4 (offload, RTX PRO 6000), AIME25 pass@3 on the four problems whose prompts reach 4 tokens per expert (10, 11, 20, 21): 3 / 4 pass; problem 20 ran all three samples into the 16384-token limit without an answer.

Tests

pytest tests/moe/test_nvfp4_backends.py on the RTX PRO 6000: 15 passed, 4 skipped (Marlin / vLLM not installed).

[jason-fxz]

Known issues, not addressed in this PR

Experts of several formats fail to load on sm_120

On sm_120 with Gluon and the automatic kernel choice, the NVFP4 layers of a checkpoint whose experts declare input_scale bind w4a4, the first entry of the NVFP4 MoE table. w4a4 keeps the per-expert input scales as GPU-resident alphas (gate_up_alpha, down_alpha). The paths for experts of several formats (for example bf16 routed experts in some layers and NVFP4 in the rest) refuse any kernel with GPU-resident alphas:

• merge_expert_banks in moe/expert_banks.py (offload, fused, and ft checkpoint)
• OffloadMoeCache.__post_init__ with layer_layouts in moe/offload_cache.py

So such a checkpoint stops at startup, where main binds triton on those layers and serves it. Other GPUs are not affected: there w4a4 is not worth it and the table still picks triton.

Reproduced on CPU with the bf16 + NVFP4 model of tests/moe/test_mixed_expert_formats.py, the quant backend left on auto and W4a4Nvfp4MoEKernel.worth_it forced to True (what sm_120 reports):

kernels: ['fused', 'w4a4', 'w4a4', 'w4a4', 'fused']
ValueError: experts of several formats cannot use the w4a4 kernel, which keeps GPU-resident alphas

Workaround: --quant-backend moe.nvfp4=triton.

FTW checkpoints keep the W4A16 prefill

w4a4 writes the same nvfp4 bank format as triton, and an FTW nvfp4 file reads back as triton (kind_kernel_for in moe/legacy_format.py, applied by _adjust_ftw_quant_backend). An FTW checkpoint therefore serves the W4A16 prefill on sm_120 too. An FTW converted on sm_120 with the automatic choice also stores the input scales as gate_up_alpha / down_alpha entries, which the triton kernel does not read.