diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index c99510301..0f8874b63 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -204,6 +204,15 @@ def is_qwen3_5(config) -> bool: ) +def is_qwen3_5_mtp_draft(config) -> bool: + """The Qwen3.5 MoE MTP draft: _config_draft_model rewrites architectures[0] to + Qwen3_5ForCausalLMMTP before quantization is resolved.""" + return ( + _hf_arch(config) == "Qwen3_5ForCausalLMMTP" + and _hf_attr(config, "model_type") == "qwen3_5_moe" + ) + + def is_deepseek_v4(config) -> bool: return _hf_arch(config) in ( "DeepseekV4ForCausalLM", @@ -1772,9 +1781,16 @@ class ModelConfig: f"Using CLI-specified quantization ({self.quantization}) which is " f"compatible with HF config quant_method ({quant_method})." ) - elif self.is_draft_model: + elif self.is_draft_model and not ( + self.is_draft_quantization_explicit + and self.quantization in REQUANTIZATION_METHODS + and is_hip() + and is_qwen3_5_mtp_draft(self.hf_config) + ): # Allow auto-detection of quantization from checkpoint for draft model - # only if the CLI quantization is not compatible + # only if the CLI quantization is not compatible. An explicit + # online-requantization request for the draft (e.g. quark_mxfp4 + # for an MTP stack the checkpoint left in bf16) is honored below. logger.info( f"Draft model quantization ({quant_method}) differs from " f"main model quantization ({self.quantization}). " diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 98da6a4cb..55b5d6b3e 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -41,7 +41,11 @@ from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator from tqdm.auto import tqdm from sglang.srt.configs.load_config import LoadConfig -from sglang.srt.configs.model_config import REQUANTIZATION_METHODS, ModelConfig +from sglang.srt.configs.model_config import ( + REQUANTIZATION_METHODS, + ModelConfig, + is_qwen3_5_mtp_draft, +) from sglang.srt.distributed import get_world_group from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config from sglang.srt.layers.quantization.fp8 import Fp8Config @@ -61,6 +65,7 @@ from sglang.srt.utils import ( BAR_FORMAT, find_local_repo_dir, is_cpu, + is_hip, log_info_on_rank0, print_warning_once, ) @@ -259,6 +264,37 @@ def _resolve_explicit_draft_quant_config( return quant_config +def _quark_draft_online_quant_config( + model_config: ModelConfig, hf_quant_config: dict +) -> Optional[QuantizationConfig]: + """Explicit ``--speculative-draft-model-quantization quark_mxfp4`` on a Quark + checkpoint whose MTP/NextN draft experts were exported in bf16 (listed under + ``exclude``): quantize the draft's routed experts online to MXFP4 instead of + running them through the bf16 MoE path. Only the draft model is affected; the + target model keeps its serialized Quark scheme.""" + if not ( + model_config.is_draft_model + and model_config.is_draft_quantization_explicit + and model_config.quantization == "quark_mxfp4" + and hf_quant_config.get("quant_method") == "quark" + # ROCm + Qwen3.5 MTP draft only (validated combination); anything else is + # left exactly as before. + and is_hip() + and is_qwen3_5_mtp_draft(model_config.hf_config) + ): + return None + excluded = hf_quant_config.get("exclude") or [] + if not any(str(name).startswith("mtp.layers.0.mlp.experts") for name in excluded): + return None + from sglang.srt.layers.quantization.quark.quark import QuarkConfig + + logger.info( + "Draft MTP experts are unquantized in the Quark checkpoint; " + "quantizing them online to MXFP4 (quark_mxfp4) for the draft model." + ) + return QuarkConfig(online_scheme="quark_mxfp4", hf_config=model_config.hf_config) + + def _modelopt_quant_section(config: dict) -> dict: """Return ModelOpt quant settings from nested or flat ``hf_quant_config.json``. @@ -313,7 +349,11 @@ def get_quant_config( # This is only used by quantization methods that support requantization (e.g. from nvfp4/fp8 to mxfp4). if model_config.quantization in REQUANTIZATION_METHODS: hf_quant_config["requantization_method"] = model_config.quantization - + draft_online = _quark_draft_online_quant_config( + model_config, hf_quant_config + ) + if draft_online is not None: + return draft_online return _resolve_explicit_draft_quant_config( model_config, quant_cls.from_config(hf_quant_config) )