[YevheniiKotyrlo]

Problem

The NVFP4 MoE prefill GEMM decodes every weight through the 16-entry table, decodes its e4m3 scale in fp32, and reads the activations twice at stride two, once for each nibble of a byte (_prefill_nvfp4_moe_kernel in kernel/triton/nvfp4_fused_moe.py):

        b_lo = tl.load(lut_ptr + lo) * scale  # [BLOCK_KB, BLOCK_N]
        b_hi = tl.load(lut_ptr + hi) * scale

        a_lo = tl.load(a_ptrs_lo, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)
        a_hi = tl.load(a_ptrs_hi, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)

At Qwen3.8-Flash-Next's prefill shapes (8,192 tokens, 512 experts, top-10) its gate_up GEMM runs at 8.6 TFLOPS on an RTX 3090 Ti, and it is most of a long prompt's compute: 8.66 s of the 11.98 s of compute in a 15,972-token prompt, traced with Nsight Systems. Every NVFP4 MoE prefill through the Triton grouped GEMM runs this kernel, offloaded or resident.

nvfp4_linear.py already decodes without the table: _nvfp4_pair_f16 places each code's bits in an fp16 field and the e4m3 scale is pre-multiplied by 128, and its docstring measures the table gather serializing the LSU. The MoE prefill kernel never took that path.

Solution

1. The prefill kernel decodes e2m1 with _nvfp4_pair_f16, once the byte's high nibble also sits at bits [19:16], and the scale with e4m3_to_f16_x128, a new e4m3_compat.py helper for the if e4m3_native_cx() branch nvfp4_linear.py inlined twice; the epilogue multiplies the global scale by 128. Every e2m1 x e4m3 product is exact in fp16 and in bf16, so each dot operand is the table's times exactly 2^-7, the epilogue's 128 restores it, and the output is bit-identical.
2. One contiguous activation tile per K step, split into the even and odd k in registers (tl.split(tl.reshape(a, (BLOCK_SIZE_M, BLOCK_SIZE_KB, 2)))), instead of two strided loads.
3. EVEN_K drops the byte and k masks when the byte tile divides K.
4. nvfp4_linear.py's two scale decodes call the helper; _nvfp4_gemm_kernel compiles to the same SASS as before in all 16 configurations I compiled (sm_86 and sm_89, BLOCK_M 16 and 64, split-K 1 and 2, even and ragged K).

Tests

• tests/moe/test_nvfp4_backends.py: every byte against every scale code through the kernel's own decode (each code's value times 2^-14, signs included, every finite product exact in bf16), and the prefill GEMM against the dequant reference past the small-batch tile, at a K the byte tile divides and at one it does not (272 and 80).
• tests/kernels/test_e4m3_compat.py: the cross-arch compile gate and the native-vs-emulated comparison also drive the dense GEMM and the MoE prefill at a ragged K. The gate's shapes had every K a multiple of the tiles, so it never compiled a ragged branch for sm_89. With an integer other in the fp8 scale load, as I first wrote it, the new cases fail for sm_89 and sm_120 with cannot cast int32[constexpr[32], constexpr[64]] to fp8e4nv and pass for sm_80 and sm_86.

A/B

RTX 3090 Ti 24 GB (PCIe 4.0 x16), Ryzen 9 9950X3D, driver 617.14, Windows 11, every arm over my three Windows fixes starting at #575.

fused_experts_nvfp4 called at Flash-Next's prefill shapes: 8,192 bf16 tokens, 512 experts of random NVFP4 weights (hidden 2,560, intermediate 640), random top-10 routing, one warm-up call and then the median of 7 calls timed with CUDA events, arms in the order main, branch, branch, main:

	main	this branch
per call	121.9 · 103.2 ms	18.3 · 18.7 ms
output SHA-256	3f995662625c2728	3f995662625c2728

End to end, main cannot hold Flash-Next's 1M window on a 24 GB card (#568), so the A/B runs on that setup with and without this change, arms before, after, after, before. Each prompt is a code review over TypeScript source cut to length, streamed with max_tokens 8 and timed to its first token; the 43,105-token one then gets a greedy 2,048-token answer:

arm	151 · 1,125 tokens	7,837	32,042	43,105	decode after the 43k prompt
without	2.62-2.70 · 2.62-2.67 s	6.02 · 6.14 s	24.80 · 24.95 s	34.4 · 34.7 s	29.12 · 29.10 tok/s
with	2.62 · 2.60-2.70 s	2.78 · 2.91 s	11.43 · 11.69 s	17.3 · 19.6 s	29.31 · 29.03 tok/s

