[2/3][sglang] feat: support mxfp8/mxfp4/nvfp4 weight refit

.github/workflows/e2e_ppo_trainer_megatron_sglang.yml

@@ -203,7 +203,7 @@ jobs:
       - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout
         run: |
           $UV_RUN ray stop --force
-          ROLLOUT_QUANTIZATION=FP8_BLOCK TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh
+          VERL_ROLLOUT_USE_DEEP_GEMM=0 ROLLOUT_QUANTIZATION=FP8_BLOCK TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh
       - name: clean up
         run: |
           rm -rf checkpoints

docs/low_precision/qat.md

@@ -1,6 +1,6 @@
 # Low-Precision Rollout and Quantization-Aware Training
 
-Last updated: 10/16/2026
+Last updated: 10/21/2026
 
 In RL training, rollout usually dominates the step time, and for large MoE models it is bound by memory capacity and bandwidth. Running rollout in low precision:
 
@@ -67,14 +67,14 @@ The trainer must be Megatron (FSDP is not yet supported for this feature). The r
 
 | Precision | Preset | QAT mode | GPU | vLLM Linear | vLLM MoE | SGLang Linear | SGLang MoE | vLLM status in verl | SGLang status in verl |
 |---|---|---|---|---|---|---|---|---|---|
-| FP8 block (E4M3, 128x128 weight blocks, dynamic 1x128 activations) | `FP8_BLOCK` | `fp8` | Hopper / Blackwell | `W8A8Fp8` | `W8A8Fp8MoE` | `W8A8Fp8` | `W8A8Fp8MoE` | ✅ Hopper | ✅ Hopper |
-| MXFP8 | `MXFP8` | `mxfp8` | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | ❌ (only whole-model `quant_method: mxfp8`) | ❌ | ✅ Blackwell | Blocked by engine support |
-| MXFP4 (MoE only) | `MXFP4` | `mxfp4` / `mxfp4_mxfp8` | Blackwell | ❌ | MXFP4×MXFP8<br>`(moe_backend=flashinfer_cutlass)` | ❌ | ❌ | ✅ Blackwell (MoE only) | Blocked by engine support |
-| NVFP4 (MoE only) | `NVFP4` | `nvfp4` | Blackwell | ❌ | W4A4<br>`(moe_backend=flashinfer_cutlass)` | ❌ | ❌ | ✅ Blackwell (MoE only) | Blocked by engine support |
+| FP8 block (E4M3, 128x128 weight blocks, dynamic 1x128 activations) | `FP8_BLOCK` | `fp8` | Hopper / Blackwell | `W8A8Fp8` | `W8A8Fp8MoE` | `W8A8Fp8`<br>`fp8_gemm_runner_backend=deep_gemm` | `W8A8Fp8MoE`<br>`moe_runner_backend=deep_gemm` | ✅ Hopper | ✅ Hopper |
+| MXFP8 | `MXFP8` | `mxfp8` | Blackwell | `W8A8Mxfp8` | `W8A8Mxfp8MoE` | native `quant_method=mxfp8`<br>`fp8_gemm_runner_backend=deep_gemm` | native `quant_method=mxfp8`<br>`moe_runner_backend=deep_gemm` | ✅ Blackwell | ✅ Blackwell |
+| MXFP4 (MoE only) | `MXFP4` | `mxfp4` / `mxfp4_mxfp8` | Blackwell | ❌ | MXFP4×MXFP8<br>`(moe_backend=flashinfer_cutlass)` | ❌ (BF16 dense) | MXFP4×MXFP8<br>`moe_runner_backend=deep_gemm` | ✅ Blackwell (MoE only) | ✅ Blackwell (MoE only) |
+| NVFP4 (MoE only) | `NVFP4` | `nvfp4` | Blackwell | ❌ | W4A4<br>`(moe_backend=flashinfer_cutlass)` | ❌ (BF16 dense) | W4A4<br>`auto` selects FlashInfer CUTLASS | ✅ Blackwell (MoE only) | ✅ Blackwell (MoE only) |
 
-MXFP8, MXFP4 and NVFP4 need Blackwell (SM100). On Hopper, vLLM falls back to Marlin weight-only kernels for these formats, whose numerics differ from Blackwell W8A8 / W4A4; verl does not support that path.
+MXFP8, MXFP4 and NVFP4 need Blackwell (SM100) for the validated paths above. On Hopper, vLLM may fall back to Marlin weight-only kernels, whose numerics differ from Blackwell W8A8 / W4A4; those paths are outside this support and validation matrix.
 
