perf(kernel): rewrite the triton nvfp4 moe prefill gemm

[jason-fxz]

Summary

NVFP4 MoE prefill goes TritonNvfp4MoEKernel.apply (the auto choice) -> fused_experts_nvfp4 -> _prefill_nvfp4_moe_kernel, in both --moe-strategy fused and offload. On H100 that kernel ran at 16 to 18 TFLOP/s on long prompts, so prefill of Gemma 4 26B-A4B NVFP4 stayed near 5k tok/s.

This PR rewrites the kernel. It is 6.6x to 16.1x faster across six NVFP4 MoE shapes at 128 to 32k tokens, with the same error. Gemma 4 prefill goes from 5,571 to 35,732 tok/s at 4.6k tokens with fused, and from 4,783 to 16,400 tok/s at 37.8k tokens. Decode does not use this kernel and is unchanged.

Changes

kernel/triton/nvfp4_fused_moe.py, _prefill_nvfp4_moe_kernel:

• Operand order: computes out^T = W @ X^T. The decoded weights become the dot's A operand, which Hopper's wgmma reads from registers (the B operand would need shared memory), and the token block becomes the MMA N dimension.
• e2m1 decode: bit arithmetic (_e2m1_f16_x2p14) instead of a lookup-table gather. The decoded value times its e4m3 block scale is exact in bf16 and fp16.
• One dot per K step: the low and high nibbles are interleaved back to K order, so one tl.dot covers the step. The old kernel issued two dots on even and odd activation columns loaded with stride 2.
• Scales: one e4m3 load per 16 values instead of one per packed byte (two values). The native and emulated e4m3 paths (e4m3_native_cx()) are kept.

moe/fused_nvfp4.py:

• _prefill_config(M, top_k, num_experts) picks BLOCK_SIZE_M from the mean tokens per expert: 128 at 64 or more, 64 at 24, 32 at 12, else 16. It uses BLOCK_N 64, 4 warps and 3 stages.
• The kernel no longer takes the e2m1 lookup table, so _prefill_gemm stops passing it.

Tested on

NVIDIA H100 80GB HBM3, Intel Xeon Platinum 8480+, driver 580.95.05, CUDA 13.0, torch 2.14.1+cu130, Triton 3.8.0.

End-to-end results

Model: nvidia/Gemma-4-26B-A4B-NVFP4.

python -m freetoken --model-path <model> --port 8137 --served-model-name default --text-model-only \
  --max-running-requests 1 --cuda-graph-max-bs 1 --max-seq-len-override 40960 --cache-type naive \
  --moe-strategy fused            # or: --moe-strategy offload --moe-cache-rate 0.5

• Prefill: one streamed /v1/chat/completions request with a 4,582-token or 37,811-token prompt and 8 output tokens. Prefill tok/s is prompt tokens / TTFT. Median of 3 runs.
• Decode: AIME25 problem 1, temperature 1.0, top_p 0.95, top_k 64, to the natural end. Median of 5 runs.

Strategy	Prompt	main tok/s (TTFT)	this PR tok/s (TTFT)
fused	4,582	5,571 (0.82 s)	35,732 (0.13 s)
fused	37,811	4,783 (7.91 s)	16,400 (2.31 s)
offload, cache rate 0.5	4,582	5,506 (0.83 s)	17,922 (0.26 s)
offload, cache rate 0.5	37,811	4,754 (7.95 s)	14,437 (2.62 s)

Strategy	main decode tok/s	this PR decode tok/s
fused	195.7	195.0
offload, cache rate 0.5	177.9	177.3

With greedy sampling and a 60-token prompt, the first 128 output tokens match main token for token, with both fused and offload.

Kernel results

fused_experts_nvfp4 alone with random weights and routing, bf16 activations, one CUDA graph replay after an L2 flush, median ms (main -> this PR).

