feat: add NVFP4 MoE training with optional 4over6 quantization

docs/advanced.md

@@ -54,6 +54,14 @@ Dense linear precision and routed-expert precision are configured independently.
 - `type = "bf16"` (default)
 - `type = "deepgemm_fp8"` (requires DeepGEMM and SM90+)
 - `type = "mxfp8"` (requires `prime-kernels`, torchao, and SM100)
+- `type = "nvfp4"` (requires a `prime-kernels` build containing `nvfp4_moe` and SM100). `backward = "dequant_bf16"` uses reconstructed forward operands for BF16 gradient GEMMs; `backward = "bf16"` uses the original BF16 operands. Gated experts require the `gate_up` fusion. LoRA and full-graph compilation are unsupported; `compile.fullgraph` defaults to `false`.
+
+For NVFP4, `four_over_six = true` enables adaptive 4/6 quantization for expert
+weights and activations. It defaults to `false`. The enabled recipe matches
+FlashInfer's defaults: 448 normalization, MAE selection, strict error scoring,
+and default candidate arithmetic. Enable the matching sampler recipe with
+`FLASHINFER_NVFP4_4OVER6 = "1"` in `[inference.env_vars]`, leaving the other
+FlashInfer NVFP4 settings at their defaults.
 
 ```toml
 [trainer.model.quantization]
@@ -81,13 +89,23 @@ For example, this selects routed experts in Qwen3's first four model layers:
 
 ```toml
 [trainer.model.moe.compute]