-MXFP4 and NVFP4 support is scoped to routed experts. The supported combinations are BF16 linear + FP4 MoE or MXFP8 linear + FP4 MoE; FP4 linear QAT, export and refit are out of scope. verl defines its `MXFP4` preset as MXFP4 expert weights with dynamic MXFP8 expert activations (group size 32 and E8M0 scales for both). This is intentionally different from compressed-tensors' `MXFP4` preset, which is W4A4, and from vLLM's native `Mxfp4Config`, whose checkpoint format is nominally W4A16.
+MXFP4 and NVFP4 support is scoped to routed experts. Both engines support BF16 linear + FP4 MoE; vLLM additionally supports MXFP8 linear + MXFP4 MoE with QAT mode `mxfp4_mxfp8`. SGLang does not support that mixed override. FP4 linear QAT, export and refit are out of scope. verl defines its `MXFP4` preset as MXFP4 expert weights with dynamic MXFP8 expert activations (group size 32 and E8M0 scales for both). This is intentionally different from compressed-tensors' `MXFP4` preset, which is W4A4, and from vLLM's native `Mxfp4Config`, whose checkpoint format is nominally W4A16.
 
 ## Configuration
 
@@ -92,9 +92,10 @@ actor_rollout_ref:
 
 - When `linear`, `moe` and `config_file` are all `null` (the default), the rollout runs in the model's native precision.
 - `linear` and `moe` can be set independently. Setting one of them to `null` keeps those layers in BF16, e.g. `linear: null, moe: FP8_BLOCK` quantizes only the routed experts.
-- `MXFP4` and `NVFP4` are valid only for `moe`. Use `linear: null` for BF16 dense layers or `linear: MXFP8` for MXFP8 dense layers.
+- `MXFP4` and `NVFP4` are valid only for `moe`. Use `linear: null` for BF16 dense layers. vLLM also supports `linear: MXFP8, moe: MXFP4`; that mixed override is not supported by SGLang.
 - The following layers are always kept in BF16: `lm_head`, embeddings, MoE routers (`mlp.gate`, `shared_expert_gate`), vision towers and MTP layers.
 - `ignore` adds more layers on top of these, using HuggingFace layer names. Layers incompatible with the selected scheme are not detected automatically: list them explicitly. For example, MXFP8 on Qwen3.5 requires `["re:.*linear_attn\\.in_proj_[ab]$"]` because these projections are too small for the vLLM kernel.
+- For an all-MXFP8 SGLang rollout, regex ignores must have a known SGLang module-name translation. Use an exact SGLang module name for other ignores; unsupported regex patterns are rejected before launch.
 
 From the command line:
 
@@ -152,15 +153,14 @@ actor_rollout_ref:
         mode: fp8   # FP8_BLOCK -> fp8, MXFP8 -> mxfp8, MXFP4 -> mxfp4, NVFP4 -> nvfp4
 ```
 
-QAT requires `rollout.quantization` to be set, and verl checks at startup that the QAT mode matches the rollout preset. Megatron QAT uses [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) for fake quantization. MXFP8 uses ModelOpt's native `MXFP8_DEFAULT_CFG`; FP8_BLOCK uses a verl backend because ModelOpt's `FP8_DEFAULT_CFG` is per-tensor and its 2D-block preset is weight-only. MXFP4 and NVFP4 use compressed-tensors-compatible QDQ backends so fake quantization is bit-exact with the exported checkpoint format. Both FP4 modes apply only to routed experts. Use `mxfp4_mxfp8` when dense linear layers also use MXFP8.
+QAT requires a quantized rollout: set `rollout.quantization` or use a checkpoint with a native `quantization_config` that verl supports for export and refit. verl checks at startup that the QAT mode matches the effective rollout configuration. Megatron QAT uses [NVIDIA ModelOpt](https://github.com/NVIDIA/Model-Optimizer) for fake quantization. MXFP8 uses ModelOpt's native `MXFP8_DEFAULT_CFG`; FP8_BLOCK uses a verl backend because ModelOpt's `FP8_DEFAULT_CFG` is per-tensor and its 2D-block preset is weight-only. MXFP4 and NVFP4 use compressed-tensors-compatible QDQ backends so fake quantization is bit-exact with the exported checkpoint format. Both FP4 modes apply only to routed experts. Use `mxfp4_mxfp8` for the vLLM combination where dense linear layers also use MXFP8.
 Its `ignore_patterns` use Megatron module names and `fnmatch` syntax. Keep them aligned with the rollout ignore list; for Qwen3.5 MXFP8, add `*linear_attn.in_proj_a` and `*linear_attn.in_proj_b`.
 
 ### Engine notes
 
 - **vLLM**: `VLLM_USE_DEEP_GEMM_E8M0` defaults to `0` for `FP8_BLOCK`, so its block scales stay in FP32. MX formats retain the E8M0 path.
 - **vLLM MXFP4**: vLLM 0.29's compressed-tensors path otherwise ignores the MXFP8 input scheme and selects W4A4 CUTLASS. verl routes this explicit W4A8(MX) scheme through vLLM's MXFP4 backend oracle with `kMxfp8Dynamic` and defaults to FlashInfer CUTLASS. The TRTLLM MXFP4×MXFP8 backend is not used because it produced a large numerical mismatch on Qwen3.5-35B-A3B.
 - **vLLM NVFP4**: verl supports only W4A4 routed experts and forces FlashInfer CUTLASS. W4A16, dense NVFP4 and other MoE backends are rejected.
-- **SGLang**: the first launch of a quantized MoE model can take several minutes while Triton compiles FP8 MoE kernels during CUDA graph capture. Later launches use the cache.
 
 ## Example
 
@@ -179,8 +179,6 @@ python3 -m verl.trainer.main_ppo \
     ...
 ```
 