Shape (E, top_k, H, I)	T = 128	T = 1024	T = 4096	T = 8192	T = 16384	T = 32768
Qwen3.6-35B-A3B (256, 8, 2048, 512)	3.12 -> 0.42	4.51 -> 0.59	13.17 -> 1.18	24.79 -> 1.99	48.06 -> 3.52	94.32 -> 6.66
Qwen3.5-122B-A10B (256, 8, 3072, 1024)	8.68 -> 1.15	12.89 -> 1.57	38.43 -> 3.08	72.48 -> 5.13	140.63 -> 9.25	276.23 -> 17.94
Qwen3.8-Flash-Next (512, 10, 2560, 640)	8.21 -> 1.13	9.01 -> 1.36	26.70 -> 2.31	48.73 -> 4.44	92.86 -> 6.79	181.28 -> 12.62
Gemma 4 26B-A4B (128, 8, 2816, 704)	2.77 -> 0.40	6.65 -> 0.71	22.58 -> 1.71	43.57 -> 3.01	85.72 -> 5.79	169.71 -> 11.33
MiniMax-M2.5 (256, 8, 3072, 1536)	13.01 -> 1.70	19.34 -> 2.34	57.67 -> 4.57	108.78 -> 7.57	211.05 -> 13.88	414.56 -> 30.42
GLM-4.7 (160, 8, 5120, 1536)	13.55 -> 1.78	27.35 -> 2.47	91.76 -> 6.41	175.93 -> 12.51	345.20 -> 21.95	683.33 -> 42.35

• At T = 32768 this PR runs at 244 to 292 TFLOP/s; main runs at 17 to 18.
• At T = 128 and 1024, the relative error against an fp32 reference over the dequantized weights is the same on both sides, 0.0038 to 0.0039.
• BLOCK_SIZE_M comes from an H100 sweep over 96 configurations (BLOCK_M, BLOCK_N, BLOCK_KB, warps, stages) on these six shapes at T = 256, 1024, 4096 and 16384. At every point, the best configuration uses the same BLOCK_M as the heuristic, and the heuristic is within 3.3% of the best.

Other checks

• PYTHONPATH=python pytest tests/moe/test_nvfp4_backends.py tests/kernels/test_e4m3_compat.py: 13 passed, 7 skipped (vllm not installed, b12x needs sm_120, ROCm only). The slow compile gates for sm_80, 86, 89 and 120 in test_e4m3_compat.py are included.
• Compile only (no run), every BLOCK_M for gate_up and down with bf16 activations: sm_80, 86, 89 and 120 use mma.sync, sm_90 wgmma and sm_100 tcgen05. Shared memory peaks at 37,120 bytes on sm_80/86/89/120, 53,504 bytes on sm_90 and 53,520 bytes on sm_100, under each arch's limit.
• No GPU other than H100 was measured.

[jason-fxz]

RTX results and per-GPU tuning (c4acce3)

The kernel table and the tile sweep were rerun on four RTX GPUs. The kernel is a clear win on all of them, but _prefill_config (from the H100 sweep) is 22% to 42% slower than the best tile on average. c4acce3 picks the tiles per GPU class.

Tested on

GPU	cc	SMs	Driver
RTX PRO 6000 Blackwell Server Edition	12.0	188	595.91
GeForce RTX 5090	12.0	170	595.71
GeForce RTX 4090	8.9	128	580.76
GeForce RTX 3080 Ti	8.6	80	595.71

torch 2.14.1+cu130, Triton 3.8.0. Same method as the PR: fused_experts_nvfp4 with random weights and routing, bf16 activations, one CUDA graph replay after an L2 flush, median ms, two rounds with the order flipped.

main -> this PR (6b4086c), kernel table shapes, T = 128 to 32768

GPU	Speedup	PR GEMMs at T = 32768	cuBLAS bf16 8192^3
RTX PRO 6000	2.4x to 6.4x	209 to 218 TFLOP/s (48% to 50% of cuBLAS)	434 TFLOP/s
RTX 5090	2.6x to 4.6x	155 to 161 TFLOP/s (68% to 70%)	228 TFLOP/s
RTX 4090	1.8x to 4.5x	108 to 115 TFLOP/s (67% to 72%)	161 TFLOP/s
RTX 3080 Ti	2.7x to 5.6x	40 to 43 TFLOP/s (66% to 70%)	61 TFLOP/s

