Skip to content

feat: add NVFP4 MoE kernels with optional 4over6 quantization - #12

Draft
S1ro1 wants to merge 6 commits into
feat/mxfp8-moe-runtimefrom
feat/nvfp4-moe
Draft

feat: add NVFP4 MoE kernels with optional 4over6 quantization#12
S1ro1 wants to merge 6 commits into
feat/mxfp8-moe-runtimefrom
feat/nvfp4-moe

Conversation

@S1ro1

@S1ro1 S1ro1 commented Sep 10, 2026

Copy link
Copy Markdown

Add the SM100 nvfp4_moe kernel package for routed-expert training, exposed through the kernel registry and prime_rl::grouped_nvfp4_gemm. The implementation packs weights with per-expert global scales and activations with per-token global scales, runs NVFP4 grouped GEMM, and supports dequant_bf16 backward using the packed forward operands or bf16 backward using the master operands.

NVFP4GroupedGemm(four_over_six=True) enables adaptive 4/6 quantization for weights and both activation inputs. It matches FlashInfer's default 4/6 recipe: 448 normalization, MAE, strict reconstruction-error scoring, and default approximate candidate arithmetic. The option defaults to False. The wrapper remains traceable with torch.compile(fullgraph=False).

Declare the CUTLASS build dependency and header paths in kernels.toml, resolving build dependencies per selected kernel. Include the source license, package documentation, and focused numerical tests.

Stacked on #9 with gh stack: main → #9#12. This PR contains only the NVFP4 changes. The tested source commit 436719d and release wheel are unchanged. Prime-RL integration: PrimeIntellect-ai/prime-rl#3511.

Validation on GB200 with CUDA 13, torch 2.13.0 and FlashInfer 0.6.18, using the locally built release wheel:

  • Eight GPU tests passed for 4/6 on/off, grouped forward, both backward policies, token locality, zero rows, and empty inputs.
  • Fifteen comparisons against FlashInfer/vLLM default 4/6 matched packed values, logical block scales, and global scales exactly, covering contraction sizes 64–7168, irregular token groups, empty experts, and tiny blocks.
  • Compiled Prime-RL backend forward and dequantized backward matched eager exactly. A configured MoE produced finite outputs and non-zero router and expert gradients.
  • A two-step Qwen3-30B-A3B math RL smoke completed on 4 trainer + 4 inference GB200 GPUs, with the first 40/48 layers quantized and 4/6 enabled on both sides. Mismatch KL was 0.006612 and 0.005918; the previous 85% run measured 0.007310 and 0.006881. Sampled batches differ. Both smoke batches had zero task reward, so this validates execution and consistency, not learning quality.
  • Lint, formatting, and whitespace checks passed. The wheel was built locally, without GitHub Actions.

@S1ro1
S1ro1 changed the base branch from main to feat/mxfp8-moe-runtime September 10, 2026 13:35
@S1ro1
S1ro1 added this pull request to stack #13 September 10, 2026 13:36
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant