[Model] Qwen4-Exp: honour per-layer FP8 for the PLE table and the MTP experts in ModelOpt mixed checkpoints

python/sglang/srt/layers/quantization/modelopt_quant.py

@@ -779,6 +779,12 @@ def embedding(self, layer: torch.nn.Module, input_: torch.Tensor) -> torch.Tenso
         return out.view(*index_shape, hidden).to(self.params_dtype)
 
 
+# FP8_PB_WO is ModelOpt's canonical 2D block-FP8 name.
+# Qwen3.8-Flash-Next-NVFP4 used FP8_BLOCK_SCALES for the same tensor layout.
+# Keep it as an alias.
+_BLOCK_FP8_ALGOS = ("FP8_PB_WO", "FP8_BLOCK_SCALES")
+
+
 class ModelOptMixedPrecisionConfig(ModelOptQuantConfig):
     """Configuration for ModelOpt MIXED_PRECISION checkpoints."""
 
@@ -875,6 +881,20 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig:
         if group_size is None:
             group_size = 16
 
+        # Block-FP8 layers carry their block size as group_size
+        # (default 128). One Fp8Config serves every such layer.
+        block_sizes = {
+            int(layer_info.get("group_size", 128))
+            for layer_info in quantized_layers.values()
+            if layer_info.get("quant_algo", "").upper() in _BLOCK_FP8_ALGOS
+        }
+        if len(block_sizes) > 1:
+            raise ValueError(
+                "MIXED_PRECISION currently requires all block-FP8 layers to "
+                f"use one group_size, got {sorted(block_sizes)}."
+            )
+        block_size = next(iter(block_sizes), 128)
+
         packed_modules_mapping = config.get("packed_modules_mapping")
         fp8_config = ModelOptFp8Config(
             is_checkpoint_fp8_serialized=True,
@@ -885,7 +905,7 @@ def from_config(cls, config: Dict[str, Any]) -> ModelOptMixedPrecisionConfig:
         fp8_pb_wo_config = Fp8Config(
             is_checkpoint_fp8_serialized=True,
             activation_scheme="dynamic",
-            weight_block_size=[128, 128],
+            weight_block_size=[block_size, block_size],
             packed_modules_mapping=packed_modules_mapping,
         )
         mxfp8_config = Fp8Config(
@@ -999,7 +1019,7 @@ def get_quant_method(
                 return UnquantizedLinearMethod()
             if quant_algo == "FP8":
                 return ModelOptFp8LinearMethod(self.fp8_config)
-            if quant_algo == "FP8_PB_WO":
+            if quant_algo in _BLOCK_FP8_ALGOS:
                 return Fp8LinearMethod(self.fp8_pb_wo_config)
             if quant_algo == "MXFP8":
                 return Fp8LinearMethod(self.mxfp8_config)
@@ -1028,6 +1048,8 @@ def get_quant_method(
                 return None
             if quant_algo == "FP8":
                 return ModelOptFp8MoEMethod(self.fp8_config)
+            if quant_algo in _BLOCK_FP8_ALGOS:
+                return Fp8MoEMethod(self.fp8_pb_wo_config)
             if quant_algo == "MXFP8":
                 return Fp8MoEMethod(self.mxfp8_config)
             if quant_algo == "NVFP4":

python/sglang/srt/models/qwen3_5_mtp.py

@@ -58,12 +58,18 @@ def _mtp_quant_config(quant_config):
     # Serialized Qwen3.5 ModelOpt checkpoints keep embedded MTP weights in
     # BF16. Disable quantization for those checkpoints; non-serialized
     # modelopt_fp4 still converts MoE expert weights on load.
+    if quant_config and quant_config.get_name() == "modelopt_mixed":
+        # Qwen3.8-Flash-Next-NVFP4 ships `mtp.*.experts` as block-FP8
+        quantized_layers = getattr(quant_config, "quantized_layers", {}) or {}
+        if any(
+            isinstance(layer, str) and layer.startswith("mtp.")
+            for layer in quantized_layers
+        ):
+            return quant_config
+        return None
     if quant_config and (
-        quant_config.get_name() == "modelopt_mixed"
-        or (
-            quant_config.get_name() == "modelopt_fp4"
-            and quant_config.is_checkpoint_nvfp4_serialized
-        )
+        quant_config.get_name() == "modelopt_fp4"
+        and quant_config.is_checkpoint_nvfp4_serialized
     ):
         return None
     if is_npu() and get_spec().speculative_draft_model_quantization is None:

python/sglang/srt/models/qwen4_exp.py

@@ -410,6 +410,32 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return (x_norm * weight).to(compute_dtype)
 
 
+def _ple_table_is_fp8(
+    config: Qwen4ExpTextConfig,
+    quant_config: Optional[QuantizationConfig],
+    ngram_prefix: str,
+) -> bool:
+    """Whether the PLE n-gram table should be stored as fp8.
+
+    True when the config pins it, or when the whole checkpoint
+    is fp8, or when a mixed-precision checkpoint lists the table
+    itself as FP8.
+    """
+    if getattr(config, "ple_embedding_dtype", None) == "float8_e4m3fn":
+        return True
+    if quant_config is None:
+        return False
+    if quant_config.get_name() == "fp8":
+        return True
+    resolve = getattr(quant_config, "_resolve_quant_algo", None)
+    if resolve is None:
+        return False
+    try:
+        return resolve(ngram_prefix) == "FP8"
+    except ValueError:
+        return False
+
+
 class Qwen4ExpNGramEmbedding(nn.Module):
     _MASK64 = (1 << 64) - 1
     _SPLITMIX_GAMMA = 0x9E3779B97F4A7C15
@@ -423,6 +449,7 @@ def __init__(
         embedding_dim: int,
         ple_layer_index: int = 0,
         quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         super().__init__()
         self.config = config
@@ -484,8 +511,7 @@ def __init__(
             self.head_dim_per_ngram,
             params_dtype=(
                 torch.float8_e4m3fn
-                if (quant_config is not None and quant_config.get_name() == "fp8")
-                or getattr(config, "ple_embedding_dtype", None) == "float8_e4m3fn"
+                if _ple_table_is_fp8(config, quant_config, f"{prefix}.ngram_embedding")
                 else torch.bfloat16
             ),
             output_dtype=torch.bfloat16,
@@ -881,6 +907,7 @@ def __init__(
             self.ple_embed_dim,
             ple_layer_index=ple_layer_index,
             quant_config=quant_config,
+            prefix=f"{prefix}.ple_embedding" if prefix else "ple_embedding",
         )
         if config.ple_offload_embedding:
             self.ple_embedding.ngram_embedding = Qwen4ExpPinnedHostEmbedding(

test/registered/unit/model_loader/test_modelopt_loader.py

@@ -22,7 +22,8 @@
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import should_apply_lm_head_quant_method
 from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
-from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
+from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod, Fp8MoEMethod
 from sglang.srt.layers.quantization.modelopt_quant import (
     ModelOptFp4Config,
     ModelOptFp4LinearMethod,
@@ -738,6 +739,55 @@ def test_fp8_pb_wo_dispatches_to_native_block_fp8(self):
         self.assertTrue(method.quant_config.is_checkpoint_fp8_serialized)
         self.assertEqual(method.quant_config.activation_scheme, "dynamic")
 
+    def _block_fp8_moe_method(self, algo, group_size=128):
+        quant_config = ModelOptMixedPrecisionConfig.from_config(
+            {
+                "quant_algo": "MIXED_PRECISION",
+                "quantized_layers": {
+                    "mtp.layers.0.mlp.experts": {
+                        "quant_algo": algo,
+                        "group_size": group_size,
+                    },
+                },
+                "packed_modules_mapping": {},
+            }
+        )
+        # Type dispatch only needs a FusedMoE instance; skip GPU weight setup.
+        layer = FusedMoE.__new__(FusedMoE)
+        return quant_config.get_quant_method(layer, "mtp.layers.0.mlp.experts")
+
+    def test_block_fp8_moe_dispatches_under_both_algo_names(self):
+        for algo in ("FP8_PB_WO", "FP8_BLOCK_SCALES"):
+            with self.subTest(algo=algo):
+                method = self._block_fp8_moe_method(algo)
+                self.assertIsInstance(method, Fp8MoEMethod)
+                self.assertEqual(method.quant_config.weight_block_size, [128, 128])
+                self.assertEqual(method.quant_config.activation_scheme, "dynamic")
+                self.assertTrue(method.quant_config.is_checkpoint_fp8_serialized)
+
+    def test_block_fp8_block_size_follows_checkpoint_group_size(self):
+        method = self._block_fp8_moe_method("FP8_BLOCK_SCALES", group_size=64)
+        self.assertEqual(method.quant_config.weight_block_size, [64, 64])
+
+    def test_block_fp8_conflicting_group_sizes_are_rejected(self):
+        with self.assertRaisesRegex(ValueError, "one group_size"):
+            ModelOptMixedPrecisionConfig.from_config(
+                {
+                    "quant_algo": "MIXED_PRECISION",
+                    "quantized_layers": {
+                        "mtp.layers.0.mlp.experts": {
+                            "quant_algo": "FP8_BLOCK_SCALES",
+                            "group_size": 128,
+                        },
+                        "model.layers.0.self_attn.q_proj": {
+                            "quant_algo": "FP8_PB_WO",
+                            "group_size": 64,
+                        },
+                    },
+                    "packed_modules_mapping": {},
+                }
+            )
+
     def test_incomplete_inline_config_falls_back_to_hf_quant_config_file(self):
         packed_modules_mapping = {
             "qkv_proj": ["q_proj", "k_proj", "v_proj"],

test/registered/unit/models/test_qwen4_exp_mixed_precision.py

@@ -0,0 +1,100 @@
+"""Regression tests for Qwen4-Exp on three-precision ModelOpt MIXED_PRECISION
+checkpoints (`nvidia/Qwen3.8-Flash-Next-NVFP4`: NVFP4 experts + FP8 PLE table
++ block-FP8 MTP experts).
+
+Each case guards a decision that used to be made before consulting the
+per-layer `quantized_layers` map:
+  1. The PLE n-gram table's storage dtype is fixed at construction. Building it
+     bf16 and switching to fp8 in load_weights is impossible once the table
+     sits in pinned host memory (`--ple-offload-embedding`, default on), so
+     the FP8 entry in `quantized_layers` must be honoured at construction.
+  2. The draft (MTP) module dropped its quant config for every modelopt_mixed
+     checkpoint. With `mtp.*.experts` listed as block-FP8 that cast the fp8
+     expert values to bf16 without their scales and silently discarded
+     `weight_scale_inv` (accept length fell from ~3.0 to ~1.6 with no error).
+"""
+
+import unittest
+
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptMixedPrecisionConfig,
+)
+from sglang.srt.models.qwen3_5_mtp import _mtp_quant_config
+from sglang.srt.models.qwen4_exp import _ple_table_is_fp8
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+PLE_PREFIX = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"
+MTP_EXPERTS = "mtp.layers.0.mlp.experts"
+
+
+def _mixed_config(quantized_layers):
+    return ModelOptMixedPrecisionConfig.from_config(
+        {
+            "quantization": {
+                "quant_algo": "MIXED_PRECISION",
+                "exclude_modules": [],
+                "group_size": 16,
+                "quantized_layers": quantized_layers,
+            }
+        }
+    )
+
+
+class _Cfg:
+    ple_embedding_dtype = None
+
+
+class TestPLETableDtype(CustomTestCase):
+    def test_mixed_precision_fp8_ple_entry_selects_fp8_storage(self):
+        quant_config = _mixed_config(
+            {
+                PLE_PREFIX: {"quant_algo": "FP8"},
+                "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
+            }
+        )
+        self.assertTrue(_ple_table_is_fp8(_Cfg(), quant_config, PLE_PREFIX))
+        # A checkpoint that quantizes only the experts keeps the table bf16.
+        self.assertFalse(
+            _ple_table_is_fp8(
+                _Cfg(),
+                _mixed_config(
+                    {
+                        "model.language_model.layers.0.mlp.experts": {
+                            "quant_algo": "NVFP4"
+                        }
+                    }
+                ),
+                PLE_PREFIX,
+            )
+        )
+
+    def test_hf_style_prefix_resolves_the_same_entry(self):
+        # Multimodal prefixes drift between `model.language_model.` and
+        # `language_model.model.`; both must find the FP8 entry.
+        quant_config = _mixed_config({PLE_PREFIX: {"quant_algo": "FP8"}})
+        drifted = PLE_PREFIX.replace("model.language_model.", "language_model.model.")
+        self.assertTrue(_ple_table_is_fp8(_Cfg(), quant_config, drifted))
+
+
+class TestMTPQuantConfig(CustomTestCase):
+    def test_mixed_precision_keeps_config_when_mtp_layers_are_quantized(self):
+        quant_config = _mixed_config(
+            {
+                "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
+                MTP_EXPERTS: {"quant_algo": "FP8_BLOCK_SCALES", "group_size": 128},
+            }
+        )
+        self.assertIs(_mtp_quant_config(quant_config), quant_config)
+
+    def test_mixed_precision_drops_config_when_mtp_ships_bf16(self):
+        quant_config = _mixed_config(
+            {"model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}}
+        )
+        self.assertIsNone(_mtp_quant_config(quant_config))
+
+
+if __name__ == "__main__":
+    unittest.main()