The relative error against the fp32 reference is the same on both sides on every GPU (0.0039). The 3080 Ti runs the emulated e4m3 path.

Why RTX needs other tiles

GeForce runs bf16 MMA with fp32 accumulate at half rate. Per SM per clock, bf16 tensor over fp32 ALU throughput (rho) is 16 on H100, 4 on RTX PRO / RTX Ada / A10 / L40S, and 2 on GeForce RTX 30/40/50 (NVIDIA whitepapers; measured 4.5 / 2.3 / 1.9 / 1.6 on the four GPUs above). Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M. With a low rho the padded MMA costs more than the extra decode, so RTX GPUs want a smaller BLOCK_SIZE_M than H100. Their best BLOCK_SIZE_KB and BLOCK_SIZE_N also differ: KB 64 and N 128 at small blocks.

Tile sweep

144 configs (BLOCK_M 16..128, BLOCK_N 32..128, BLOCK_KB 32/64, 4/8 warps, 2..4 stages) on the six shapes at T = 256, 512, 1024, 2048, 4096 and 16384. Loss against the best config at each point, mean (max):

GPU	H100 _prefill_config	c4acce3
RTX PRO 6000	+24.4% (+90%)	+1.8% (+12%)
RTX 5090	+22.0% (+89%)	+2.3% (+16%)
RTX 4090	+42.3% (+122%)	+1.9% (+8%)
RTX 3080 Ti	+24.7% (+61%)	+3.5% (+13%)

Limiting the SM count with a green context (24 to 188 SMs on the RTX PRO 6000) does not change the best tile, so the SM count is not an input.

c4acce3

• rho comes from the compute capability and the device name: GeForce RTX 30/40/50 -> 2, RTX A / RTX Ada / RTX PRO / A10 / A40 / L4 / L40S -> 4, A100 / H100 -> 16, H20 -> 3.4. The name patterns were checked against every product name in the 580.76, 580.95 and 595.71 driver lists, plus names without "GeForce" ("NVIDIA RTX 4090") and vGPU names. An unknown name logs a warning and uses 4.
• rho sets the BLOCK_SIZE_M thresholds on tokens per expert: (12, 24, 64) for rho >= 8 (unchanged), (16, 36, 150) for RTX PRO, (16, 150, never 128) for GeForce.
• cc 8.6 / 8.9 / 12.0 use the RTX PRO or GeForce tile table. H100, H20 and A100 keep the PR's tiles. AMD keeps the PR's config (TODO).

this PR -> c4acce3

Geometric mean over the six shapes at T = 128, 256, 512, 1024, 2048, 4096 and 16384:

GPU	All	T <= 1024	T >= 2048	Min	Max
RTX PRO 6000	1.17x	1.26x	1.06x	0.91x	1.91x
RTX 5090	1.15x	1.19x	1.11x	0.88x	1.84x
RTX 4090	1.33x	1.45x	1.18x	1.04x	2.15x
RTX 3080 Ti	1.15x	1.14x	1.16x	0.93x	1.58x

The RTX PRO 6000 numbers use the second round only; another process used the GPU during the first round. The slower points depend on the model shape: GLM-4.7 at T = 2048 (best BLOCK_SIZE_M is 128 there; 0.88x on the 5090, 0.92x on the RTX PRO 6000), Qwen3.8-Flash-Next at T = 4096 on the RTX PRO 6000 (0.91x), and Qwen3.6-35B-A3B at T <= 256 on the 3080 Ti (0.93x).

Other checks

• pytest tests/moe/test_nvfp4_backends.py tests/kernels/test_e4m3_compat.py on the RTX PRO 6000: 15 passed, 5 skipped (vllm not installed, ROCm only).
• H100, H20 and A100 were not measured. On H100 the config is the same as the PR's.