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

Implement online nvfp4 quantization (#26083)

Browse Source
This commit is contained in:
Ziang Li
2026-06-10 00:26:51 -07:00
committed by GitHub
parent e76e4959b5
commit 01f10acd06
12 changed files with 783 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/configs/model_config.py
+3
View File
@@ -1120,6 +1120,7 @@ class ModelConfig:
"modelopt",
"modelopt_fp8",
"modelopt_fp4",
"nvfp4_online",
"modelopt_mixed",
]
modelopt_quantization_specified = (
@@ -1164,6 +1165,7 @@ class ModelConfig:
"modelopt_fp8",
"modelopt_fp4",
"modelopt_mixed",
"nvfp4_online",
"gptq_marlin_24",
"gptq_marlin",
"awq_marlin",
@@ -1185,6 +1187,7 @@ class ModelConfig:
"modelopt_fp8": ["modelopt"],
"modelopt_fp4": ["modelopt"],
"modelopt_mixed": ["modelopt"],
"nvfp4_online": ["fp8"],
"petit_nvfp4": ["modelopt"],
"w8a8_int8": ["compressed-tensors", "compressed_tensors"],
"w8a8_fp8": ["compressed-tensors", "compressed_tensors"],
python/sglang/srt/environ.py
+2 -1
View File
@@ -494,13 +494,14 @@ class Envs:
SGLANG_NVFP4_CKPT_FP8_NEXTN_MOE = EnvBool(False)
SGLANG_QUANT_ALLOW_DOWNCASTING = EnvBool(False)
SGLANG_FP8_IGNORED_LAYERS = EnvStr("")
SGLANG_FP4_IGNORED_LAYERS = EnvStr("")
# Flashinfer
SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True)
SGLANG_FLASHINFER_USE_PAGED = EnvBool(False)
# Default to the pick from flashinfer
SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024)
# Enable per-token NVFP4 activation scaling path for FlashInfer TRT-LLM MoE.
# Enable NVFP4 per-token activation scaling path for FlashInfer TRT-LLM MoE.
SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION = EnvBool(False)
# SGLang needs to know FlashInfer NVFP4 4over6 config to compute the global scale factor.
FLASHINFER_NVFP4_4OVER6 = EnvBool(False)
python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py
+2 -1
View File
@@ -812,6 +812,7 @@ class FlashInferTrtllmFp4MoeQuantInfo(MoeQuantInfo):
intermediate_size_per_partition: int
routing_method_type: int
use_per_token_activation: bool = False
def quantize_hidden_states_fp4(
@@ -875,7 +876,7 @@ def fused_experts_none_to_flashinfer_trtllm_fp4(
topk_output = dispatch_output.topk_output
# Quantize hidden states to FP4
if envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get():
if quant_info.use_per_token_activation:
from flashinfer import SfLayout, nvfp4_quantize
e4m3_max = 448.0
python/sglang/srt/layers/quantization/__init__.py
+2
View File
@@ -44,6 +44,7 @@ from sglang.srt.layers.quantization.modelopt_quant import (
from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig
from sglang.srt.layers.quantization.moe_wna16 import MoeWNA16Config
from sglang.srt.layers.quantization.mxfp4 import Mxfp4Config
from sglang.srt.layers.quantization.nvfp4_online import NvFp4OnlineConfig
from sglang.srt.layers.quantization.petit import PetitNvFp4Config
from sglang.srt.layers.quantization.qoq import QoQConfig
from sglang.srt.layers.quantization.quark.quark import QuarkConfig
@@ -74,6 +75,7 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
"modelopt": ModelOptFp8Config, # Auto-detect, defaults to FP8
"modelopt_fp8": ModelOptFp8Config,
"modelopt_fp4": ModelOptFp4Config,
"nvfp4_online": NvFp4OnlineConfig,
"modelopt_mixed": ModelOptMixedPrecisionConfig,
"w8a8_int8": W8A8Int8Config,
"w8a8_fp8": W8A8Fp8Config,
python/sglang/srt/layers/quantization/modelopt_quant.py
+42 -5
View File
@@ -287,6 +287,7 @@ class ModelOptQuantConfig(QuantizationConfig):
self.packed_modules_mapping = packed_modules_mapping
self.exclude_modules = exclude_modules or []
self.kv_cache_quant_algo = kv_cache_quant_algo
self.use_per_token_activation = False
def _get_quant_method(
self,
@@ -1177,6 +1178,7 @@ class ModelOptFp4Config(ModelOptQuantConfig):
group_size: int = None,
exclude_modules: List[str] = None,
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
use_per_token_activation: Optional[bool] = None,
) -> None:
super().__init__(kv_cache_quant_algo, exclude_modules, packed_modules_mapping)
self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
@@ -1186,6 +1188,10 @@ class ModelOptFp4Config(ModelOptQuantConfig):
"format is experimental and subject to change."
)
self.group_size = group_size
self.use_per_token_activation = (
use_per_token_activation
or envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get()
)
@classmethod
def override_quantization_method(cls, hf_quant_config, user_quant):
@@ -1746,11 +1752,23 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
params_dtype: torch.dtype,
**extra_weight_attrs,
):
if not self.quant_config.is_checkpoint_nvfp4_serialized:
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."
)
# `nvfp4_online` is not a serialized checkpoint format, but after the
# online loader converts each expert it uses the same packed NVFP4
# weights, block scales, and per-tensor scales as serialized ModelOpt
# NVFP4 checkpoints. Reuse this layout and swap only the weight loader.
if is_nvfp4_online:
if not self.enable_flashinfer_trtllm_moe:
raise ValueError(
"--quantization nvfp4_online supports only "
"--moe-runner-backend flashinfer_trtllm or "
"flashinfer_trtllm_routed."
)
# TODO(ch-wan): check if this is needed
layer.intermediate_size_per_partition = intermediate_size_per_partition
@@ -1760,6 +1778,8 @@ 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)
# GEMM 1
num_shards = 2 if layer.moe_runner_config.is_gated else 1
@@ -1858,6 +1878,23 @@ 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:
# 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.
w13_source_weight_scale_inv = PerTensorScaleParameter(
data=torch.empty(0, dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter(
"w13_weight_scale_inv", w13_source_weight_scale_inv
)
w2_source_weight_scale_inv = PerTensorScaleParameter(
data=torch.empty(0, dtype=torch.float32),
weight_loader=weight_loader,
)
layer.register_parameter("w2_weight_scale_inv", w2_source_weight_scale_inv)
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
)
@@ -1945,10 +1982,9 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32)
w2_input_scale = layer.w2_input_scale
if (
self.enable_flashinfer_trtllm_moe
and envs.SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION.get()
):
if self.quant_config.use_per_token_activation:
# FlashInfer computes activation scales dynamically per token, so
# the static checkpoint activation scale is intentionally neutral.
w13_input_scale = torch.ones_like(w13_input_scale, dtype=torch.float32)
w2_input_scale = torch.ones_like(w2_input_scale, dtype=torch.float32)
@@ -2273,6 +2309,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase):
local_num_experts=layer.num_local_experts,
intermediate_size_per_partition=layer.intermediate_size_per_partition,
routing_method_type=routing_method_type,
use_per_token_activation=self.quant_config.use_per_token_activation,
)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/layers/quantization/nvfp4_online.py
+583
View File
@@ -0,0 +1,583 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import logging
import re
import threading
from typing import Any, Dict, List, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.quantization.fp8_utils import (
block_quant_dequant,
inverse_transform_scale_ue8m0,
)
from sglang.srt.layers.quantization.modelopt_quant import (
ModelOptNvFp4FusedMoEMethod,
ModelOptQuantConfig,
)
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.layers.quantization.utils import (
is_layer_skipped,
per_tensor_dequantize,
)
logger = logging.getLogger(__name__)
class NvFp4OnlineConfig(ModelOptQuantConfig):
"""Config for `--quantization nvfp4_online`.
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.
"""
# Marker consumed by the ModelOpt FP4 layout and the model loader. Serialized
# NVFP4 checkpoints use ModelOptFp4Config instead.
is_nvfp4_online = True
is_checkpoint_nvfp4_serialized = False
group_size = 16
@staticmethod
def _normalize_ignored_layers(
ignored_layers: Optional[List[str]],
) -> List[str]:
if not ignored_layers:
return []
normalized_ignored_layers = []
for layer in ignored_layers:
base = layer.removeprefix("model.")
normalized_ignored_layers.append(base)
normalized_ignored_layers.append(f"model.{base}")
return list(dict.fromkeys(normalized_ignored_layers))
def __init__(
self,
exclude_modules: Optional[List[str]] = None,
packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
is_checkpoint_fp8_serialized: bool = False,
activation_scheme: str = "dynamic",
weight_block_size: Optional[List[int]] = None,
use_mxfp8: bool = False,
) -> None:
source_ignored_layers = self._normalize_ignored_layers(exclude_modules)
fp4_ignored_layers = list(source_ignored_layers)
if ignored_layers_str := envs.SGLANG_FP4_IGNORED_LAYERS.get():
fp4_ignored_layers.extend(
layer.strip()
for layer in ignored_layers_str.split(",")
if layer.strip()
)
fp4_ignored_layers = self._normalize_ignored_layers(fp4_ignored_layers)
super().__init__(
kv_cache_quant_algo=None,
exclude_modules=source_ignored_layers,
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
self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
self.is_fp4_experts = False
self.activation_scheme = activation_scheme
self.weight_block_size = weight_block_size
self.use_mxfp8 = use_mxfp8
@classmethod
def get_name(cls) -> str:
return "nvfp4_online"
@classmethod
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
return [torch.bfloat16, torch.half]
@classmethod
def get_min_capability(cls) -> int:
return 100
@classmethod
def get_config_filenames(cls) -> List[str]:
return []
@classmethod
def from_config(cls, config: Dict[str, Any]) -> NvFp4OnlineConfig:
quant_method = str(config.get("quant_method", "")).lower()
use_mxfp8 = "mxfp8" in quant_method
is_checkpoint_fp8_serialized = "fp8" in quant_method or use_mxfp8
ignored_layers = config.get("ignored_layers") or config.get(
"modules_to_not_convert"
)
if isinstance(ignored_layers, str):
ignored_layers = [ignored_layers]
return cls(
exclude_modules=ignored_layers,
packed_modules_mapping=config.get("packed_modules_mapping"),
is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized,
activation_scheme=config.get("activation_scheme", "dynamic"),
weight_block_size=config.get("weight_block_size"),
use_mxfp8=use_mxfp8,
)
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.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
if isinstance(layer, LinearBase):
if is_layer_skipped(
prefix, self.exclude_modules, self.packed_modules_mapping
) or self.is_layer_excluded(prefix):
return UnquantizedLinearMethod()
if self.is_checkpoint_fp8_serialized:
return Fp8LinearMethod(self)
return UnquantizedLinearMethod()
if isinstance(layer, FusedMoE):
if is_layer_skipped(
prefix, self.exclude_modules, self.packed_modules_mapping
) or self.is_layer_excluded(prefix):
return None
if is_layer_skipped(
prefix, self.fp4_ignored_layers, self.packed_modules_mapping
):
if self.is_checkpoint_fp8_serialized:
return Fp8MoEMethod(self)
return None
return ModelOptNvFp4OnlineFusedMoEMethod(self, prefix)
return None
class ModelOptNvFp4OnlineFusedMoEMethod(ModelOptNvFp4FusedMoEMethod):
"""MoE method that converts source expert weights to NVFP4 during loading."""
def __init__(self, quant_config: NvFp4OnlineConfig, layer_prefix: str):
super().__init__(quant_config)
self.layer_prefix = layer_prefix
layer_match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", layer_prefix)
self.layer_log_name = (
f"layer {layer_match.group(1)} ({layer_prefix})"
if layer_match is not None
else layer_prefix
)
if not self.enable_flashinfer_trtllm_moe:
raise ValueError(
"--quantization nvfp4_online supports only "
"--moe-runner-backend flashinfer_trtllm or "
"flashinfer_trtllm_routed."
)
@staticmethod
def _quantize_weight_nvfp4(
weight: torch.Tensor,
weight_scale_2: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Return packed NVFP4 weight, block scales, and per-tensor decode scale.
The weight scale is static and per tensor. Callers pass an existing
scale when multiple shards must share one global scale, for example the
gated w1/w3 pair.
"""
from flashinfer import SfLayout, nvfp4_quantize
if weight.ndim != 2:
raise ValueError(
"--quantization nvfp4_online expects 2D expert weights, "
f"got shape {tuple(weight.shape)}."
)
if weight.shape[-1] % 16 != 0:
raise ValueError(
"--quantization nvfp4_online requires expert weight K to be "
f"a multiple of 16, got shape {tuple(weight.shape)}."
)
if weight_scale_2 is None:
# weight_scale_2 is the NVFP4 decode scale. FlashInfer consumes its
# reciprocal as the global encode scale, matching 448 * 6 / amax.
weight_amax = (
weight.abs()
.nan_to_num()
.amax()
.to(device=weight.device, dtype=torch.float32)
)
e4m3_max = (
256.0
if envs.FLASHINFER_NVFP4_4OVER6.get()
and envs.FLASHINFER_NVFP4_4OVER6_E4M3_USE_256.get()
else float(torch.finfo(torch.float8_e4m3fn).max)
)
fp8_fp4_max = e4m3_max * 6.0
weight_scale_2 = torch.where(
weight_amax > 0,
weight_amax / fp8_fp4_max,
torch.ones_like(weight_amax),
)
else:
weight_scale_2 = weight_scale_2.to(
device=weight.device, dtype=torch.float32
)
fp4_weight, weight_sf = nvfp4_quantize(
weight.contiguous(),
1.0 / weight_scale_2,
sfLayout=SfLayout.layout_linear,
backend="cuda",
)
rows, cols = weight.shape
weight_sf = weight_sf.view(torch.float8_e4m3fn).reshape(rows, cols // 16)
return (
fp4_weight.reshape(rows, cols // 2),
weight_sf.contiguous(),
weight_scale_2,
)
@staticmethod
def _is_fp8_weight(weight: torch.Tensor) -> bool:
fp8_dtypes = {
dtype
for dtype in (
getattr(torch, "float8_e4m3fn", None),
getattr(torch, "float8_e5m2", None),
)
if dtype is not None
}
return weight.dtype in fp8_dtypes
@staticmethod
def _is_fp8_weight_scale_name(weight_name: str) -> bool:
return "weight_scale" in weight_name and "weight_scale_2" not in weight_name
def _dequantize_fp8_weight(
self,
weight: torch.Tensor,
weight_scale: torch.Tensor,
device: torch.device,
) -> torch.Tensor:
if self.quant_config.use_mxfp8:
raise ValueError(
"--quantization nvfp4_online does not support online "
"requantization from MXFP8 expert checkpoints."
)
weight = weight.to(device).contiguous()
weight_scale = weight_scale.to(device=device).contiguous()
if weight_scale.dtype == torch.int32:
weight_scale = inverse_transform_scale_ue8m0(
weight_scale, mn=weight.shape[-2]
)
weight_scale = weight_scale.to(dtype=torch.float32).contiguous()
if weight_scale.numel() == 1 or self.quant_config.weight_block_size is None:
return (
per_tensor_dequantize(weight, weight_scale)
.to(torch.bfloat16)
.contiguous()
)
return block_quant_dequant(
weight,
weight_scale,
self.quant_config.weight_block_size,
torch.bfloat16,
).contiguous()
@staticmethod
def _should_skip_loaded_expert(
layer: torch.nn.Module,
param: torch.nn.Parameter,
expert_id: Optional[int],
) -> bool:
if expert_id is None:
return False
if getattr(param, "_sglang_require_global_experts", False):
return False
# With EPLB or explicit expert placement, logical expert IDs can map to
# one or more physical experts. Let the canonical MoE loader do that
# mapping instead of pre-skipping from the trivial EP layout.
from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
if get_global_expert_location_metadata() is not None:
return False
return layer._map_global_expert_id_to_local_expert_id(expert_id) == -1
@staticmethod
def _scale_weight_name(weight_name: str) -> str:
if "weight" in weight_name:
prefix, suffix = weight_name.rsplit("weight", 1)
return f"{prefix}weight_scale{suffix}"
return f"{weight_name}.weight_scale"
@staticmethod
def _scale_2_weight_name(weight_name: str) -> str:
if "weight" in weight_name:
prefix, suffix = weight_name.rsplit("weight", 1)
return f"{prefix}weight_scale_2{suffix}"
return f"{weight_name}.weight_scale_2"
def get_online_weight_loader(self, layer, original_weight_loader):
"""Wrap the normal MoE loader with load-time NVFP4 conversion.
The wrapper quantizes each eligible expert shard as soon as the loader
sees enough source data, which avoids materializing and then converting
the full checkpoint. FP8 checkpoints stream weight and scale tensors
separately, so those pairs are staged until both sides have arrived.
"""
pending_w13_weights = {}
pending_w13_lock = threading.Lock()
pending_fp8_weights = {}
pending_fp8_weight_scales = {}
pending_fp8_lock = threading.Lock()
quantization_log_lock = threading.Lock()
did_log_quantization = False
def log_quantization_start() -> None:
nonlocal did_log_quantization
if did_log_quantization:
return
with quantization_log_lock:
if did_log_quantization:
return
logger.info(
"Running online NVFP4 quantization for MoE expert weights in %s.",
self.layer_log_name,
)
did_log_quantization = True
def store_quantized_weight(
param: torch.nn.Parameter,
fp4_weight: torch.Tensor,
weight_scale: torch.Tensor,
weight_scale_2: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: Optional[int],
) -> None:
original_weight_loader(
param,
fp4_weight,
weight_name=weight_name,
shard_id=shard_id,
expert_id=expert_id,
)
scale_param = (
layer.w13_weight_scale
if shard_id in ("w1", "w3")
else layer.w2_weight_scale
)
original_weight_loader(
scale_param,
weight_scale,
weight_name=self._scale_weight_name(weight_name),
shard_id=shard_id,
expert_id=expert_id,
)
scale_2_param = (
layer.w13_weight_scale_2
if shard_id in ("w1", "w3")
else layer.w2_weight_scale_2
)
original_weight_loader(
scale_2_param,
weight_scale_2,
weight_name=self._scale_2_weight_name(weight_name),
shard_id=shard_id,
expert_id=expert_id,
)
def process_loaded_weight(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: Optional[int],
) -> None:
log_quantization_start()
if shard_id == "w2":
loaded_weight = loaded_weight.to(param.device)
fp4_weight, weight_scale, weight_scale_2 = self._quantize_weight_nvfp4(
loaded_weight
)
store_quantized_weight(
param,
fp4_weight,
weight_scale,
weight_scale_2,
weight_name,
shard_id,
expert_id,
)
return
pending_key = expert_id
current = (
param,
loaded_weight,
weight_name,
shard_id,
expert_id,
)
with pending_w13_lock:
pending = pending_w13_weights.pop(pending_key, None)
if pending is None:
pending_w13_weights[pending_key] = current
return
(
pending_param,
pending_weight,
pending_name,
pending_shard_id,
pending_eid,
) = pending
if pending_shard_id == shard_id:
raise ValueError(
"--quantization nvfp4_online expects paired w1/w3 expert "
f"weights, got two {shard_id} tensors for expert {expert_id}."
)
pending_weight = pending_weight.to(param.device)
loaded_weight = loaded_weight.to(param.device)
pending_rows = pending_weight.shape[0]
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(
torch.cat([pending_weight, loaded_weight], dim=0)
)
pending_fp4_weight, loaded_fp4_weight = fp4_weight.split(
[pending_rows, loaded_rows], dim=0
)
pending_weight_scale, loaded_weight_scale = weight_scale.split(
[pending_rows, loaded_rows], dim=0
)
store_quantized_weight(
pending_param,
pending_fp4_weight.contiguous(),
pending_weight_scale.contiguous(),
weight_scale_2,
pending_name,
pending_shard_id,
pending_eid,
)
store_quantized_weight(
param,
loaded_fp4_weight.contiguous(),
loaded_weight_scale.contiguous(),
weight_scale_2,
weight_name,
shard_id,
expert_id,
)
def process_fp8_weight(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: Optional[int],
) -> None:
if not self._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:
raise ValueError(
"--quantization nvfp4_online received an FP8 expert "
"weight, but the checkpoint quantization config does not "
"declare serialized FP8 weights."
)
key = (expert_id, shard_id)
with pending_fp8_lock:
weight_scale = pending_fp8_weight_scales.pop(key, None)
if weight_scale is None:
pending_fp8_weights[key] = (
param,
loaded_weight,
weight_name,
shard_id,
expert_id,
)
return
log_quantization_start()
loaded_weight = self._dequantize_fp8_weight(
loaded_weight, weight_scale, param.device
)
process_loaded_weight(
param, loaded_weight, weight_name, shard_id, expert_id
)
def process_fp8_weight_scale(
loaded_weight: torch.Tensor,
shard_id: str,
expert_id: Optional[int],
) -> None:
key = (expert_id, shard_id)
with pending_fp8_lock:
pending = pending_fp8_weights.pop(key, None)
if pending is None:
pending_fp8_weight_scales[key] = loaded_weight
return
log_quantization_start()
(
pending_param,
pending_weight,
pending_name,
pending_shard_id,
pending_eid,
) = pending
loaded_weight = self._dequantize_fp8_weight(
pending_weight, loaded_weight, pending_param.device
)
process_loaded_weight(
pending_param,
loaded_weight,
pending_name,
pending_shard_id,
pending_eid,
)
def nvfp4_online_weight_loader(
param: torch.nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: Optional[int],
):
if shard_id not in ("w1", "w2", "w3"):
original_weight_loader(
param,
loaded_weight,
weight_name=weight_name,
shard_id=shard_id,
expert_id=expert_id,
)
return
if self._should_skip_loaded_expert(layer, param, expert_id):
return
if self._is_fp8_weight_scale_name(weight_name):
process_fp8_weight_scale(loaded_weight, shard_id, expert_id)
return
if "weight" in weight_name:
process_fp8_weight(
param, loaded_weight, weight_name, shard_id, expert_id
)
return
original_weight_loader(
param,
loaded_weight,
weight_name=weight_name,
shard_id=shard_id,
expert_id=expert_id,
)
return nvfp4_online_weight_loader
python/sglang/srt/model_loader/loader.py
+17 -2
View File
@@ -90,7 +90,7 @@ from sglang.srt.utils.common import is_cuda_alike
DEFAULT_GPU_MEMORY_FRACTION_FOR_CALIBRATION = (
0.8 # Reserve 20% GPU memory headroom for ModelOpt calibration
)
from sglang.srt.environ import envs
from sglang.srt.environ import envs, temp_set_env
from sglang.srt.model_loader.weight_utils import (
buffered_multi_thread_safetensors_weights_iterator,
download_safetensors_index_file_from_hf,
@@ -760,7 +760,22 @@ class DefaultModelLoader(BaseModelLoader):
target_device.type, gpu_id=torch.cuda.current_device()
)
model.load_weights(weights)
quant_config = getattr(model, "quant_config", None)
is_nvfp4_online = getattr(quant_config, "is_nvfp4_online", False)
if is_nvfp4_online:
# Scope exact FP4 quantization math to load-time conversion only;
# restore the original environment before serving starts.
with temp_set_env(
TRTLLM_DISABLE_FP4_QUANT_FAST_MATH="1",
FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH="1",
):
model.load_weights(weights)
if target_device.type == "cuda":
torch.cuda.synchronize()
torch.cuda.empty_cache()
else:
model.load_weights(weights)
# Used in tests to verify memory savings when using online quantization.
if is_cuda_alike():
python/sglang/srt/server_args.py
+22 -2
View File
@@ -152,6 +152,7 @@ QUANTIZATION_CHOICES = [
"modelopt",
"modelopt_fp8",
"modelopt_fp4",
"nvfp4_online",
"modelopt_mixed",
"petit_nvfp4",
"w8a8_int8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
@@ -3477,6 +3478,23 @@ class ServerArgs:
), "Please enable dp attention when setting enable_dp_lm_head. "
def _handle_moe_kernel_config(self):
if self.quantization == "nvfp4_online":
if not is_sm100_supported():
raise ValueError(
"--quantization nvfp4_online is supported only on "
"NVIDIA Blackwell SM100/SM103 GPUs."
)
if self.moe_runner_backend == "auto":
self.moe_runner_backend = "flashinfer_trtllm"
elif self.moe_runner_backend not in [
"flashinfer_trtllm",
"flashinfer_trtllm_routed",
]:
raise ValueError(
"--quantization nvfp4_online supports only "
"--moe-runner-backend flashinfer_trtllm or "
"flashinfer_trtllm_routed."
)
if self.quantization == "mxfp8":
if self.moe_runner_backend == "auto":
self.moe_runner_backend = "flashinfer_trtllm"
@@ -3539,13 +3557,14 @@ class ServerArgs:
if self.moe_runner_backend in ["flashinfer_trtllm", "experimental_sgl_trtllm"]:
assert self.quantization in [
"modelopt_fp4",
"nvfp4_online",
"fp8",
"mxfp8",
"modelopt_fp8",
"modelopt_mixed",
"compressed-tensors",
None,
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)."
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'nvfp4_online', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)."
self.disable_shared_experts_fusion = True
logger.warning(
"FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set."
@@ -3556,8 +3575,9 @@ class ServerArgs:
"fp8",
"mxfp8",
"modelopt_fp4",
"nvfp4_online",
None,
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', or bfloat16 (None)."
], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', 'nvfp4_online', or bfloat16 (None)."
self.disable_shared_experts_fusion = True
logger.warning(
"FlashInfer TRTLLM routed MoE is enabled. --disable-shared-experts-fusion is automatically set."