From 65d5a0ec2596d314fbcd88bf179d2ea912aaaa02 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:44:32 +0800 Subject: [PATCH] Support ModelOpt MXFP8 checkpoints (#32538) --- python/sglang/srt/configs/model_config.py | 28 +++++- .../srt/layers/quantization/base_config.py | 4 +- python/sglang/srt/layers/quantization/fp8.py | 6 ++ .../srt/layers/quantization/modelopt_quant.py | 8 +- python/sglang/srt/model_loader/loader.py | 9 ++ .../unit/model_loader/test_modelopt_loader.py | 90 ++++++++++++++++++- 6 files changed, 137 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 3d80a871c..ded01222c 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1249,8 +1249,34 @@ class ModelConfig: return {"quant_method": "w4afp8", "quant_algo": quant_algo} elif quant_algo and ("FP4" in quant_algo or "NVFP4" in quant_algo): return {"quant_method": "modelopt_fp4", "quant_algo": quant_algo} - elif quant_algo and "FP8" in quant_algo: + elif quant_algo == "FP8": return {"quant_method": "modelopt_fp8", "quant_algo": quant_algo} + elif quant_algo == "MXFP8": + group_size = json_quant_configs.get("group_size", 32) + ignored_layers = json_quant_configs.get( + "exclude_modules", json_quant_configs.get("ignore") + ) + kv_cache_quant_algo = json_quant_configs.get("kv_cache_quant_algo") + if kv_cache_quant_algo is None: + kv_cache_scheme = json_quant_configs.get("kv_cache_scheme") + if ( + isinstance(kv_cache_scheme, dict) + and kv_cache_scheme.get("type") == "float" + and kv_cache_scheme.get("num_bits") == 8 + ): + kv_cache_quant_algo = "FP8" + parsed = { + "quant_method": "mxfp8", + "quant_algo": quant_algo, + "activation_scheme": "dynamic", + "weight_block_size": [1, group_size], + "scale_fmt": "ue8m0", + } + if ignored_layers is not None: + parsed["modules_to_not_convert"] = ignored_layers + if kv_cache_quant_algo is not None: + parsed["kv_cache_quant_algo"] = kv_cache_quant_algo + return parsed else: return None diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index ccce6169f..56e95c50b 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -190,7 +190,9 @@ class QuantizationConfig(ABC): # If user specified generic "modelopt", auto-detect the specific method if user_quant == "modelopt": - if "FP8" in quant_algo: + if quant_algo == "MXFP8": + return "mxfp8" + elif quant_algo == "FP8": return "modelopt_fp8" elif "NVFP4" in quant_algo or "FP4" in quant_algo: return "modelopt_fp4" diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index f94ce61ae..bee7f5523 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -235,6 +235,7 @@ class Fp8Config(QuantizationConfig): packed_modules_mapping: Optional[Dict[str, List[str]]] = None, use_mxfp8: bool = False, is_fp4_experts: bool = False, + kv_cache_quant_algo: Optional[str] = None, ) -> None: super().__init__() # DSV4 mxfp4-packed (True) vs converted FP8 (False); injected by @@ -258,6 +259,7 @@ class Fp8Config(QuantizationConfig): ) self.packed_modules_mapping = packed_modules_mapping or {} self.use_mxfp8 = use_mxfp8 + self.kv_cache_quant_algo = kv_cache_quant_algo if weight_block_size is not None: if not is_checkpoint_fp8_serialized: raise ValueError( @@ -322,6 +324,9 @@ class Fp8Config(QuantizationConfig): normalized.append(f"model.{base}") ignored_layers = normalized weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) + kv_cache_quant_algo = cls.get_from_keys_or( + config, ["kv_cache_quant_algo"], None + ) if use_mxfp8: # MXFP8 (OCP) spec fixes block size to [1, 32]; ckpt field is metadata only. if weight_block_size is not None and weight_block_size != [1, 32]: @@ -337,6 +342,7 @@ class Fp8Config(QuantizationConfig): weight_block_size=weight_block_size, packed_modules_mapping=packed_modules_mapping, use_mxfp8=use_mxfp8, + kv_cache_quant_algo=kv_cache_quant_algo, ) def get_quant_method( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 45c0eef71..cd2e6e67a 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -465,11 +465,11 @@ class ModelOptFp8Config(ModelOptQuantConfig): raise ValueError( "Cannot find 'quant_algo' in the model's quantization config. " ) - if "FP8" not in quant_method: + if quant_method != "FP8": raise ValueError( - "ModelOptFp8Config only supports static FP8 quantization in SGLang. " - "For FP4 quantization, use ModelOptFp4Config. " - "Check the quantization config for your model's configuration." + "ModelOptFp8Config only supports regular FP8 quantization, " + f"but found {quant_method!r}. Use the native 'mxfp8' " + "quantization method for MXFP8 or ModelOptFp4Config for FP4." ) return cls( diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 2cc109864..88662bfee 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -835,6 +835,15 @@ class DefaultModelLoader(BaseModelLoader): quant_config = getattr(model, "quant_config", None) is_nvfp4_online = getattr(quant_config, "is_nvfp4_online", False) + is_mxfp8 = quant_config is not None and quant_config.get_name() == "mxfp8" + if is_mxfp8: + weights = ( + ( + f"{name}_inv" if name.endswith(".weight_scale") else name, + loaded_weight, + ) + for name, loaded_weight in weights + ) if is_nvfp4_online: # Scope exact FP4 quantization math to load-time conversion only; diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 246a26802..69bfcc706 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -17,14 +17,15 @@ from sglang.srt.configs.model_config import ModelConfig 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 Fp8LinearMethod +from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptFp4Config, ModelOptFp4LinearMethod, + ModelOptFp8Config, ModelOptMixedPrecisionConfig, ModelOptNvFp4A16LinearMethod, ) -from sglang.srt.model_loader.loader import ModelOptModelLoader +from sglang.srt.model_loader.loader import DefaultModelLoader, ModelOptModelLoader from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device @@ -462,6 +463,7 @@ class TestParseQuantHfConfig(CustomTestCase): ({"quant_algo": "NVFP4_AWQ"}, "modelopt_fp4"), ({"quant_method": "modelopt", "quant_algo": "MIXED_PRECISION"}, "w4afp8"), ({"quant_algo": "FP8"}, "modelopt_fp8"), + ({"quant_algo": "MXFP8"}, "mxfp8"), ({"quant_algo": "FP4"}, "modelopt_fp4"), ({"quant_algo": "MIXED_PRECISION"}, "w4afp8"), ({"quant_method": "modelopt"}, "modelopt"), @@ -509,6 +511,90 @@ class TestParseQuantHfConfig(CustomTestCase): self.assertEqual(cfg.group_size, 16) self.assertTrue(cfg.is_awq) + def test_modelopt_mxfp8_config(self): + """ModelOpt MXFP8 metadata must select block scales and retain FP8 KV policy.""" + model_config = ModelConfig.__new__(ModelConfig) + for kv_cache_config in ( + {"kv_cache_quant_algo": "FP8"}, + {"kv_cache_scheme": {"type": "float", "num_bits": 8}}, + ): + with self.subTest(kv_cache_config=kv_cache_config): + result = model_config._parse_modelopt_quant_config( + { + "quantization": { + "quant_algo": "MXFP8", + "group_size": 32, + "exclude_modules": ["lm_head"], + **kv_cache_config, + } + } + ) + self.assertEqual(result["quant_method"], "mxfp8") + self.assertEqual(result["scale_fmt"], "ue8m0") + + quant_config = Fp8Config.from_config(result) + self.assertEqual(quant_config.get_name(), "mxfp8") + self.assertEqual(quant_config.activation_scheme, "dynamic") + self.assertEqual(quant_config.weight_block_size, [1, 32]) + self.assertIn("lm_head", quant_config.ignored_layers) + self.assertEqual(quant_config.kv_cache_quant_algo, "FP8") + + def test_modelopt_mxfp8_override(self): + """Generic ModelOpt selection must not route MXFP8 to scalar FP8.""" + self.assertEqual( + ModelOptFp8Config.override_quantization_method( + {"quant_algo": "MXFP8"}, "modelopt" + ), + "mxfp8", + ) + + def test_modelopt_mxfp8_weight_loading(self): + """ModelOpt MXFP8 block scales must reach native scale parameters.""" + weight = torch.empty(1) + weights = [ + ("model.q_proj.weight_scale", weight), + ("model.q_proj.input_weight_scale", weight), + ("model.q_proj.weight_scale_inv", weight), + ] + + def load_names(quant_config): + model = nn.Module() + model.quant_config = quant_config + loaded_names = [] + model.load_weights = lambda weights: loaded_names.extend( + name for name, _ in weights + ) + with patch( + "sglang.srt.model_loader.loader.is_cuda_alike", return_value=False + ): + DefaultModelLoader.load_weights_and_postprocess( + model, iter(weights), torch.device("cpu") + ) + return loaded_names + + mxfp8_config = Fp8Config( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[1, 32], + use_mxfp8=True, + ) + self.assertEqual( + load_names(mxfp8_config), + [ + "model.q_proj.weight_scale_inv", + "model.q_proj.input_weight_scale", + "model.q_proj.weight_scale_inv", + ], + ) + self.assertEqual( + load_names(Fp8Config(is_checkpoint_fp8_serialized=True)), + [ + "model.q_proj.weight_scale", + "model.q_proj.input_weight_scale", + "model.q_proj.weight_scale_inv", + ], + ) + def test_non_modelopt_quant_method_unchanged(self): """Non-modelopt quant_method (e.g. 'gptq') must NOT enter the modelopt path.""" self.model_config.hf_config.quantization_config = {