chore: import upstream snapshot with attribution

python/sglang/srt/layers/quantization/nvfp4_online.py

@@ -0,0 +1,584 @@
+# 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.dequant_fp4_to_fp8 = 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