-Use `actor_rollout_ref.rollout.name=sglang` for SGLang. For complete scripts, see the [low-precision recipe directory](https://github.com/verl-project/verl-recipe/tree/main/low_precision).
-
 ## Measuring the mismatch
 
 `scripts/rollout_logprob_diff.py` runs one rollout, recomputes the log-probs with the Megatron actor, and reports the gap (`kl`, `k3_kl`, absolute log-prob difference, etc.). Use it to check a configuration before a long run:
@@ -194,9 +192,11 @@ python scripts/rollout_logprob_diff.py --model-path ~/models/Qwen3.5-35B-A3B --e
     --r3-router-replay
 ```
 
-Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arguments. Reference `k3_kl` on Qwen3.5-35B-A3B (EP=8, rollout TP=2):
+Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arguments. Reference `k3_kl` on
+Qwen3.5-35B-A3B (EP=8; vLLM and Blackwell SGLang rollout TP=2; Blackwell
+SGLang uses the default CUDA Graph configuration):
 
-| GPU | Engine | R3 | Rollout | MoE backend | k3_kl |
+| GPU | Engine | R3 | Rollout | Backend | k3_kl |
 |---|---|---|---|---|---|
 | Hopper | vLLM | ✅ | BF16 | - | 0.00031 |
 | Hopper | SGLang | ✅ | BF16 | - | 0.00031 |
@@ -206,18 +206,23 @@ Pass `--engine sglang` for SGLang, and any extra Hydra overrides as trailing arg
 | Hopper | SGLang | ✅ | FP8_BLOCK + FP8 QAT | - | 0.00188 |
 | Blackwell | vLLM | ❌ | MXFP8 linear + MoE | auto | 0.00148 |
 | Blackwell | vLLM | ❌ | MXFP8 linear + MoE + MXFP8 QAT | auto | 0.00315 |
+| Blackwell | SGLang | ❌ | MXFP8 linear + MoE | DeepGEMM (Linear + MoE) | 0.00324 |
+| Blackwell | SGLang | ❌ | MXFP8 linear + MoE + MXFP8 QAT | DeepGEMM (Linear + MoE) | 0.00327 |
 | Blackwell | vLLM | ❌ | MXFP4 MoE | FlashInfer CUTLASS | 0.00858 |
 | Blackwell | vLLM | ❌ | MXFP4 MoE + MXFP4 QAT | FlashInfer CUTLASS | 0.00210 |
+| Blackwell | SGLang | ❌ | MXFP4 MoE | DeepGEMM (MoE; BF16 Linear) | 0.00571 |
+| Blackwell | SGLang | ❌ | MXFP4 MoE + MXFP4 QAT | DeepGEMM (MoE; BF16 Linear) | 0.00200 |
 | Blackwell | vLLM | ✅ | NVFP4 MoE | FlashInfer CUTLASS | 0.00813 |
 | Blackwell | vLLM | ✅ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS | 0.00574 |
+| Blackwell | SGLang | ❌ | NVFP4 MoE | FlashInfer CUTLASS (MoE; BF16 Linear) | 0.00946 |
+| Blackwell | SGLang | ❌ | NVFP4 MoE + NVFP4 QAT | FlashInfer CUTLASS (MoE; BF16 Linear) | 0.00653 |
 
 ## Limitations
 
-- `FP8_BLOCK` is supported on Hopper. `MXFP8` with vLLM is supported on Blackwell; SGLang has no per-layer MXFP8 support.
+- SGLang MXFP8, MXFP4 and NVFP4 custom refit requires the default bucketed
+  weight sync so all derived kernel layouts can be rebuilt after the final
+  bucket. The checkpoint engine's `delta_flush` wire format is not supported.
 - Native vLLM `Fp8Config` rollout is supported without a rollout override. Standard
   128x128 block-FP8 models are quantized by the trainer-side exporter; DeepSeek-V4
   keeps the private mixed FP8/FP4 layout already exported by Megatron-Bridge.
-- Native compressed-tensors configs are also supported without an override. The
-  same trainer-side dispatcher handles native FP8, native compressed-tensors and
-  `rollout.quantization` compressed-tensors overrides.
 - Megatron QAT and the rollout use separate ignore lists and naming conventions. Users must keep them aligned.

verl/utils/qat/__init__.py

@@ -36,6 +36,7 @@
     Fp8ConfigWeightQuantizer,
     build_compressed_tensors_config,
     dispatch_rollout_weight_quantizer,
+    is_all_mxfp8_config,
     qat_modes_for_config,
 )
 from verl.utils.qat.core import (
@@ -56,5 +57,6 @@
     "Fp8ConfigWeightQuantizer",
     "build_compressed_tensors_config",
     "dispatch_rollout_weight_quantizer",
+    "is_all_mxfp8_config",
     "qat_modes_for_config",
 ]

verl/utils/qat/compressed_tensors.py

@@ -49,6 +49,7 @@
     "Fp8ConfigWeightQuantizer",
     "dispatch_rollout_weight_quantizer",
     "build_compressed_tensors_config",
+    "is_all_mxfp8_config",
     "qat_modes_for_config",
 ]
 
@@ -142,6 +143,13 @@ def qat_modes_for_config(ct_config: Mapping[str, Any]) -> set[Optional[str]]:
     return {SUPPORTED_SCHEMES[preset] for preset in presets}
 
 
+def is_all_mxfp8_config(ct_config: Mapping[str, Any]) -> bool:
+    """Whether every compressed-tensors group uses the MXFP8 preset."""
+    if ct_config.get("quant_method") != "compressed-tensors" or not ct_config.get("config_groups"):
+        return False
+    return all(_preset_of(scheme) == "MXFP8" for scheme in _group_schemes(ct_config))
+
+
 def _matches(name: str, targets: Iterable[str]) -> bool:
     for target in targets:
         if target.startswith("re:"):
@@ -521,7 +529,7 @@ def dispatch_rollout_weight_quantizer(
     model_config: Any,
     rollout_config: Any,
     qat_config: Any = None,
-) -> Optional[CompressedTensorsWeightQuantizer | Fp8ConfigWeightQuantizer]:
+) -> Optional[Any]:
     """Select the trainer-side exporter for native or overridden rollout quantization.
 
     The rollout engine keeps its native ``quantization_config`` unless
@@ -565,7 +573,12 @@ def dispatch_rollout_weight_quantizer(
     quant_method = str(quantization_config.get("quant_method", "")).lower()
     if quant_method == "compressed-tensors":
         qat_modes = qat_modes_for_config(quantization_config) if qat_mode is not None else set()
-        quantizer = CompressedTensorsWeightQuantizer(quantization_config)
+        if getattr(rollout_config, "name", None) == "sglang" and is_all_mxfp8_config(quantization_config):
+            from verl.utils.sglang.sglang_quant_utils import SGLangMxfp8WeightQuantizer
+
+            quantizer = SGLangMxfp8WeightQuantizer(quantization_config)
+        else:
+            quantizer = CompressedTensorsWeightQuantizer(quantization_config)
     elif "fp8" in quant_method:
         qat_modes = {"fp8"}
         model_types = {