minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] fix: reject unsupported modelopt checkpoint algorithms (#35182)

Browse Source
This commit is contained in:
Mick
2026-08-20 09:13:24 +08:00
committed by GitHub
parent 1cf2b8c54d
commit ab203663c4
5 changed files with 185 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/utils/quantization_utils.py
+30 -5
View File
@@ -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}")
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+86
View File
@@ -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(
python/sglang/srt/layers/modelopt_utils.py
+34
View File
@@ -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())
python/sglang/srt/layers/quantization/base_config.py
+5 -6
View File
@@ -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.