From ab203663c4bf2297d7e97cbe3310d03db4085ec3 Mon Sep 17 00:00:00 2001 From: Mick Date: Thu, 20 Aug 2026 09:13:24 +0800 Subject: [PATCH] [diffusion] fix: reject unsupported modelopt checkpoint algorithms (#35182) --- .../runtime/utils/quantization_utils.py | 35 ++++++-- .../test/unit/test_transformer_quant.py | 86 +++++++++++++++++++ python/sglang/srt/layers/modelopt_utils.py | 34 ++++++++ .../srt/layers/quantization/base_config.py | 11 ++- .../unit/test_checkpoint_quantization.py | 30 +++++++ 5 files changed, 185 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index e3a49f3f4..9bfb72112 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -13,6 +13,7 @@ from sglang.multimodal_gen.runtime.layers.quantization import ( get_quantization_config, ) from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo from sglang.srt.model_loader.checkpoint_quantization import ( resolve_checkpoint_quant_spec, ) @@ -106,7 +107,8 @@ def _resolve_quant_method_name(quant_cfg: dict) -> str: quant_method = quant_cfg.get("quant_method") if quant_method == "bitsandbytes": return "bitsandbytes" - if quant_method != "modelopt": + modelopt_methods = {"modelopt", "modelopt_fp8", "modelopt_fp4"} + if quant_method not in modelopt_methods: return quant_method quant_algo = ( @@ -114,14 +116,37 @@ def _resolve_quant_method_name(quant_cfg: dict) -> str: or quant_cfg.get("quantization", {}).get("quant_algo") or "" ).upper() + if quant_method != "modelopt" and not quant_algo: + # Preserve explicit legacy configs that select the backend directly. + # When an algorithm is present below, validate that it agrees. + return quant_method if quant_algo == "MIXED_PRECISION": raise ValueError( "ModelOpt mixed precision is not supported by the current SGLang diffusion runtime." ) - if "FP8" in quant_algo: - return "modelopt_fp8" - if "FP4" in quant_algo or "NVFP4" in quant_algo: - return "modelopt_fp4" + canonical_method = canonicalize_modelopt_quant_algo(quant_algo) + if ( + quant_method != "modelopt" + and canonical_method is not None + and quant_method != canonical_method + ): + raise ValueError( + f"ModelOpt config declares quant_method={quant_method!r}, but " + f"quant_algo={quant_algo!r} maps to {canonical_method!r}." + ) + supported_algorithms = { + "FP8": "modelopt_fp8", + "NVFP4": "modelopt_fp4", + } + runtime_method = supported_algorithms.get(quant_algo) + if runtime_method is not None: + return runtime_method + if canonical_method is not None: + raise ValueError( + f"ModelOpt quant_algo={quant_algo!r} maps to {canonical_method!r}, but " + "that checkpoint algorithm is not supported by the SGLang diffusion runtime. " + "Supported ModelOpt checkpoint algorithms are FP8 and NVFP4." + ) raise ValueError(f"Unsupported ModelOpt quant_algo for diffusion: {quant_algo}") diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index cb85d61de..62e1e034d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -75,6 +75,7 @@ from sglang.multimodal_gen.runtime.models.dits.flux import FluxSingleTransformer from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum from sglang.multimodal_gen.runtime.platforms.interface import DeviceCapability from sglang.multimodal_gen.runtime.utils.quantization_utils import ( + _resolve_quant_method_name, build_nvfp4_config_from_safetensors_list, get_quant_config, ) @@ -647,6 +648,91 @@ class TestTransformerQuantHelpers(unittest.TestCase): self.assertIsInstance(config, ModelOptFp8Config) self.assertEqual(config.exclude_modules, ["proj_out"]) + def test_modelopt_checkpoint_algorithm_admission(self): + cases = [ + ("modelopt", "FP8", {"ignore": []}, ModelOptFp8Config, None), + ("modelopt_fp8", "FP8", {"ignore": []}, ModelOptFp8Config, None), + ( + "modelopt", + "NVFP4", + {"group_size": 16, "ignore": []}, + ModelOptFp4Config, + None, + ), + ( + "modelopt_fp4", + "NVFP4", + {"group_size": 16, "ignore": []}, + ModelOptFp4Config, + None, + ), + ("modelopt", "MXFP8", {}, None, "maps to 'mxfp8'"), + ("modelopt", "FP4", {}, None, "maps to 'modelopt_fp4'"), + ("modelopt", "NVFP4_AWQ", {}, None, "maps to 'modelopt_fp4'"), + ("modelopt", "W4A16_NVFP4", {}, None, "maps to 'modelopt_fp4'"), + ( + "modelopt", + "MIXED_PRECISION", + {}, + None, + "mixed precision is not supported", + ), + ( + "modelopt", + "FP8_FAKE", + {}, + None, + "Unsupported ModelOpt quant_algo for diffusion: FP8_FAKE", + ), + ( + "modelopt_fp8", + "MXFP8", + {}, + None, + "declares quant_method='modelopt_fp8'.*maps to 'mxfp8'", + ), + ( + "modelopt_fp4", + "FP8", + {}, + None, + "declares quant_method='modelopt_fp4'.*maps to 'modelopt_fp8'", + ), + ] + for ( + quant_method, + quant_algo, + extra_metadata, + expected_type, + expected_error, + ) in cases: + with self.subTest(quant_algo=quant_algo): + metadata = { + "quant_method": quant_method, + "quant_algo": quant_algo, + **extra_metadata, + } + if expected_error is not None: + with self.assertRaisesRegex(ValueError, expected_error): + get_quant_config( + {"quantization_config": metadata}, + "/unused/component/path", + ) + else: + config = get_quant_config( + {"quantization_config": metadata}, + "/unused/component/path", + ) + self.assertIsInstance(config, expected_type) + + def test_explicit_modelopt_method_without_algorithm_is_preserved(self): + for quant_method in ("modelopt_fp8", "modelopt_fp4"): + with self.subTest(quant_method=quant_method): + self.assertEqual( + _resolve_quant_method_name({"quant_method": quant_method}), + quant_method, + ) + @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_rank", return_value=0) @patch("sglang.multimodal_gen.runtime.layers.linear.get_group_size", return_value=1) @patch( diff --git a/python/sglang/srt/layers/modelopt_utils.py b/python/sglang/srt/layers/modelopt_utils.py index 8e9d84351..e964af74f 100644 --- a/python/sglang/srt/layers/modelopt_utils.py +++ b/python/sglang/srt/layers/modelopt_utils.py @@ -2,6 +2,8 @@ ModelOpt related constants """ +from typing import Literal, TypeAlias + QUANT_CFG_CHOICES = { "fp8": "FP8_DEFAULT_CFG", "int4_awq": "INT4_AWQ_CFG", # TODO: add support for int4_awq @@ -9,3 +11,35 @@ QUANT_CFG_CHOICES = { "nvfp4": "NVFP4_DEFAULT_CFG", "nvfp4_awq": "NVFP4_AWQ_LITE_CFG", # TODO: add support for nvfp4_awq } + + +ModelOptQuantMethod: TypeAlias = Literal[ + "modelopt_fp8", + "modelopt_fp4", + "mxfp8", +] + + +_MODELOPT_QUANT_ALGO_TO_METHOD: dict[str, ModelOptQuantMethod] = { + "FP8": "modelopt_fp8", + "MXFP8": "mxfp8", + "FP4": "modelopt_fp4", + "NVFP4": "modelopt_fp4", + "NVFP4_AWQ": "modelopt_fp4", + "W4A16_NVFP4": "modelopt_fp4", +} + + +def canonicalize_modelopt_quant_algo( + quant_algo: object, +) -> ModelOptQuantMethod | None: + """Map a known ModelOpt algorithm name to its SGLang runtime family. + + This is intentionally an exact allowlist. In particular, ``MXFP8`` must not + be mistaken for ordinary ``FP8`` merely because its name contains that + substring. Runtime-specific capability checks remain with each consumer. + """ + + if not isinstance(quant_algo, str): + return None + return _MODELOPT_QUANT_ALGO_TO_METHOD.get(quant_algo.upper()) diff --git a/python/sglang/srt/layers/quantization/base_config.py b/python/sglang/srt/layers/quantization/base_config.py index 509e2cc07..f702bd8f8 100644 --- a/python/sglang/srt/layers/quantization/base_config.py +++ b/python/sglang/srt/layers/quantization/base_config.py @@ -10,6 +10,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type import torch from torch import nn +from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo + if TYPE_CHECKING: from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo @@ -198,12 +200,9 @@ class QuantizationConfig(ABC): # If user specified generic "modelopt", auto-detect the specific method if user_quant == "modelopt": - 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" + canonical_method = canonicalize_modelopt_quant_algo(quant_algo) + if canonical_method is not None: + return canonical_method # The hf_quant_config may be a parsed quant config, so we need to check the # quant_method. diff --git a/test/registered/unit/test_checkpoint_quantization.py b/test/registered/unit/test_checkpoint_quantization.py index aae75d58d..8b6b11c84 100644 --- a/test/registered/unit/test_checkpoint_quantization.py +++ b/test/registered/unit/test_checkpoint_quantization.py @@ -2,6 +2,8 @@ import unittest +from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo +from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.model_loader.checkpoint_quantization import ( CheckpointQuantSpec, resolve_checkpoint_quant_spec, @@ -26,6 +28,34 @@ class _QuantConfigObject: class TestResolveCheckpointQuantSpec(CustomTestCase): + def test_modelopt_quant_algo_canonicalization(self): + cases = { + "FP8": "modelopt_fp8", + "mxfp8": "mxfp8", + "NVFP4": "modelopt_fp4", + "NVFP4_AWQ": "modelopt_fp4", + "W4A16_NVFP4": "modelopt_fp4", + "FP8_FAKE": None, + "MIXED_PRECISION": None, + None: None, + } + for quant_algo, expected in cases.items(): + with self.subTest(quant_algo=quant_algo): + self.assertEqual(canonicalize_modelopt_quant_algo(quant_algo), expected) + + def test_srt_modelopt_override_uses_the_exact_algorithm_allowlist(self): + self.assertEqual( + QuantizationConfig._modelopt_override_quantization_method( + {"quant_algo": "NVFP4"}, "modelopt" + ), + "modelopt_fp4", + ) + self.assertIsNone( + QuantizationConfig._modelopt_override_quantization_method( + {"quant_algo": "NVFP4_FAKE"}, "modelopt" + ) + ) + def test_top_level_quantization_config(self): config = { "quantization_config": {