From 4ad990ba7d75bb9f948f5f6bd8d79a66b5d3fd63 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 6 Aug 2026 11:01:55 -0700 Subject: [PATCH] [ModelOpt FP4] Support online MoE weight quantization (#33115) Co-authored-by: Brayden Zhong --- docs/docs/advanced_features/quantization.mdx | 10 +- .../docs/references/environment_variables.mdx | 4 +- python/sglang/srt/configs/model_config.py | 11 +- python/sglang/srt/layers/moe/ep_moe/layer.py | 3 +- .../srt/layers/quantization/modelopt_quant.py | 132 ++++++++++----- .../srt/layers/quantization/nvfp4_online.py | 155 ++++++++++++++---- python/sglang/srt/model_loader/loader.py | 36 +++- .../sglang/srt/model_loader/weight_utils.py | 39 ++++- python/sglang/srt/models/qwen3_5_mtp.py | 13 +- python/sglang/srt/server_args.py | 10 ++ .../models_e2e/test_qwen35_fp4_mtp.py | 2 + .../quantization/test_modelopt_nvfp4.py | 57 +++++++ .../unit/model_loader/test_modelopt_loader.py | 94 ++++++++++- .../unit/server_args/test_server_args.py | 12 ++ 14 files changed, 472 insertions(+), 106 deletions(-) diff --git a/docs/docs/advanced_features/quantization.mdx b/docs/docs/advanced_features/quantization.mdx index 73d11444c..16d39f9c6 100644 --- a/docs/docs/advanced_features/quantization.mdx +++ b/docs/docs/advanced_features/quantization.mdx @@ -159,14 +159,14 @@ The following table summarizes quantization method support across NVIDIA and AMD Yes (SM80-SM90 via Marlin; SM100+ native FP4) No No - NVIDIA ModelOpt; use Marlin W4A16 fallback on Ampere/Hopper and native FP4 backends on Blackwell + NVIDIA ModelOpt; use Marlin W4A16 fallback on Ampere/Hopper and native FP4 backends on Blackwell; supports load-time BF16/FP16/FP8 MoE conversion with per-tensor FP32 activation scales nvfp4_online Yes (Blackwell/SM100 or SM103) No No - Online MoE-only NVFP4 weight quantization with runtime per-token activation scaling for BF16/FP16/FP8 checkpoints; requires flashinfer_trtllm or flashinfer_trtllm_routed + Online MoE-only NVFP4 weight quantization with per-token FP32 activation scales for BF16/FP16/FP8 checkpoints; use modelopt_fp4 for per-tensor FP32 activation scales; requires flashinfer_trtllm or flashinfer_trtllm_routed petit_nvfp4 @@ -851,18 +851,18 @@ python3 -m sglang.launch_server \ ### `nvfp4_online` online quantization method -Use `--quantization nvfp4_online` when you have a BF16, FP16, or FP8 MoE checkpoint and want SGLang to convert eligible MoE expert weights to NVFP4 while loading the model. This mode is for online conversion from higher-precision or FP8 checkpoints. It is not the serving path for already serialized NVFP4 checkpoints; use the existing ModelOpt FP4 path for those checkpoints. +Use `--quantization nvfp4_online` to convert eligible BF16, FP16, or FP8 MoE expert weights to NVFP4 at load time with per-token FP32 activation scales. Use `modelopt_fp4` for serialized NVFP4 checkpoints or the same load-time conversion with per-tensor FP32 activation scales. The design separates weight quantization from activation scaling: - **Weights:** SGLang quantizes each eligible MoE expert weight tensor as it is loaded, using standard 2D NVFP4 weight quantization. The generated NVFP4 weights use static E4M3 block scales plus static per-tensor FP32 scales derived from the weight amax. For gated MoE experts, the w1/w3 pair shares one per-tensor FP32 scale. -- **Activations:** FlashInfer computes activation FP32 scales dynamically per token at runtime. Because activations are scaled per token, this mode does not need calibrated static activation FP32 scales from the checkpoint. +- **Activations:** FlashInfer computes and propagates one FP32 scale per token at runtime. Backends that accept one per-tensor FP32 activation scale must use `modelopt_fp4`. - **FP8 checkpoints:** If an eligible expert weight is stored as FP8, SGLang first dequantizes that tensor with the checkpoint scale and then requantizes it to NVFP4 during loading. - **Other layers:** Dense linear layers stay in their source checkpoint precision or checkpoint quantization path. Only `--moe-runner-backend flashinfer_trtllm` and `--moe-runner-backend flashinfer_trtllm_routed` are supported. If `--moe-runner-backend` is omitted, SGLang selects `flashinfer_trtllm`. Tensor parallelism is supported; activation per-token scales are computed locally on each TP rank, while online weight quantization still uses the loaded expert tensor's per-tensor amax-derived FP32 scale. -FlashInfer TRTLLM MoE backends disable shared-expert fusion, so online quantization applies to routed MoE experts while shared experts stay in the checkpoint precision. To keep specific routed MoE layers out of online FP4 conversion, include their module path in `SGLANG_FP4_IGNORED_LAYERS`; FP8 checkpoints keep those ignored experts in FP8 checkpoint precision. +FlashInfer TRTLLM MoE backends disable shared-expert fusion, so online quantization applies to routed MoE experts while shared experts stay in the checkpoint precision. Both online modes honor `SGLANG_FP4_IGNORED_LAYERS`; for FP8 source checkpoints, listed experts remain FP8 instead of being converted to NVFP4. ```bash Command python3 -m sglang.launch_server \ diff --git a/docs/docs/references/environment_variables.mdx b/docs/docs/references/environment_variables.mdx index fb84b5c76..3af663055 100644 --- a/docs/docs/references/environment_variables.mdx +++ b/docs/docs/references/environment_variables.mdx @@ -761,7 +761,7 @@ SGLang supports various environment variables that can be used to configure its SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION - Enable FlashInfer TRTLLM NVFP4 per-token activation scaling; ignores checkpoint activation FP32 scale by treating it as 1 + Enable FlashInfer TRTLLM NVFP4 per-token activation scaling for serialized modelopt_fp4 checkpoints; checkpoint FP32 activation scales are treated as 1 false @@ -791,7 +791,7 @@ SGLang supports various environment variables that can be used to configure its SGLANG_FP4_IGNORED_LAYERS - A comma-separated list of layer names to keep out of FP4 online quantization, including nvfp4_online. For example: model.layers.40,model.layers.41. + A comma-separated list of layer names to keep out of online FP4 conversion for modelopt_fp4 or nvfp4_online. For example: model.layers.40,model.layers.41. "" diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 4b2172937..554ea487e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -267,12 +267,14 @@ class ModelConfig: disable_hybrid_swa_memory: bool = False, model_config_parser: str = "auto", speculative_algorithm: Optional[str] = None, + is_draft_quantization_explicit: bool = False, ) -> None: # Parse args self.model_path = model_path self.revision = revision self.quantization = quantization self.is_draft_model = is_draft_model + self.is_draft_quantization_explicit = is_draft_quantization_explicit self.speculative_algorithm = speculative_algorithm self.model_impl = model_impl self.sampling_defaults = sampling_defaults @@ -568,6 +570,10 @@ class ModelConfig: language_only=server_args.language_only, encoder_only=server_args.encoder_only, is_draft_model=is_draft_model, + is_draft_quantization_explicit=( + is_draft_model + and server_args._speculative_draft_quantization_explicitly_set + ), disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory, model_config_parser=server_args.model_config_parser, speculative_algorithm=server_args.speculative_algorithm, @@ -1425,7 +1431,9 @@ class ModelConfig: ] compatible_quantization_methods = { "modelopt_fp8": ["modelopt"], - "modelopt_fp4": ["modelopt"], + # Keep explicit or inherited modelopt_fp4 for literal FP8 checkpoints + # so eligible MoE experts are requantized online. + "modelopt_fp4": ["modelopt", "fp8"], "modelopt_mixed": ["modelopt"], "nvfp4_online": ["fp8"], "petit_nvfp4": ["modelopt"], @@ -1467,7 +1475,6 @@ class ModelConfig: and self.quantization == "nvfp4_online" and quant_method == "modelopt_fp4" ) - # Detect which checkpoint is it if not preserve_online_draft_quantization: for _, method in QUANTIZATION_METHODS.items(): diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index a30b6bd73..513ba0aaa 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -121,8 +121,7 @@ class DeepEPMoE(FusedMoE): elif ( get_moe_runner_backend().is_flashinfer_cutedsl() and quant_config is not None - and quant_config.get_name() - in ("modelopt_fp4", "modelopt_mixed", "nvfp4_online") + and quant_config.get_name() in ("modelopt_fp4", "modelopt_mixed") ): self.deprecate_flag = True elif ( diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index c99dad370..6422d82d1 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -16,8 +16,6 @@ from sglang.srt.layers.moe import ( MoeRunner, MoeRunnerBackend, MoeRunnerConfig, - get_deepep_mode, - get_moe_a2a_backend, get_moe_runner_backend, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo @@ -1212,7 +1210,18 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): class ModelOptFp4Config(ModelOptQuantConfig): - """Config class for FP4.""" + """Supported ModelOpt FP4 paths: + + - Serialized + per-tensor FP32 activation scales: load packed NVFP4 weights + and checkpoint-provided scales. + - Serialized + per-token FP32 activation scales: set + `SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION=1`; use + `flashinfer_trtllm` or `flashinfer_trtllm_routed`. + - BF16/FP16/FP8 MoE + per-tensor FP32 activation scales: quantize expert + weights on load, keep dense weights in source precision or FP8, and use + 1.0 when the checkpoint has no NVFP4 activation scale. + - BF16/FP16/FP8 MoE + per-token FP32 activation scales: use `nvfp4_online`. + """ def __init__( self, @@ -1233,11 +1242,20 @@ class ModelOptFp4Config(ModelOptQuantConfig): ) self.is_awq = is_awq self.group_size = group_size - self.use_per_token_activation = ( - envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() - if use_per_token_activation is None - else use_per_token_activation - ) + if not is_checkpoint_nvfp4_serialized: + if use_per_token_activation: + raise ValueError( + "Non-serialized modelopt_fp4 uses per-tensor FP32 " + "activation scales. Use nvfp4_online for online per-token " + "FP32 activation scales." + ) + self.use_per_token_activation = False + else: + self.use_per_token_activation = ( + envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get() + if use_per_token_activation is None + else use_per_token_activation + ) @classmethod def override_quantization_method(cls, hf_quant_config, user_quant): @@ -1248,6 +1266,18 @@ class ModelOptFp4Config(ModelOptQuantConfig): def get_name(cls) -> str: return "modelopt_fp4" + @classmethod + def for_online_weight_quantization( + cls, + packed_modules_mapping: Optional[Dict[str, List[str]]] = None, + ) -> QuantizationConfig: + """Use per-tensor FP32 activation scales for load-time MoE quantization.""" + from sglang.srt.layers.quantization.nvfp4_online import ( + make_modelopt_fp4_online_config, + ) + + return make_modelopt_fp4_online_config(packed_modules_mapping) + @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: return [torch.bfloat16, torch.half, torch.float8_e4m3fn] @@ -1290,13 +1320,21 @@ class ModelOptFp4Config(ModelOptQuantConfig): return next(iter(sizes)) @classmethod - def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config: + def from_config(cls, config: Dict[str, Any]) -> QuantizationConfig: # Handle two different config formats: # 1. hf_quant_config.json format: {"quantization": {"quant_algo": "NVFP4", ...}} # 2. config.json quantization_config format: {"quant_algo": "NVFP4", ...} # In future modelopt will deprecate hf_quant_config.json, and only keep config.json. # For legacy reasons, we keep hf_quant_config.json for now. + quant_method = str(config.get("quant_method", "")).lower() + if quant_method == "fp8": + from sglang.srt.layers.quantization.nvfp4_online import ( + make_modelopt_fp4_online_config_from_fp8, + ) + + return make_modelopt_fp4_online_config_from_fp8(config) + # Initialize variables kv_cache_quant_algo = None group_size = None @@ -1385,6 +1423,20 @@ class ModelOptFp4Config(ModelOptQuantConfig): ) def get_quant_method(self, layer: torch.nn.Module, prefix: str): + from sglang.srt.layers.linear import LinearBase + from sglang.srt.layers.moe.fused_moe_triton import FusedMoE + from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead + + if not self.is_checkpoint_nvfp4_serialized: + if isinstance(layer, (LinearBase, ParallelLMHead)): + # Load-time quantization applies only to MoE weights. + return UnquantizedLinearMethod() + if isinstance(layer, FusedMoE): + if self.is_layer_excluded(prefix): + return None + return ModelOptNvFp4FusedMoEMethod(self) + return None + return self._get_quant_method( layer, prefix, @@ -1503,6 +1555,7 @@ class ModelOptFp4LinearMethod(LinearMethodBase): input_scale = _make_per_tensor_scale_parameter( (len(output_partition_sizes),), weight_loader=weight_loader, + fill_value=1.0, needs_scalar_to_array=True, ) layer.register_parameter("input_scale", input_scale) @@ -1994,17 +2047,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): return get_moe_runner_backend().is_flashinfer_cutedsl() - @property - def supports_nvfp4_online_moe(self) -> bool: - a2a_backend = get_moe_a2a_backend() - return self.enable_flashinfer_trtllm_moe or ( - self.enable_flashinfer_cutedsl_moe - and ( - a2a_backend.is_flashinfer() - or (a2a_backend.is_deepep() and get_deepep_mode().is_low_latency()) - ) - ) - # ----- CuteDSL v1 vs v2 path helpers ----- # # "v1": cutedsl + deepep low-latency. @@ -2028,6 +2070,22 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): """CuteDSL v2 standard path (a2a=none or flashinfer, uses CuteDslMoEWrapper).""" return self.enable_flashinfer_cutedsl_moe and not self._is_cutedsl_v1_deepep + def prepare_weight_loader(self, layer, weight_loader): + if self.quant_config.is_checkpoint_nvfp4_serialized: + return weight_loader + from sglang.srt.layers.quantization.nvfp4_online import ( + make_nvfp4_online_weight_loader, + ) + + return make_nvfp4_online_weight_loader( + layer=layer, + original_weight_loader=weight_loader, + ) + + def _uses_serialized_fp8_source(self) -> bool: + # nvfp4_online overrides this for serialized FP8 source weights. + return False + def create_weights( self, layer: torch.nn.Module, @@ -2037,22 +2095,6 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ): - is_nvfp4_online = getattr(self.quant_config, "is_nvfp4_online", False) - if not self.quant_config.is_checkpoint_nvfp4_serialized and not is_nvfp4_online: - raise ValueError( - "NVFP4 quantization was selected, " - " dynamic quantization is not supported." - ) - # Online conversion changes only weight loading; downstream tensors use - # the same packed weight and scale layout as serialized ModelOpt NVFP4. - if is_nvfp4_online: - if not self.supports_nvfp4_online_moe: - raise ValueError( - "--quantization nvfp4_online supports flashinfer_trtllm, " - "flashinfer_trtllm_routed, or flashinfer_cutedsl with " - "FlashInfer A2A or DeepEP low_latency." - ) - # TODO(ch-wan): check if this is needed layer.intermediate_size_per_partition = intermediate_size_per_partition layer.params_dtype = params_dtype @@ -2060,9 +2102,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): weight_dtype = torch.uint8 weight_scale_dtype = torch.float8_e4m3fn - weight_loader = extra_weight_attrs.get("weight_loader") - if is_nvfp4_online: - weight_loader = self.get_online_weight_loader(layer, weight_loader) + weight_loader = self.prepare_weight_loader( + layer, extra_weight_attrs.get("weight_loader") + ) # GEMM 1 num_shards = 2 if layer.moe_runner_config.is_gated else 1 @@ -2161,7 +2203,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) - if is_nvfp4_online and self.quant_config.is_checkpoint_fp8_serialized: + if self._uses_serialized_fp8_source(): # FP8 checkpoints usually store expert scales as weight_scale_inv. # Online NVFP4 consumes them in the loader and writes the generated # NVFP4 scales into w*_weight_scale / w*_weight_scale_2 instead. @@ -2182,7 +2224,10 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value} ) - input_scale_fill = 1.0 if is_nvfp4_online else None + is_nvfp4_online = self.quant_config.get_name() == "nvfp4_online" + # nvfp4_online installs per-token activation scales after loading; + # per-tensor paths default to 1.0 here. + input_scale_fill = 1.0 if not is_nvfp4_online else None w13_input_scale = _make_per_tensor_scale_parameter( (layer.num_experts, num_shards), weight_loader=weight_loader, @@ -2200,10 +2245,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): layer.register_parameter("w2_input_scale", w2_input_scale) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - """Process FP4 MoE weights after loading from serialized checkpoint. - - Only supports pre-quantized checkpoints with FP8 weights and scales. - """ + """Transform packed FP4 MoE weights and scales for the selected backend.""" if getattr(layer, "inference_moe_w13_interleaved", False) and not getattr( layer, "_w13_deinterleaved", False ): diff --git a/python/sglang/srt/layers/quantization/nvfp4_online.py b/python/sglang/srt/layers/quantization/nvfp4_online.py index ce607386f..84f9d11c2 100644 --- a/python/sglang/srt/layers/quantization/nvfp4_online.py +++ b/python/sglang/srt/layers/quantization/nvfp4_online.py @@ -6,7 +6,7 @@ from __future__ import annotations import logging import re import threading -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional import torch @@ -16,6 +16,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( inverse_transform_scale_ue8m0, ) from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, ModelOptNvFp4FusedMoEMethod, ModelOptQuantConfig, ) @@ -29,12 +30,13 @@ logger = logging.getLogger(__name__) class NvFp4OnlineConfig(ModelOptQuantConfig): - """Config for `--quantization nvfp4_online`. + """Load-time NVFP4 with online per-token FP32 activation scales. - This mode is a load-time conversion path, not a serialized NVFP4 checkpoint - format. It reuses the ModelOpt NVFP4 MoE parameter layout and fills those - parameters by converting BF16/FP16/FP8 expert tensors as they are loaded. - Dense layers stay in the source checkpoint precision or quantization path. + `--quantization nvfp4_online` exclusively means online per-token FP32 + activation scaling. Use `modelopt_fp4` for per-tensor FP32 activation scales + or serialized NVFP4 checkpoints. This path converts BF16/FP16/FP8 MoE expert + weights as they load; dense layers retain their source precision or + quantization. """ # Marker consumed by the ModelOpt FP4 layout and the model loader. Serialized @@ -42,6 +44,7 @@ class NvFp4OnlineConfig(ModelOptQuantConfig): is_nvfp4_online = True is_checkpoint_nvfp4_serialized = False group_size = 16 + _use_per_token_activation = True @staticmethod def _normalize_ignored_layers( @@ -80,9 +83,9 @@ class NvFp4OnlineConfig(ModelOptQuantConfig): packed_modules_mapping=packed_modules_mapping or {}, ) self.fp4_ignored_layers = fp4_ignored_layers - # Weights use static NVFP4 scales, while FlashInfer computes activation - # FP32 scales dynamically per token at runtime. - self.use_per_token_activation = True + # NVFP4 weight scales are fixed at load time; FlashInfer computes one + # FP32 activation scale per token at runtime. + self.use_per_token_activation = self._use_per_token_activation self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized self.is_fp4_experts = False self.dequant_fp4_to_fp8 = False @@ -153,6 +156,39 @@ class NvFp4OnlineConfig(ModelOptQuantConfig): return None +class _ModelOptFp4OnlineConfig(NvFp4OnlineConfig): + """`modelopt_fp4` adapter for online per-tensor activation scaling.""" + + is_nvfp4_online = False + _use_per_token_activation = False + + @classmethod + def get_name(cls) -> str: + return "modelopt_fp4" + + @classmethod + def get_supported_act_dtypes(cls) -> List[torch.dtype]: + return ModelOptFp4Config.get_supported_act_dtypes() + + @classmethod + def get_min_capability(cls) -> int: + return ModelOptFp4Config.get_min_capability() + + +def make_modelopt_fp4_online_config( + packed_modules_mapping: Optional[Dict[str, List[str]]] = None, +) -> ModelOptQuantConfig: + return _ModelOptFp4OnlineConfig( + packed_modules_mapping=packed_modules_mapping, + ) + + +def make_modelopt_fp4_online_config_from_fp8( + config: Dict[str, Any], +) -> ModelOptQuantConfig: + return _ModelOptFp4OnlineConfig.from_config(config) + + class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): """MoE method that converts source expert weights to NVFP4 during loading.""" @@ -165,13 +201,34 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): if layer_match is not None else layer_prefix ) - if not self.supports_nvfp4_online_moe: + if ( + quant_config.use_per_token_activation + and not self.enable_flashinfer_trtllm_moe + ): raise ValueError( - "--quantization nvfp4_online supports flashinfer_trtllm, " - "flashinfer_trtllm_routed, or flashinfer_cutedsl with " - "FlashInfer A2A or DeepEP low_latency." + "--quantization nvfp4_online requires online per-token FP32 " + "activation scales and supports only flashinfer_trtllm or " + "flashinfer_trtllm_routed. Use --quantization modelopt_fp4 " + "for per-tensor FP32 activation scales." ) + def prepare_weight_loader(self, layer, weight_loader): + """Wrap the MoE weight loader with load-time NVFP4 conversion.""" + fp8_dequantizer = ( + self._dequantize_fp8_weight + if self.quant_config.is_checkpoint_fp8_serialized + else None + ) + return self.get_online_weight_loader( + layer, + weight_loader, + layer_log_name=self.layer_log_name, + fp8_dequantizer=fp8_dequantizer, + ) + + def _uses_serialized_fp8_source(self) -> bool: + return self.quant_config.is_checkpoint_fp8_serialized + @staticmethod def _quantize_weight_nvfp4( weight: torch.Tensor, @@ -187,18 +244,18 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): if weight.ndim != 2: raise ValueError( - "--quantization nvfp4_online expects 2D expert weights, " + "Online NVFP4 weight conversion expects 2D expert weights, " f"got shape {tuple(weight.shape)}." ) if not weight.is_floating_point(): raise ValueError( - "--quantization nvfp4_online expects floating-point source " + "Online NVFP4 weight conversion expects floating-point source " f"expert weights, got dtype {weight.dtype}. Serialized packed " "FP4 weights must use --quantization modelopt_fp4." ) if weight.shape[-1] % 16 != 0: raise ValueError( - "--quantization nvfp4_online requires expert weight K to be " + "Online NVFP4 weight conversion requires expert weight K to be " f"a multiple of 16, got shape {tuple(weight.shape)}." ) @@ -265,7 +322,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): ) -> torch.Tensor: if self.quant_config.use_mxfp8: raise ValueError( - "--quantization nvfp4_online does not support online " + "Online NVFP4 weight conversion does not support " "requantization from MXFP8 expert checkpoints." ) @@ -324,7 +381,15 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): return f"{prefix}weight_scale_2{suffix}" return f"{weight_name}.weight_scale_2" - def get_online_weight_loader(self, layer, original_weight_loader): + @classmethod + def get_online_weight_loader( + cls, + layer, + original_weight_loader, + *, + layer_log_name: str, + fp8_dequantizer: Optional[Callable], + ): """Wrap the normal MoE loader with load-time NVFP4 conversion. The wrapper quantizes each eligible expert shard as soon as the loader @@ -349,7 +414,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): return logger.info( "Running online NVFP4 quantization for MoE expert weights in %s.", - self.layer_log_name, + layer_log_name, ) did_log_quantization = True @@ -378,7 +443,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): original_weight_loader( scale_param, weight_scale, - weight_name=self._scale_weight_name(weight_name), + weight_name=cls._scale_weight_name(weight_name), shard_id=shard_id, expert_id=expert_id, ) @@ -390,7 +455,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): original_weight_loader( scale_2_param, weight_scale_2, - weight_name=self._scale_2_weight_name(weight_name), + weight_name=cls._scale_2_weight_name(weight_name), shard_id=shard_id, expert_id=expert_id, ) @@ -403,9 +468,9 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): expert_id: Optional[int], ) -> None: log_quantization_start() - if shard_id == "w2": + if shard_id == "w2" or not layer.moe_runner_config.is_gated: loaded_weight = loaded_weight.to(param.device) - fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4( + fp4_weight, weight_scale, weight_scale_2 = cls._quantize_weight_nvfp4( loaded_weight ) store_quantized_weight( @@ -442,7 +507,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): ) = pending if pending_shard_id == shard_id: raise ValueError( - "--quantization nvfp4_online expects paired w1/w3 expert " + "Online NVFP4 weight conversion expects paired w1/w3 expert " f"weights, got two {shard_id} tensors for expert {expert_id}." ) pending_weight = pending_weight.to(param.device) @@ -451,7 +516,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): loaded_rows = loaded_weight.shape[0] # Quantize the gated pair together so w1/w3 share one amax-derived # per-tensor FP32 scale, matching the serialized NVFP4 convention. - fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4( + fp4_weight, weight_scale, weight_scale_2 = cls._quantize_weight_nvfp4( torch.cat([pending_weight, loaded_weight], dim=0) ) pending_fp4_weight, loaded_fp4_weight = fp4_weight.split( @@ -486,14 +551,14 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): shard_id: str, expert_id: Optional[int], ) -> None: - if not self._is_fp8_weight(loaded_weight): + if not cls._is_fp8_weight(loaded_weight): process_loaded_weight( param, loaded_weight, weight_name, shard_id, expert_id ) return - if not self.quant_config.is_checkpoint_fp8_serialized: + if fp8_dequantizer is None: raise ValueError( - "--quantization nvfp4_online received an FP8 expert " + "Online NVFP4 weight conversion received an FP8 expert " "weight, but the checkpoint quantization config does not " "declare serialized FP8 weights." ) @@ -512,9 +577,7 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): return log_quantization_start() - loaded_weight = self._dequantize_fp8_weight( - loaded_weight, weight_scale, param.device - ) + loaded_weight = fp8_dequantizer(loaded_weight, weight_scale, param.device) process_loaded_weight( param, loaded_weight, weight_name, shard_id, expert_id ) @@ -539,7 +602,13 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): pending_shard_id, pending_eid, ) = pending - loaded_weight = self._dequantize_fp8_weight( + if fp8_dequantizer is None: + raise ValueError( + "Online NVFP4 weight conversion received an FP8 expert " + "weight, but the checkpoint quantization config does not " + "declare serialized FP8 weights." + ) + loaded_weight = fp8_dequantizer( pending_weight, loaded_weight, pending_param.device ) process_loaded_weight( @@ -566,10 +635,14 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): expert_id=expert_id, ) return - if self._should_skip_loaded_expert(layer, param, expert_id): + if cls._should_skip_loaded_expert(layer, param, expert_id): return - if self._is_fp8_weight_scale_name(weight_name): + # FP8 activation scales do not describe the requantized NVFP4 input. + if fp8_dequantizer is not None and "input_scale" in weight_name: + return + + if cls._is_fp8_weight_scale_name(weight_name): process_fp8_weight_scale(loaded_weight, shard_id, expert_id) return @@ -588,3 +661,17 @@ class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod): ) return nvfp4_online_weight_loader + + +def make_nvfp4_online_weight_loader( + *, + layer: torch.nn.Module, + original_weight_loader: Callable, +) -> Callable: + """Wrap a MoE weight loader with load-time NVFP4 conversion.""" + return ModelOptNvFp4OnlineFusedMoEMethod.get_online_weight_loader( + layer, + original_weight_loader, + layer_log_name="MoE layer", + fp8_dequantizer=None, + ) diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 88662bfee..74168210e 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -835,6 +835,11 @@ class DefaultModelLoader(BaseModelLoader): quant_config = getattr(model, "quant_config", None) is_nvfp4_online = getattr(quant_config, "is_nvfp4_online", False) + is_modelopt_fp4_online = ( + quant_config is not None + and quant_config.get_name() == "modelopt_fp4" + and not quant_config.is_checkpoint_nvfp4_serialized + ) is_mxfp8 = quant_config is not None and quant_config.get_name() == "mxfp8" if is_mxfp8: weights = ( @@ -845,7 +850,7 @@ class DefaultModelLoader(BaseModelLoader): for name, loaded_weight in weights ) - if is_nvfp4_online: + if is_nvfp4_online or is_modelopt_fp4_online: # Scope exact FP4 quantization math to load-time conversion only; # restore the original environment before serving starts. with temp_set_env( @@ -4108,12 +4113,29 @@ def get_model_loader( logger.info("Using IncModelLoader due to AutoRound quantization config.") return IncModelLoader(load_config) - # ModelOptModelLoader's local-copy quantize-and-export workflow doesn't apply - # to non-local loaders. These loaders own their weight transport path and still - # initialize the model with ModelOpt quantization config where applicable. - model_optloader_allowed = model_config and load_config.load_format not in ( - LoadFormat.RUNAI_STREAMER, - LoadFormat.REMOTE_INSTANCE, + modelopt_config = load_config.modelopt_config + modelopt_workflow_requested = modelopt_config is not None and any( + ( + modelopt_config.checkpoint_restore_path, + modelopt_config.checkpoint_save_path, + modelopt_config.export_path, + ) + ) + + # Online modelopt_fp4 converts weights through DefaultModelLoader unless the + # caller explicitly requests ModelOpt calibration/checkpoint/export work. + # Non-local loaders still own their weight transport path. + modelopt_fp4_online = ( + model_config + and model_config.quantization == "modelopt_fp4" + and not model_config._is_already_quantized() + and not modelopt_workflow_requested + ) + model_optloader_allowed = ( + model_config + and not modelopt_fp4_online + and load_config.load_format + not in (LoadFormat.RUNAI_STREAMER, LoadFormat.REMOTE_INSTANCE) ) if model_optloader_allowed and ( diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index 506d359a4..5ff5bafa6 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -234,6 +234,27 @@ class DisabledTqdm(tqdm): super().__init__(*args, **kwargs) +def _resolve_explicit_draft_quant_config( + model_config: ModelConfig, + quant_config: QuantizationConfig, +) -> QuantizationConfig: + if not ( + model_config.is_draft_model and model_config.is_draft_quantization_explicit + ): + return quant_config + + if model_config.quantization == "modelopt_fp4" and ( + isinstance(quant_config, ModelOptFp4Config) + and quant_config.is_checkpoint_nvfp4_serialized + and quant_config.is_layer_excluded("mtp.layers.0.mlp.experts") + ): + return ModelOptFp4Config.for_online_weight_quantization( + quant_config.packed_modules_mapping + ) + + return quant_config + + # TODO(woosuk): Move this to other place. def get_quant_config( model_config: ModelConfig, @@ -282,7 +303,9 @@ def get_quant_config( if model_config.quantization in REQUANTIZATION_METHODS: hf_quant_config["requantization_method"] = model_config.quantization - return quant_cls.from_config(hf_quant_config) + return _resolve_explicit_draft_quant_config( + model_config, quant_cls.from_config(hf_quant_config) + ) # In case of bitsandbytes/QLoRA, get quant config from the adapter model. if model_config.quantization == "bitsandbytes": @@ -332,6 +355,12 @@ def get_quant_config( f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) ] if len(quant_config_files) == 0: + if model_config.quantization == "modelopt_fp4": + # Without serialized metadata, quantize MoE expert weights online; + # leave dense layers in source precision. + return ModelOptFp4Config.for_online_weight_quantization( + packed_modules_mapping + ) raise ValueError(f"Cannot find the config file for {model_config.quantization}") if len(quant_config_files) > 1: raise ValueError( @@ -367,8 +396,12 @@ def get_quant_config( elif quant_algo == "FP8" or model_config.quantization == "modelopt_fp8": return ModelOptFp8Config.from_config(config) elif "FP4" in quant_algo: - return ModelOptFp4Config.from_config(config) - return quant_cls.from_config(config) + return _resolve_explicit_draft_quant_config( + model_config, ModelOptFp4Config.from_config(config) + ) + return _resolve_explicit_draft_quant_config( + model_config, quant_cls.from_config(config) + ) def _check_index_files_exist(snapshot_dir: str) -> Tuple[bool, Optional[str]]: diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index ded5ef3b8..fd146e1a6 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -61,10 +61,15 @@ class Qwen3_5ForCausalLMMTP(nn.Module): # Deep-copy so MTP mutations below don't leak into the target's config. config = copy.deepcopy(config) - # The MTP model is unquantized in the nvfp4 checkpoint. - if quant_config and quant_config.get_name() in ( - "modelopt_fp4", - "modelopt_mixed", + # 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" + or ( + quant_config.get_name() == "modelopt_fp4" + and quant_config.is_checkpoint_nvfp4_serialized + ) ): quant_config = None if is_npu() and get_spec().speculative_draft_model_quantization is None: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 849bff48a..c166bb21c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2165,6 +2165,12 @@ class ServerArgs: ), NS("spec"), ] = None + # Internal provenance used after the public draft quantization inherits the + # target value. It is a dataclass field so ServerArgs round-trips preserve + # whether the user explicitly set the draft option; it has no CLI surface. + _speculative_draft_quantization_explicitly_set: A[ + Optional[bool], Arg(no_cli=True), NS("spec") + ] = None speculative_skip_dp_mlp_sync: A[ bool, "Skip the extra MLP sync that the scheduler performs before merging a new batch when speculative decoding + DP attention are both enabled.", @@ -4073,6 +4079,10 @@ class ServerArgs: # In speculative scenario: # - If `speculative_draft_model_quantization` is specified, the draft model uses this quantization method. # - Otherwise, the draft model defaults to the same quantization as the target model. + if self._speculative_draft_quantization_explicitly_set is None: + self._speculative_draft_quantization_explicitly_set = ( + self.speculative_draft_model_quantization is not None + ) if self.speculative_draft_model_quantization is None: self.speculative_draft_model_quantization = self.quantization diff --git a/test/registered/models_e2e/test_qwen35_fp4_mtp.py b/test/registered/models_e2e/test_qwen35_fp4_mtp.py index 3428af1af..9628b40f3 100644 --- a/test/registered/models_e2e/test_qwen35_fp4_mtp.py +++ b/test/registered/models_e2e/test_qwen35_fp4_mtp.py @@ -38,6 +38,8 @@ MTP_BASE_ARGS = [ "trtllm_mha", "--quantization", "modelopt_fp4", + "--speculative-draft-model-quantization", + "modelopt_fp4", "--speculative-algorithm", "NEXTN", "--speculative-num-steps", diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py index 7773928c7..cb2f4bdd0 100644 --- a/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4.py @@ -1,9 +1,16 @@ import unittest +from unittest.mock import patch import torch +import torch.nn as nn from sglang.srt.layers.linear import MergedColumnParallelLinear, QKVParallelLinear from sglang.srt.layers.parameter import PerTensorScaleParameter +from sglang.srt.layers.quantization.modelopt_quant import ( + ModelOptFp4Config, + ModelOptFp4LinearMethod, +) +from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -75,6 +82,56 @@ class TestModelOptNvfp4(CustomTestCase): torch.testing.assert_close(scale, torch.tensor([0.25, 0.5])) + def test_missing_input_scale_defaults_to_one_and_checkpoint_overwrites(self): + config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + use_per_token_activation=False, + ) + layer = nn.Module() + ModelOptFp4LinearMethod(config).create_weights( + layer, + input_size_per_partition=16, + output_partition_sizes=[16], + input_size=16, + output_size=16, + params_dtype=torch.bfloat16, + weight_loader=default_weight_loader, + ) + + torch.testing.assert_close(layer.input_scale, torch.ones(1)) + default_weight_loader(layer.input_scale, torch.tensor(0.25)) + torch.testing.assert_close(layer.input_scale, torch.tensor([0.25])) + + @patch( + "sglang.srt.layers.quantization.modelopt_quant.envs." + "SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get", + return_value=True, + ) + def test_modelopt_fp4_per_token_activation_contract(self, _): + # Serialized ModelOpt FP4 retains the existing environment-controlled + # per-token activation path. + serialized_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + ) + # Online modelopt_fp4 always uses per-tensor activation scaling, even + # when the serialized-checkpoint environment switch is enabled. + online_config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=False, + group_size=16, + ) + + self.assertTrue(serialized_config.use_per_token_activation) + self.assertFalse(online_config.use_per_token_activation) + # nvfp4_online is the public interface for online per-token scaling. + with self.assertRaisesRegex(ValueError, "Use nvfp4_online"): + ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=False, + group_size=16, + use_per_token_activation=True, + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_loader/test_modelopt_loader.py b/test/registered/unit/model_loader/test_modelopt_loader.py index 69bfcc706..b6fc958ef 100644 --- a/test/registered/unit/model_loader/test_modelopt_loader.py +++ b/test/registered/unit/model_loader/test_modelopt_loader.py @@ -6,6 +6,7 @@ applies NVIDIA Model Optimizer quantization to models during loading. """ import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch import torch @@ -25,7 +26,12 @@ from sglang.srt.layers.quantization.modelopt_quant import ( ModelOptMixedPrecisionConfig, ModelOptNvFp4A16LinearMethod, ) -from sglang.srt.model_loader.loader import DefaultModelLoader, ModelOptModelLoader +from sglang.srt.model_loader.loader import ( + DefaultModelLoader, + ModelOptModelLoader, + get_model_loader, +) +from sglang.srt.model_loader.weight_utils import get_quant_config from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM from sglang.srt.models.utils import WeightsMapper from sglang.srt.utils import get_device @@ -605,6 +611,86 @@ class TestParseQuantHfConfig(CustomTestCase): self.assertEqual(result["quant_method"], "gptq") self.assertNotIn("quant_algo", result) + def test_inherited_draft_modelopt_fp4_accepts_fp8_checkpoint(self): + # ServerArgs has already copied the target's modelopt_fp4 request to the + # draft. Compatible FP8 metadata must not replace it with plain fp8. + self.model_config.quantization = "modelopt_fp4" + self.model_config.is_draft_model = True + self.model_config.is_draft_quantization_explicit = False + with ( + patch.object( + self.model_config, + "_parse_quant_hf_config", + return_value={"quant_method": "fp8"}, + ), + patch.object( + self.model_config, + "_find_quant_modelslim_config", + return_value=None, + ), + ): + self.model_config._verify_quantization() + + # Keeping modelopt_fp4 selects online FP8-to-NVFP4 conversion for + # eligible MoE experts; this test stops at quantization-method routing. + self.assertEqual(self.model_config.quantization, "modelopt_fp4") + + +class TestModelOptFp4LoaderSelection(CustomTestCase): + def test_draft_modelopt_fp4_uses_checkpoint_exclusions(self): + cases = ( + # Excluded MTP experts are unpacked, so an explicit draft request + # replaces the serialized config with online weight quantization. + ("explicit embedded draft", True, ["mtp.layers.0*"], False), + # MTP experts present in the serialized checkpoint stay serialized. + ("explicit serialized draft", True, [], True), + # Inherited target quantization does not override draft exclusions. + ("inherited embedded draft", False, ["mtp.layers.0*"], True), + ) + for name, is_explicit, ignored_layers, is_serialized in cases: + with self.subTest(name=name): + model_config = SimpleNamespace( + model_path="target-model", + quantization="modelopt_fp4", + is_draft_model=True, + is_draft_quantization_explicit=is_explicit, + hf_config=SimpleNamespace( + quantization_config={ + "quant_algo": "NVFP4", + "group_size": 16, + "ignore": ignored_layers, + } + ), + ) + + config = get_quant_config(model_config, LoadConfig(), {}) + + self.assertEqual(config.get_name(), "modelopt_fp4") + self.assertEqual(config.is_checkpoint_nvfp4_serialized, is_serialized) + + def test_unquantized_modelopt_fp4_preserves_modelopt_workflows(self): + model_config = SimpleNamespace( + quantization="modelopt_fp4", + _is_already_quantized=lambda: False, + ) + + # Online conversion runs through the regular per-layer weight loaders. + online_loader = get_model_loader(LoadConfig(), model_config) + self.assertIsInstance(online_loader, DefaultModelLoader) + self.assertNotIsInstance(online_loader, ModelOptModelLoader) + + # Explicit ModelOpt checkpoint/export workflows still need its loader. + for option in ( + "modelopt_checkpoint_restore_path", + "modelopt_checkpoint_save_path", + "modelopt_export_path", + ): + with self.subTest(option=option): + loader = get_model_loader( + LoadConfig(**{option: "/tmp/modelopt"}), model_config + ) + self.assertIsInstance(loader, ModelOptModelLoader) + class TestModelOptMixedPrecisionConfig(CustomTestCase): def test_minimax_mixed_precision_resolves_runtime_names_and_mxfp8(self): @@ -734,7 +820,11 @@ class TestModelOptMixedPrecisionConfig(CustomTestCase): return_value=True, ) def test_explicit_nvfp4_per_token_activation_false_overrides_env(self, _): - config = ModelOptFp4Config(use_per_token_activation=False) + config = ModelOptFp4Config( + is_checkpoint_nvfp4_serialized=True, + group_size=16, + use_per_token_activation=False, + ) self.assertFalse(config.use_per_token_activation) diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 5a8f44223..6fbdbeee9 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1,3 +1,4 @@ +import dataclasses import importlib import json import os @@ -81,6 +82,17 @@ class TestPrepareServerArgs(CustomTestCase): return_hidden_states_mode="lst", ) + def test_draft_quantization_explicitness_survives_asdict_round_trip(self): + inherited = ServerArgs(model_path="dummy", quantization="modelopt_fp4") + inherited._handle_missing_default_values() + self.assertEqual(inherited.speculative_draft_model_quantization, "modelopt_fp4") + self.assertFalse(inherited._speculative_draft_quantization_explicitly_set) + + reconstructed = ServerArgs(**dataclasses.asdict(inherited)) + reconstructed._handle_missing_default_values() + + self.assertFalse(reconstructed._speculative_draft_quantization_explicitly_set) + def test_config_nested_dict_args_are_json(self): with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: f.write("mm-process-config:\n image:\n resize: 128\n")