-type = "mxfp8"
+type = "nvfp4"
 apply_to = [0, 1, 2, 3]
 ```
 
 Backend shape checks and token alignment apply only to the selected compute path.
 
-In RL runs, configure the same precision selection for rollouts. Inference module names can differ from the trainer's names, and inference precision is configured explicitly, not inferred from `apply_to`. Check the selected modules on both sides before comparing trainer and rollout logprobs.
+Configure the same precision selection for rollouts using vLLM's native online quantization targets. For the Qwen3 example above:
+
+```toml
+[inference.vllm]
+quantization = "online"
+
+[inference.vllm.quantization_config.targets]
+"model.layers.[0-3].mlp.experts" = "nvfp4_per_token"
+```
+
+Inference targets use vLLM's module names, which can differ from the trainer's names on other architectures. They are configured explicitly, not inferred from `apply_to`. Check the selected modules on both sides before comparing trainer and rollout logprobs. Online quantization targets preserve the selection when weights are reloaded.
 
 GLM-5.2 adds IndexShare: the DSA sparse-attention indexer runs only on a subset of layers and the remaining layers reuse the cached top-k indices. The trainer reads this schedule from the model's `indexer_types` config field and enables the index cache automatically, so no extra config is needed. To override the schedule manually, set `[trainer.model.index_cache]` (`topk_freq` or `topk_pattern`).
 

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -228,8 +228,18 @@ class MXFP8MoEComputeConfig(MoEComputeConfigBase):
     """MXFP8 expert-compute recipe."""
 
 
+class NVFP4MoEComputeConfig(MoEComputeConfigBase):
+    """Run routed-expert forwards in NVFP4 with BF16 gradient GEMMs."""
+
+    type: Literal["nvfp4"] = "nvfp4"
+    backward: Literal["dequant_bf16", "bf16"] = "dequant_bf16"
+    """Use dequantized forward operands or the original BF16 operands for backward."""
+    four_over_six: bool = False
+    """Use FlashInfer's default 4/6 recipe: 448 normalization and MAE scoring."""
+
+
 MoEComputeConfig: TypeAlias = Annotated[
-    BF16MoEComputeConfig | DeepGemmFP8MoEComputeConfig | MXFP8MoEComputeConfig,
+    BF16MoEComputeConfig | DeepGemmFP8MoEComputeConfig | MXFP8MoEComputeConfig | NVFP4MoEComputeConfig,
     Field(discriminator="type"),
 ]
 
@@ -408,6 +418,11 @@ def quantization_only_with_custom_impl(self):
 
     @model_validator(mode="after")
     def validate_moe_runtime(self):
+        if isinstance(self.moe.compute, NVFP4MoEComputeConfig):
+            if self.lora is not None:
+                raise ValueError("NVFP4 expert compute does not support LoRA.")
+            if self.compile is not None and self.compile.fullgraph:
+                raise ValueError("NVFP4 expert compute requires compile.fullgraph=false.")
         if self.ep == 1:
             return self
 

pyproject.toml

@@ -98,7 +98,7 @@ quack = [
     "quack-kernels>=0.4.1 ; sys_platform == 'linux'",
 ]
 # Prebuilt `prime-kernels` wheels from a release (see the `kernels` skill). Compiled for
-# sm_100a; on other GPUs the package installs fine and reports its kernels unavailable.
+# SM100a and SM103a where supported by each kernel.
 kernels = [
     "prime-kernels ; sys_platform == 'linux'",
 ]
@@ -289,7 +289,7 @@ vllm-router = [
 ]
 vllm = [
     { url = "https://github.com/S1ro1/vllm/releases/download/nvfp4-hisparse-b300-20260921-90836b3/vllm-0.28.1.dev0%2Bg90836b3.b300.hisparse1-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://wheels.vllm.ai/c7e9816c6ab0731165a134fd0a9defed6ab1d748/vllm-0.28.1rc1.dev588%2Bgc7e9816c6-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/S1ro1/vllm/releases/download/nvfp4-20260910-bc4755a/vllm-0.28.1rc1.dev588%2Bnvfp4.gbc4755a.cu130-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-ep = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
@@ -300,8 +300,8 @@ deep-gemm = [
     { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 prime-kernels = [
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/v0.1.1/prime_kernels-0.1.0+cu130torch2.13.0-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/nvfp4-b300-20260917-2802a7a/prime_kernels-0.1.0%2Bcu130torch2.13.0.b300.1-cp312-cp312-linux_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://github.com/PrimeIntellect-ai/prime-kernels/releases/download/nvfp4-20260910-436719d/prime_kernels-0.1.0%2Bcu130torch2.13.0.nvfp4.g436719d-cp312-cp312-linux_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 flash-linear-attention = { git = "https://github.com/fla-org/flash-linear-attention" }
 mamba-ssm = { git = "https://github.com/state-spaces/mamba", rev = "f577286" }

src/prime_rl/trainer/activation_checkpointing.py

@@ -67,6 +67,7 @@
         "prime_rl::dsv4_sparse_attn",
         "prime_rl::fp8_blockwise_mm",
         "prime_rl::grouped_fp8_gemm",
+        "prime_rl::grouped_nvfp4_gemm",
         "prime_rl::sparse_mla",
     }
 )

src/prime_rl/trainer/model.py

@@ -910,14 +910,14 @@ def _validate_flash_attn_4_installed() -> None:
 def resolve_auto_attn(config: ModelConfig) -> None:
     """Resolve ``attn='auto'`` to a concrete flash attention implementation based on GPU architecture.
 
-    FA4 on datacenter Blackwell (SM100), FA3 on Hopper (SM90), FA2 otherwise.
+    FA4 on datacenter Blackwell (SM100/SM103), FA3 on Hopper (SM90), FA2 otherwise.
     Workstation Blackwell GPUs (e.g. RTX PRO 6000, SM120) lack FA4 kernels and
     can't run the Hopper-only FA3 kernels, so they fall back to FA2.
     """
     if config.attn != "auto":
         return
     major, minor = torch.cuda.get_device_capability()
-    if (major, minor) == (10, 0):
+    if (major, minor) in ((10, 0), (10, 3)):
         resolved = "flash_attention_4"
     elif major == 9:
         resolved = "flash_attention_3"

src/prime_rl/trainer/models/layers/grouped_gemm.py

@@ -66,3 +66,22 @@ def __call__(
             offs,
             high_precision_wgrad=self.high_precision_wgrad,
         )
+
+
+@dataclass(frozen=True)
+class NVFP4GroupedGemm:
+    kernel: ModuleType
+    backward: str
+    token_group_alignment: int
+    four_over_six: bool = False
+
+    def __call__(
+        self,
+        x: torch.Tensor,
+        weight_t: torch.Tensor,
+        *,
+        offs: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.kernel.grouped_gemm(
+            x, weight_t, offs=offs, backward=self.backward, four_over_six=self.four_over_six
+        )

src/prime_rl/trainer/moe_runtime.py

@@ -11,6 +11,7 @@
     ModelConfig,
     MoERuntimeConfig,
     MXFP8MoEComputeConfig,
+    NVFP4MoEComputeConfig,
     TorchMoEDispatchConfig,
 )
 from prime_rl.trainer.distributed.expert_parallel import ExpertWeightParallel
@@ -24,6 +25,7 @@
     DeepGemmFP8GroupedGemm,
     GroupedGemm,
     MXFP8GroupedGemm,
+    NVFP4GroupedGemm,
 )
 from prime_rl.trainer.models.layers.moe import MoE
 from prime_rl.trainer.parallel_dims import ParallelDims
@@ -53,6 +55,18 @@ def _resolve_grouped_gemm(config: ModelConfig) -> GroupedGemm:
             high_precision_wgrad=compute.recipe == "mxfp8_rceil_wgrad_with_hp",
             token_group_alignment=kernel.TOKEN_GROUP_ALIGNMENT,
         )
+    if isinstance(compute, NVFP4MoEComputeConfig):
+        import prime_kernels
+
+        if "nvfp4_moe" not in prime_kernels.KERNELS:
+            raise RuntimeError("NVFP4 requires a prime-kernels build containing nvfp4_moe.")
+        kernel = prime_kernels.load("nvfp4_moe")
+        return NVFP4GroupedGemm(
+            kernel=kernel,
+            backward=compute.backward,
+            token_group_alignment=kernel.TOKEN_GROUP_ALIGNMENT,
+            four_over_six=compute.four_over_six,
+        )
     raise TypeError(f"Unsupported MoE compute config: {type(compute).__name__}")
 
 
@@ -84,6 +98,14 @@ def configure_moe_runtime(model: nn.Module, config: ModelConfig, parallel_dims:
 
     for moe in moe_layers:
         grouped_gemm = selected_grouped_gemm if moe in selected_moes else bf16_grouped_gemm
+        if isinstance(grouped_gemm, NVFP4GroupedGemm):
+            if moe.experts.gate_proj is not None:
+                raise ValueError("NVFP4 gated experts require the gate_up runtime fusion for shared weight scales.")
+            reason = grouped_gemm.kernel.unsupported_shape_reason(
+                moe.experts.down_proj.shape[-2], moe.experts.hidden_dim
+            )
+            if reason is not None:
+                raise ValueError(reason)
         if ep_mesh is not None and moe.experts.num_experts % parallel_dims.ep:
             raise ValueError(
                 f"MoE expert count {moe.experts.num_experts} must be divisible by model.ep={parallel_dims.ep}."

tests/unit/test_configs.py

@@ -208,6 +208,10 @@ def test_moe_runtime_defaults_are_independent_from_dense_quantization():
     [
         {"moe": {"compute": {"type": "deepgemm_fp8"}}},
         {"moe": {"compute": {"type": "mxfp8", "recipe": "mxfp8_rceil_wgrad_with_hp"}}},
+        {"moe": {"compute": {"type": "nvfp4"}}},
+        {"moe": {"compute": {"type": "nvfp4", "backward": "bf16"}}},
+        {"moe": {"compute": {"type": "nvfp4", "four_over_six": True}}},
+        {"ep": 2, "moe": {"compute": {"type": "nvfp4"}, "dispatch": {"type": "deepep"}}},
         {"ep": 2, "moe": {"dispatch": {"type": "deepep", "num_sms": 16, "token_chunk_size": 1024}}},
         {
             "ep": 2,
@@ -223,7 +227,7 @@ def test_supported_moe_runtime_configs(model):
     assert config.moe.compute.resolve_layers(43) == set(range(43))
 
 
-@pytest.mark.parametrize("backend", ["bf16", "deepgemm_fp8", "mxfp8"])
+@pytest.mark.parametrize("backend", ["bf16", "deepgemm_fp8", "mxfp8", "nvfp4"])
 @pytest.mark.parametrize(
     ("selection", "num_layers", "selected"),
     [
@@ -284,6 +288,18 @@ def test_moe_compute_apply_to_cli(selection, selected):
             },
             "does not support DeepEP",
         ),
+        (
+            {"moe": {"compute": {"type": "nvfp4"}}, "lora": {}},
+            "does not support LoRA",
+        ),
+        (
+            {"moe": {"compute": {"type": "nvfp4"}}, "compile": {"fullgraph": True}},
+            "requires compile.fullgraph=false",
+        ),
+        (
+            {"moe": {"compute": {"type": "nvfp4", "backward": "fp4"}}},
+            "Input should be",
+        ),
     ],
 )
 def test_invalid_moe_runtime_configs_fail(model, message):

tests/unit/train/models/test_moe.py

@@ -39,14 +39,15 @@ def test_moe_detection_for_text_and_vlm(config_cls, expected):
     assert is_tt_moe_model(model) is expected
 
 
+@pytest.mark.parametrize("backend", ["mxfp8", "nvfp4"])
 @pytest.mark.parametrize("selection", [[], [0], "0%", "50%"])
-def test_unselected_moe_uses_bf16_without_loading_quantization_backend(selection):
+def test_unselected_moe_uses_bf16_without_loading_quantization_backend(backend, selection):
     moe = MoE.from_args(MoEArgs(num_experts=2), dim=4, hidden_dim=8, shared_expert=None)
     parameters = dict(moe.named_parameters())
     model = torch.nn.Module()
     model.model = torch.nn.Module()
     model.model.layers = torch.nn.ModuleList([torch.nn.Identity(), moe])
-    config = ModelConfig.model_validate({"moe": {"compute": {"type": "mxfp8", "apply_to": selection}}})
+    config = ModelConfig.model_validate({"moe": {"compute": {"type": backend, "apply_to": selection}}})
     dims = ParallelDims(dp_replicate=1, dp_shard=1, cp=1, pp=1, ep=1, world_size=1)
 
     configure_moe_runtime(model, config, dims)