diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index ec7a6bf4f536..634220afe88f 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/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": diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index f116e12c9a4b..f38bc72e0335 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/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: diff --git a/python/sglang/srt/models/qwen4_exp.py b/python/sglang/srt/models/qwen4_exp.py index 32a51803ef1d..d027cc749f1e 100644 --- a/python/sglang/srt/models/qwen4_exp.py +++ b/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( diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 97c6011ab4be..09260906cbf6 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/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"], diff --git a/test/registered/unit/models/test_qwen4_exp_mixed_precision.py b/test/registered/unit/models/test_qwen4_exp_mixed_precision.py new file mode 100644 index 000000000000..4708bb7b91ab --- /dev/null +++ b/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()