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

[ModelOpt FP4] Support online MoE weight quantization (#33115)

Browse Source 
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
This commit is contained in:
Ziang Li
2026-08-06 11:01:55 -07:00
committed by GitHub
co-authored by Brayden Zhong
parent 05c7ebf64c
commit 4ad990ba7d
14 changed files with 472 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+9 -2
View File
@@ -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():
python/sglang/srt/layers/moe/ep_moe/layer.py
+1 -2
View File
@@ -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 (
python/sglang/srt/layers/quantization/modelopt_quant.py
+87 -45
View File
@@ -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
):
python/sglang/srt/layers/quantization/nvfp4_online.py
+121 -34
View File
@@ -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,
)
python/sglang/srt/model_loader/loader.py
+29 -7
View File
@@ -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 (
python/sglang/srt/model_loader/weight_utils.py
+36 -3
View File
@@ -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]]:
python/sglang/srt/models/qwen3_5_mtp.py
+9 -4
View File
@@ -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:
python/sglang/srt/server_args.py
+10
View File
@@ -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