From 01f10acd0669b37cde1d4b248edcb3e1f986d062 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Wed, 10 Jun 2026 00:26:51 -0700 Subject: [PATCH] Implement online nvfp4 quantization (#26083) --- .../docs/advanced_features/quantization.mdx | 31 + .../advanced_features/server_arguments.mdx | 2 +- .../docs/references/environment_variables.mdx | 11 +- python/sglang/srt/configs/model_config.py | 3 + python/sglang/srt/environ.py | 3 +- .../moe/moe_runner/flashinfer_trtllm.py | 3 +- .../srt/layers/quantization/__init__.py | 2 + .../srt/layers/quantization/modelopt_quant.py | 47 +- .../srt/layers/quantization/nvfp4_online.py | 583 ++++++++++++++++++ python/sglang/srt/model_loader/loader.py | 19 +- python/sglang/srt/server_args.py | 24 +- .../test_flashinfer_trtllm_gen_moe_backend.py | 72 ++- 12 files changed, 783 insertions(+), 17 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/nvfp4_online.py diff --git a/docs_new/docs/advanced_features/quantization.mdx b/docs_new/docs/advanced_features/quantization.mdx index ed4fc6b1a..96203d352 100644 --- a/docs_new/docs/advanced_features/quantization.mdx +++ b/docs_new/docs/advanced_features/quantization.mdx @@ -147,6 +147,13 @@ The following table summarizes quantization method support across NVIDIA and AMD No NVIDIA ModelOpt; use Marlin W4A16 fallback on Ampere/Hopper and native FP4 backends on Blackwell + + 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 + petit_nvfp4 No @@ -805,6 +812,30 @@ python3 -m sglang.launch_server \ --port 30000 --host 0.0.0.0 ``` +### `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. + +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. +- **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. + +```bash Command +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-30B-A3B-Instruct-2507 \ + --tp-size 2 \ + --ep-size 2 \ + --quantization nvfp4_online \ + --port 30000 --host 0.0.0.0 +``` + ### `quark_int4fp8_moe` online quantization method SGLang running on AMD GPUs (CDNA3 or CDNA4 architecture) supports the quantization method `--quantization quark_int4fp8_moe`, that will replace [MoE layers](https://github.com/sgl-project/sglang/blob/v0.4.8/python/sglang/srt/layers/moe/fused_moe_triton/layer.py#L271) originally in high precision (bfloat16, float16 or float32) to use weights dynamically quantized to int4, that are upcasted to float8 during inference to run compute in float8 precision with activations dynamically quantized on the fly to float8. diff --git a/docs_new/docs/advanced_features/server_arguments.mdx b/docs_new/docs/advanced_features/server_arguments.mdx index a111d1fd2..42954c153 100644 --- a/docs_new/docs/advanced_features/server_arguments.mdx +++ b/docs_new/docs/advanced_features/server_arguments.mdx @@ -278,7 +278,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s `--quantization` The quantization method. `None` - awq, fp8, gptq, marlin, gptq_marlin, awq_marlin, bitsandbytes, gguf, modelopt, modelopt_fp8, modelopt_fp4, petit_nvfp4, w8a8_int8, w8a8_fp8, moe_wna16, qoq, w4afp8, mxfp4, mxfp8, auto-round, compressed-tensors, modelslim, quark_int4fp8_moe + awq, fp8, gptq, marlin, gptq_marlin, awq_marlin, bitsandbytes, gguf, modelopt, modelopt_fp8, modelopt_fp4, nvfp4_online, petit_nvfp4, w8a8_int8, w8a8_fp8, moe_wna16, qoq, w4afp8, mxfp4, mxfp8, auto-round, compressed-tensors, modelslim, quark_int4fp8_moe `--quantization-param-path` diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 4bbcbf8b0..5cd4f7266 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -576,17 +576,17 @@ SGLang supports various environment variables that can be used to configure its SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION - Enable FlashInfer TRTLLM per-token NVFP4 activation scaling; ignores checkpoint activation FP32 scale by treating it as 1 + Enable FlashInfer TRTLLM NVFP4 per-token activation scaling; ignores checkpoint activation FP32 scale by treating it as 1 false FLASHINFER_NVFP4_4OVER6 - Enable FlashInfer NVFP4 4over6 scaling for the per-token activation path; effective only with SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION + Enable FlashInfer NVFP4 4over6 scaling for NVFP4 per-token activation and online NVFP4 MoE weight quantization paths false FLASHINFER_NVFP4_4OVER6_E4M3_USE_256 - Use 256 as the E4M3 scale maximum for FlashInfer NVFP4 4over6 per-token activation scaling; otherwise uses 448 + Use 256 as the E4M3 scale maximum for FlashInfer NVFP4 4over6 scaling; otherwise uses 448 false @@ -604,6 +604,11 @@ SGLang supports various environment variables that can be used to configure its A comma-separated list of layer names to ignore during FP8 quantization. For example: model.layers.0,model.layers.1.,qkv_proj. "" + + 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. + "" + diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 75b631023..db705eae9 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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"], diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9b84337e3..b71a9e149 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py index 38c8b0c5a..d374f5cd8 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_trtllm.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 82668bded..d104c8f4f 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/python/sglang/srt/layers/quantization/__init__.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 20a699661..d362b764e 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/nvfp4_online.py b/python/sglang/srt/layers/quantization/nvfp4_online.py new file mode 100644 index 000000000..89103ab0f --- /dev/null +++ b/python/sglang/srt/layers/quantization/nvfp4_online.py @@ -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 diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 32ff5bac2..5ce23f861 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -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(): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 03f81a6b8..83e37c2c3 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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." diff --git a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py index d0238f200..ca66e6f9a 100644 --- a/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py +++ b/test/registered/backends/test_flashinfer_trtllm_gen_moe_backend.py @@ -12,7 +12,7 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=600, suite="nightly-4-gpu-b200", nightly=True) +register_cuda_ci(est_time=800, suite="nightly-4-gpu-b200", nightly=True) class FlashinferTrtllmGenMoeBackendFP8Base: @@ -243,6 +243,58 @@ class FlashinferTrtllmGenMoeBackendNVFP4Base: self.assertGreater(metrics["score"], 0.89) +class FlashinferTrtllmGenMoeBackendNvFp4OnlineBase: + backend = None + extra_env = {} + + @classmethod + def setUpClass(cls): + cls.model = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8" + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + env={**os.environ, **cls.extra_env, "SGLANG_ENABLE_JIT_DEEPGEMM": "False"}, + other_args=[ + "--attention-backend", + "triton", + "--moe-runner-backend", + cls.backend, + "--cuda-graph-max-bs", + "128", + "--tp-size", + "4", + "--ep-size", + "2", + "--quantization", + "nvfp4_online", + "--mem-fraction-static", + "0.7", + "--mamba-ssm-dtype", + "bfloat16", + ], + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.90) + + class TestFlashinferTrtllmGenMoeBackendFP8( FlashinferTrtllmGenMoeBackendFP8Base, CustomTestCase ): @@ -273,12 +325,28 @@ class TestFlashinferTrtllmGenMoeBackendBF16Routed( backend = "flashinfer_trtllm_routed" -class TestFlashinferTrtllmGenMoeBackendPerTokenNVFP4Routed( +class TestFlashinferTrtllmGenMoeBackendNvFp4PerTokenActivationRouted( FlashinferTrtllmGenMoeBackendNVFP4Base, CustomTestCase ): extra_env = {"SGLANG_FLASHINFER_NVFP4_PER_TOKEN_ACTIVATION": "1"} backend = "flashinfer_trtllm_routed" +class TestFlashinferTrtllmGenMoeBackendNvFp4Online( + FlashinferTrtllmGenMoeBackendNvFp4OnlineBase, CustomTestCase +): + extra_env = { + "FLASHINFER_NVFP4_4OVER6": "1", + "FLASHINFER_NVFP4_4OVER6_ERR_MODE": "MSE", + "FLASHINFER_NVFP4_4OVER6_ERR_USE_FAST_MATH": "1", + "FLASHINFER_NVFP4_4OVER6_E4M3_USE_256": "1", + "SGLANG_FP4_IGNORED_LAYERS": ",".join( + ["shared_expert"] + + [f"model.layers.{layer_id}" for layer_id in range(40, 48)] + ), + } + backend = "flashinfer_trtllm" + + if __name__ == "__main__": unittest.main()