On this card the old kernel took longer than a chunk's 2.6 s expert stream (6.0-6.1 s for the 7,837-token prompt, one chunk), so time to first token falls 1.8-2.2x on prompts of thousands of tokens. A prompt of about a thousand tokens or fewer is bound by that stream and does not move, and neither does decode. oscarlius's RTX 4090 run in this thread measures the MoE compute 4.2x faster with bit-identical output; there the old kernel already finished under the stream of a 4,096-token chunk, so only a single-pass prefill shows it in time to first token. These rows are from the first version of this change; the branch then took nvfp4_linear.py's helpers, 24 fewer SASS instructions on sm_86, and the kernel A/B above is the branch's.

The end-to-end server command

ft serve --model RadixArk/Qwen3.8-Flash-Next-NVFP4 --moe-strategy offload --text-model-only \
  --max-running-requests 1 --num-tokens 1048576 --kv-reserve-tokens 1048576 \
  --max-seq-len-override 1048576 --max-output-tokens 1048576 --memory-ratio 0.85 \
  --moe-cache-policy lrfu --moe-collect-stats --vram-reserve-mb 2000 \
  --kv-placement host --token-embedding-placement host --mamba-host-slots 128 \
  --pin-reasoning-effort xhigh --bind-when-ready \
  --hf-overrides '{"text_config":{"rope_parameters":{"rope_type":"yarn","rope_theta":10000000,"partial_rotary_factor":0.25,"mrope_interleaved":true,"mrope_section":[11,11,10],"factor":4,"original_max_position_embeddings":262144}}}'

--hf-overrides is #573 and --vram-reserve-mb #574; --kv-placement, --token-embedding-placement, --mamba-host-slots, --moe-cache-policy lrfu, --moe-collect-stats, --pin-reasoning-effort and --bind-when-ready come from my local build.

Verification

tests/moe/test_nvfp4_backends.py and tests/kernels/test_e4m3_compat.py: 12 passed, 11 skipped (marlin needs vllm, b12x needs sm_120, ROCm, native-fp8 references) on Linux over main (WSL2 Ubuntu 24.04, Python 3.12, torch 2.14.0, triton 3.8.0) and on Windows 11 over the Windows fixes (triton-windows 3.8.0), where main fails the new test_triton_prefill_dequant_reads_every_code_and_scale_exactly on both platforms; the compile gate passes for sm_80, 86, 89 and 120. tests/moe, tests/kernels and tests/models on Linux: 767 passed with this branch and 764 on main, with the same 3 failures on both: test_adjust_config_converts_moe_cache_rate_to_cache_size needs flashinfer, which that environment lacks, and test_dsv41_pack.py's two fp4_e4m3_b16 roundtrips come back 0.015625 off the reference quantizer on this sm_86 card.

No issue or PR covers this kernel's speed: searched the tracker for nvfp4 prefill, e2m1, LUT gather, fused_nvfp4, prefill_nvfp4 and prefill throughput.

Known limits

Untested on ROCm. #260 routes HIP to the serial kernel because this grouped kernel trips a memory-aperture violation on gfx1151; this change alters the kernel's loads, and I have no AMD card to check it on.

[oscarlius]

Validated this PR on an RTX 4090 (PCIe 4.0 x16) with Qwen3.8-Flash-Next NVFP4, cherry-picked onto a d3512b4-based fork.

Cold 32K prompt (32,012 tokens), prefilled in a single pass over the experts (we run a local layer-major prefill patch, so each layer's experts are streamed once):

	without #578	with #578
MoE compute, summed over the 48 layers	9.42 s	2.22 s (about 4.2×)
Expert H2D bytes per pass	68,136,468,480	68,136,468,480
Expert copy time	2.87 s	2.86 s
TTFT	13.28 s	6.13 s
Decode	37.10 tok/s	37.13 tok/s

The 128-token output on the same prompt was bit-identical to the unpatched build.

With the default 4096-token chunked prefill (8 passes for 32K), TTFT barely moved (about 21–25 s either way). Each chunk re-streams every expert (#579), so that path stays transfer-bound; the kernel speedup shows once prefill is down to one pass.

Thanks for the work!

[YevheniiKotyrlo]

Thanks for the 4090 run. The chunked result fits what I measured: on my 3090 Ti the old kernel took about 6 s for a 7.8k-token chunk, longer than the 2.6 s expert stream, so time to first token fell about 2x; on your 4090 it already finished under the stream, so only the single pass shows the gain.

[YevheniiKotyrlo]

#723 replaced the table decode, the per-element scale decode and the stride-2 activation loads this PR targeted, and #733 tuned the kernel further on RTX GPUs, so I'm closing this as superseded.