diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index e963ac69a..321df36a9 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -291,6 +291,13 @@ the fused ConvRot kernel without requantizing them. TP1/2/4 and sequence or layerwise offload are supported; TP8 violates the checkpoint's 256-element ConvRot group boundary, and FSDP is rejected. +Self-describing community MXFP8 files use the same +`--transformer-weights-path ` form and need no +quantization flag. SGLang reads their global format metadata, keeps unmarked H3 +layers in their original dtype, and reuses SRT's MXFP8 dense kernels. The +selected SRT backend must support MXFP8 on the target GPU; FSDP is rejected for +this mixed per-layer layout. + ### Advanced: precomputed AdaLN cache The [model card](https://huggingface.co/MiniMaxAI/MiniMax-H3) notes that about diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index a6ea52029..642987a1d 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -176,6 +176,14 @@ backend. comfy-kitchen CUDA; auto-detected; uses the fused Kitchen INT8 kernel and validates weight/scale layout before model construction. TP requires every row-parallel input shard to preserve the checkpoint's ConvRot group boundary. The H3 256-group DiT supports TP1/2/4, not TP8; its Qwen3-VL encoder keeps TP8 by replicating only incompatible row projections. Offload is supported; FSDP is not. + + mxfp8 + Self-describing serialized weights, or BF16/FP16 weights for online quantization + --transformer-weights-path, or --quantization mxfp8 for online quantization + Native diffusion linears; mixed MiniMax-H3 checkpoints are selected per layer + SRT's platform MXFP8 backend + Serialized metadata is auto-detected. NVIDIA and ROCm reuse SRT's dense MXFP8 kernels; Ascend keeps its native online path. Mixed per-layer checkpoints do not support FSDP. + qvg-kv Unquantized model with runtime causal KV-cache compression @@ -1124,11 +1132,13 @@ MindStudio-ModelSlim (msModelSlim) is a model offline quantization compression t - [x] ```W4A4_MXFP4``` / ```W4A4_MXFP4_DUALSCALE``` linear with offline quantization (msmodelslim pre-quantized weights) - [x] ```mxfp4_npu``` linear with online quantization (`--quantization mxfp4_npu`) -## MXFP8 Online Quantization +## MXFP8 Quantization -For online MXFP8 quantization, load the original FP16/BF16 model and add `--quantization mxfp8`. -Weights are quantized at load time via `npu_dynamic_mx_quant`, and activations are quantized per-token -during inference with `npu_quant_matmul` (block_size=32). +Self-describing serialized MXFP8 transformers are inferred from their +checkpoint metadata. For online MXFP8 quantization, load the original +FP16/BF16 model and add `--quantization mxfp8`. NVIDIA and ROCm reuse SRT's +dense MXFP8 implementation; Ascend uses `npu_dynamic_mx_quant` and +`npu_quant_matmul` with a block size of 32. ```bash sglang generate \ @@ -1138,7 +1148,9 @@ sglang generate \ --save-output ``` -> **Hardware requirement:** Ascend A5 series or newer. `npu_dynamic_mx_quant` is not available on A2/A3. +The selected SRT backend must provide an MXFP8 kernel on NVIDIA or ROCm. On +Ascend, the hardware requirement remains A5 series or newer; +`npu_dynamic_mx_quant` is not available on A2/A3. ## MXFP8 Offline Quantization (msmodelslim) diff --git a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py index 376fd87b2..3bdbc3e68 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/configs/models/dits/minimax_h3.py @@ -13,6 +13,7 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig): # H3 fuses Q/K/V, so split projections are stacked for the fused LoRA layer param_names_mapping: dict = field( default_factory=lambda: { + r"^(.*)\.weight_scale$": r"\1.weight_scale_inv", r"^(.*\.lora_[AB])\.[^.]+$": r"\1", r"^base_model\.model\.(.*\.lora_[AB])$": r"\1", r"^transformer\.(.*\.lora_[AB])$": r"\1", diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index 25b1dc1da..da90353dd 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -24,7 +24,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.mxfp4 import Mxfp4Config from sglang.multimodal_gen.runtime.layers.quantization.mxfp4_npu import ( NPUMXFP4Config, ) -from sglang.multimodal_gen.runtime.layers.quantization.mxfp8_npu import MXFP8Config +from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config QuantizationMethods = Literal[ "fp8", diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py b/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py index 36482f0bc..bcc53f635 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py @@ -6,7 +6,6 @@ from typing import Any, Optional import torch import torch.nn as nn -from packaging import version from sglang.multimodal_gen.runtime.layers.linear import ( LinearBase, @@ -18,39 +17,17 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config impor QuantizeMethodBase, ) from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs +from sglang.srt.layers.quantization.bitsandbytes import ( + BitsAndBytesConfig as SRTBitsAndBytesConfig, +) +from sglang.srt.layers.quantization.bitsandbytes import ( + calculate_quant_ratio, + is_layer_skipped_bnb, + require_bitsandbytes, +) -def _require_bitsandbytes() -> None: - try: - import bitsandbytes - - if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"): - raise ImportError( - "bitsandbytes version is wrong. Please install bitsandbytes>=0.46.1." - ) - except ImportError as err: - raise ImportError( - "Please install bitsandbytes>=0.46.1 via " - "`pip install bitsandbytes>=0.46.1` to use bitsandbytes quantizer." - ) from err - - -def _calculate_quant_ratio(dtype: torch.dtype) -> int: - if dtype.is_floating_point: - return torch.finfo(dtype).bits // torch.iinfo(torch.uint8).bits - return torch.iinfo(dtype).bits // torch.iinfo(torch.uint8).bits - - -def _is_layer_skipped(prefix: str, skipped_modules: list[str]) -> bool: - components = prefix.split(".") - if any(module_name in components for module_name in skipped_modules): - return True - - prefixes = {".".join(components[: i + 1]) for i in range(len(components))} - return bool(set(skipped_modules) & prefixes) - - -class BitsAndBytesConfig(QuantizationConfig): +class BitsAndBytesConfig(SRTBitsAndBytesConfig, QuantizationConfig): """Config class for pre-quantized bitsandbytes 4-bit checkpoints.""" def __init__( @@ -66,79 +43,26 @@ class BitsAndBytesConfig(QuantizationConfig): llm_int8_skip_modules: list[str] | None = None, llm_int8_threshold: float = 6.0, ) -> None: - super().__init__() - self.load_in_8bit = load_in_8bit - self.load_in_4bit = load_in_4bit - self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype - self.bnb_4bit_quant_storage = bnb_4bit_quant_storage - self.bnb_4bit_quant_type = bnb_4bit_quant_type - self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant - self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload - self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight - self.llm_int8_skip_modules = llm_int8_skip_modules or [] - self.llm_int8_threshold = llm_int8_threshold - + super().__init__( + load_in_8bit=load_in_8bit, + load_in_4bit=load_in_4bit, + bnb_4bit_compute_dtype=bnb_4bit_compute_dtype, + bnb_4bit_quant_storage=bnb_4bit_quant_storage, + bnb_4bit_quant_type=bnb_4bit_quant_type, + bnb_4bit_use_double_quant=bnb_4bit_use_double_quant, + llm_int8_enable_fp32_cpu_offload=llm_int8_enable_fp32_cpu_offload, + llm_int8_has_fp16_weight=llm_int8_has_fp16_weight, + llm_int8_skip_modules=llm_int8_skip_modules, + llm_int8_threshold=llm_int8_threshold, + ) if self.load_in_8bit or not self.load_in_4bit: raise ValueError("SGLang diffusion only supports bitsandbytes 4-bit.") - if self.bnb_4bit_quant_storage != "uint8": - raise ValueError( - f"Unsupported bnb_4bit_quant_storage: {self.bnb_4bit_quant_storage}" - ) - - @classmethod - def get_name(cls) -> str: - return "bitsandbytes" - - def get_scaled_act_names(self) -> list[str]: - return [] - - @classmethod - def get_supported_act_dtypes(cls) -> list[torch.dtype]: - return [torch.float32, torch.float16, torch.bfloat16] - - @classmethod - def get_min_capability(cls) -> int: - return 70 - - @staticmethod - def get_config_filenames() -> list[str]: - return [] - - @classmethod - def from_config(cls, config: dict[str, Any]) -> BitsAndBytesConfig: - def get_safe_value(keys, default_value=None): - try: - value = QuantizationConfig.get_from_keys(config, keys) - return value if value is not None else default_value - except ValueError: - return default_value - - return cls( - load_in_8bit=get_safe_value(["load_in_8bit"], False), - load_in_4bit=get_safe_value(["load_in_4bit"], True), - bnb_4bit_compute_dtype=get_safe_value( - ["bnb_4bit_compute_dtype"], "float32" - ), - bnb_4bit_quant_storage=get_safe_value(["bnb_4bit_quant_storage"], "uint8"), - bnb_4bit_quant_type=get_safe_value(["bnb_4bit_quant_type"], "fp4"), - bnb_4bit_use_double_quant=get_safe_value( - ["bnb_4bit_use_double_quant"], False - ), - llm_int8_enable_fp32_cpu_offload=get_safe_value( - ["llm_int8_enable_fp32_cpu_offload"], False - ), - llm_int8_has_fp16_weight=get_safe_value( - ["llm_int8_has_fp16_weight"], False - ), - llm_int8_skip_modules=get_safe_value(["llm_int8_skip_modules"], []), - llm_int8_threshold=get_safe_value(["llm_int8_threshold"], 6.0), - ) def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> Optional[QuantizeMethodBase]: if isinstance(layer, LinearBase): - if _is_layer_skipped(prefix, self.llm_int8_skip_modules): + if is_layer_skipped_bnb(prefix, self.llm_int8_skip_modules): return UnquantizedLinearMethod() return BitsAndBytesLinearMethod(self) return None @@ -148,7 +72,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase): """Linear method for pre-quantized bitsandbytes 4-bit weights.""" def __init__(self, quant_config: BitsAndBytesConfig): - _require_bitsandbytes() + require_bitsandbytes() self.quant_config = quant_config def create_weights( @@ -161,7 +85,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase): params_dtype: torch.dtype, **extra_weight_attrs, ) -> None: - quant_ratio = _calculate_quant_ratio(params_dtype) + quant_ratio = calculate_quant_ratio(params_dtype) output_size_per_partition = sum(output_partition_sizes) total_size = input_size_per_partition * output_size_per_partition if total_size % quant_ratio != 0: @@ -255,11 +179,11 @@ class BitsAndBytes4BitLinear(nn.Module): compute_dtype: torch.dtype | None = None, ) -> None: super().__init__() - _require_bitsandbytes() + require_bitsandbytes() self.in_features = in_features self.out_features = out_features self.compute_dtype = compute_dtype - quant_ratio = _calculate_quant_ratio(compute_dtype or torch.get_default_dtype()) + quant_ratio = calculate_quant_ratio(compute_dtype or torch.get_default_dtype()) total_size = in_features * out_features if total_size % quant_ratio != 0: raise ValueError( diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py index e758a12f4..ffe08c181 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py @@ -4,51 +4,21 @@ # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/quantization/base_config.py import inspect -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any import torch -from torch import nn -if TYPE_CHECKING: - from sglang.multimodal_gen.runtime.layers.quantization import QuantizationMethods -else: - QuantizationMethods = str +from sglang.srt.layers.quantization.base_config import ( + QuantizationConfig as SRTQuantizationConfig, +) +from sglang.srt.layers.quantization.base_config import ( + QuantizeMethodBase as SRTQuantizeMethodBase, +) -class QuantizeMethodBase(ABC): - """Base class for different quantized methods.""" - - @abstractmethod - def create_weights( - self, layer: torch.nn.Module, *weight_args, **extra_weight_attrs - ): - """Create weights for a layer. - - The weights will be set as attributes of the layer.""" - raise NotImplementedError - - @abstractmethod - def apply(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: - """Apply the weights in layer to the input tensor. - - Expects create_weights to have been called before on the layer.""" - raise NotImplementedError - - # Not required functions +class QuantizeMethodBase(SRTQuantizeMethodBase): def embedding(self, layer: torch.nn.Module, *args, **kwargs) -> torch.Tensor: - """Gather embeddings in the layer based on indices in the input tensor. - - Expects create_weights to have been called before on the layer.""" raise NotImplementedError - def process_weights_after_loading(self, layer: nn.Module) -> None: - """Process the weight after loading. - - This can be used for example, to transpose weights for computation. - """ - return - def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: """ @@ -62,98 +32,13 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> return class_embedding is not None and class_embedding is not base_embedding -class QuantizationConfig(ABC): - """Base class for quantization configs.""" - +class QuantizationConfig(SRTQuantizationConfig): # for quantization frameworks with a separate quantized model provided, e.g. Nunchaku quantized_model_path: str | None = None checkpoint_uses_native_qkv_layout: bool = False - def __init__(self): - super().__init__() - # mapping is updated by models as they initialize - self.packed_modules_mapping: dict[str, list[str]] = dict() - - @abstractmethod - def get_name(self) -> QuantizationMethods: - """Name of the quantization method.""" - raise NotImplementedError - - @abstractmethod - def get_supported_act_dtypes(self) -> list[torch.dtype]: - """List of supported activation dtypes.""" - raise NotImplementedError - - @classmethod - @abstractmethod - def get_min_capability(cls) -> int: - """Minimum GPU capability to support the quantization method. - - E.g., 70 for Volta, 75 for Turing, 80 for Ampere. - This requirement is due to the custom CUDA kernels used by the - quantization method. - """ - raise NotImplementedError - - @staticmethod - @abstractmethod - def get_config_filenames() -> list[str]: - """List of filenames to search for in the model directory.""" - raise NotImplementedError - - @classmethod - @abstractmethod - def from_config(cls, config: dict[str, Any]) -> "QuantizationConfig": - """Create a config class from the model's quantization config.""" - raise NotImplementedError - - @classmethod - def override_quantization_method( - cls, hf_quant_cfg, user_quant - ) -> QuantizationMethods | None: - """ - Detects if this quantization method can support a given checkpoint - format by overriding the user specified quantization method -- - this method should only be overwritten by subclasses in exceptional - circumstances - """ - return None - - @staticmethod - def get_from_keys(config: dict[str, Any], keys: list[str]) -> Any: - """Get a value from the model's quantization config.""" - for key in keys: - if key in config: - return config[key] - raise ValueError( - f"Cannot find any of {keys} in the model's " "quantization config." - ) - - @staticmethod - def get_from_keys_or(config: dict[str, Any], keys: list[str], default: Any) -> Any: - """Get a optional value from the model's quantization config.""" - try: - return QuantizationConfig.get_from_keys(config, keys) - except ValueError: - return default - - @abstractmethod - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> QuantizeMethodBase | None: - """Get the quantize method to use for the quantized layer. - - Args: - layer: The layer for the quant method. - prefix: The full name of the layer in the state dict - Returns: - The quantize method. None if the given layer doesn't support quant - method. - """ - raise NotImplementedError - - def get_cache_scale(self, name: str) -> str | None: - return None + def get_scaled_act_names(self) -> list[str]: + return [] def supports_input_partition( self, prefix: str, input_size_per_partition: int diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py index e6a95a628..ca39b2f0e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union +from typing import TYPE_CHECKING, List, Optional, Union import torch from torch.nn import Module @@ -15,6 +15,7 @@ from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_tensor_model_parallel_world_size, ) from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, LinearMethodBase, UnquantizedLinearMethod, ) @@ -35,6 +36,7 @@ from sglang.multimodal_gen.runtime.utils.common import ( use_intel_amx_backend, ) from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config from sglang.srt.layers.quantization.fp8_utils import ( apply_fp8_linear, can_auto_enable_marlin_fp8, @@ -69,89 +71,23 @@ if USE_AITER or _use_hip_int4: pass -ACTIVATION_SCHEMES = ["static", "dynamic"] - logger = logging.getLogger(__name__) -class Fp8Config(QuantizationConfig): +class Fp8Config(SRTFp8Config, QuantizationConfig): """Config class for FP8. No-arg ``Fp8Config()`` selects online (post-load) weight quantization: ``is_checkpoint_fp8_serialized=False`` with ``activation_scheme="dynamic"``. """ - def __init__( - self, - is_checkpoint_fp8_serialized: bool = False, - activation_scheme: str = "dynamic", - ignored_layers: Optional[List[str]] = None, - weight_block_size: List[int] = None, - packed_modules_mapping: Optional[Dict[str, List[str]]] = None, - ) -> None: - self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized - if is_checkpoint_fp8_serialized: - logger.info("Detected fp8 checkpoint.") - if activation_scheme not in ACTIVATION_SCHEMES: - raise ValueError(f"Unsupported activation scheme {activation_scheme}") - self.activation_scheme = activation_scheme - self.ignored_layers = ignored_layers or [] - self.packed_modules_mapping = packed_modules_mapping or {} - if weight_block_size is not None: - if not is_checkpoint_fp8_serialized: - raise ValueError( - "The block-wise quantization only supports fp8-serialized checkpoint for now." - ) - if len(weight_block_size) != 2: - raise ValueError( - f"The quantization block size of weight must have 2 dimensions, but got {len(weight_block_size)} dimensions." - ) - if activation_scheme != "dynamic": - raise ValueError( - f"The block-wise quantization only supports dynamic activation scheme for now, but got {activation_scheme} activation scheme." - ) - self.weight_block_size = weight_block_size - - @classmethod - def get_name(cls) -> str: - return "fp8" - - @classmethod - def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return [torch.bfloat16, torch.half] - @classmethod def get_min_capability(cls) -> int: return 80 - @classmethod - def get_config_filenames(cls) -> List[str]: - return [] - - @classmethod - def from_config(cls, config: Dict[str, Any]) -> Fp8Config: - quant_method = cls.get_from_keys(config, ["quant_method"]) - is_checkpoint_fp8_serialized = "fp8" in quant_method - activation_scheme = cls.get_from_keys(config, ["activation_scheme"]) - ignored_layers = cls.get_from_keys_or( - config, ["ignored_layers", "modules_to_not_convert"], None - ) - if ignored_layers: - # hacking ministral - ignored_layers = [layer.replace("model.", "") for layer in ignored_layers] - weight_block_size = cls.get_from_keys_or(config, ["weight_block_size"], None) - return cls( - is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, - activation_scheme=activation_scheme, - ignored_layers=ignored_layers, - weight_block_size=weight_block_size, - ) - def get_quant_method( self, layer: torch.nn.Module, prefix: str ) -> Optional[QuantizeMethodBase]: - from sglang.multimodal_gen.runtime.layers.linear import LinearBase - if isinstance(layer, LinearBase): if is_layer_skipped( prefix, @@ -162,9 +98,6 @@ class Fp8Config(QuantizationConfig): return Fp8LinearMethod(self) return None - def get_scaled_act_names(self) -> List[str]: - return [] - class Fp8LinearMethod(LinearMethodBase): """Linear method for FP8. diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py index 1bc49779d..79202676e 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py @@ -53,7 +53,8 @@ class ModelSlimMXFP8Scheme(ModelSlimLinearScheme): # msmodelslim exports weight_scale as uint8, shape [out, in/32]. # NOTE: This parameter is intentionally named "weight_scale" (not - # "weight_scale_inv" as used in mxfp8_npu.py) because the weight loader + # "weight_scale_inv" as used by SRT's NPUMXFP8LinearMethod) because the + # weight loader # matches parameter names to checkpoint keys, and msmodelslim checkpoints # store this tensor under the key ".weight_scale". scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py new file mode 100644 index 000000000..3a2396650 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MXFP8 diffusion adapter backed by SRT's dense linear kernels.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.layers.linear import ( + LinearBase, + UnquantizedLinearMethod, +) +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.multimodal_gen.runtime.platforms import current_platform +from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + NPUMXFP8LinearMethod, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config +from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod as SRTFp8LinearMethod + + +class MXFP8Config(SRTFp8Config, QuantizationConfig): + """Route diffusion linears through SRT's MXFP8 implementation.""" + + def __init__( + self, + *, + is_checkpoint_fp8_serialized: bool = False, + layer_markers: dict[str, dict[str, Any]] | None = None, + ignored_layers: list[str] | None = None, + ) -> None: + if current_platform.is_mps(): + raise ValueError("MXFP8 is not supported on MPS") + super().__init__( + is_checkpoint_fp8_serialized=is_checkpoint_fp8_serialized, + activation_scheme="dynamic", + ignored_layers=ignored_layers, + weight_block_size=[1, 32] if is_checkpoint_fp8_serialized else None, + use_mxfp8=True, + ) + self.layer_markers = layer_markers + self.checkpoint_uses_native_qkv_layout = layer_markers is not None + + @classmethod + def get_min_capability(cls) -> int: + return 90 + + @classmethod + def from_config(cls, config: dict[str, Any]) -> MXFP8Config: + quant_method = str(cls.get_from_keys(config, ["quant_method"])).lower() + if "mxfp8" not in quant_method: + raise ValueError(f"Expected an MXFP8 checkpoint, got {quant_method!r}") + activation_scheme = cls.get_from_keys_or( + config, ["activation_scheme"], "dynamic" + ) + if activation_scheme != "dynamic": + raise ValueError("MXFP8 only supports dynamic activation scaling") + ignored_layers = cls.get_from_keys_or( + config, ["ignored_layers", "modules_to_not_convert"], None + ) + return cls( + is_checkpoint_fp8_serialized=True, + ignored_layers=ignored_layers, + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> QuantizeMethodBase | None: + if not isinstance(layer, LinearBase): + return None + if self.layer_markers is not None and prefix not in self.layer_markers: + return UnquantizedLinearMethod() + if current_platform.is_npu(): + return NPUMXFP8LinearMethod(self) + return SRTFp8LinearMethod(self) + + +__all__ = ["MXFP8Config"] diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py b/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py deleted file mode 100644 index 2ba40d1af..000000000 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py +++ /dev/null @@ -1,176 +0,0 @@ -"""Online MXFP8 quantization for Diffusion models on Ascend NPU. - -Provides ``MXFP8Config`` (registered as ``"mxfp8"``) and -``NPUMXFP8DiffusionLinearMethod`` which quantise FP16/BF16 weights to MXFP8 -at load time and use ``npu_dynamic_mx_quant`` + ``npu_quant_matmul`` for -inference, mirroring the LLM-side ``NPUMXFP8LinearMethod``. -""" - -from __future__ import annotations - -from typing import Any, Dict, List, Optional - -import torch -from torch.nn.parameter import Parameter - -from sglang.multimodal_gen.runtime.platforms import current_platform - -_is_npu = current_platform.is_npu() - -if _is_npu: - import torch_npu - -from sglang.multimodal_gen.runtime.layers.linear import LinearBase, LinearMethodBase -from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( - QuantizationConfig, - QuantizeMethodBase, -) -from sglang.multimodal_gen.runtime.models.parameter import ModelWeightParameter -from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger - -logger = init_logger(__name__) - -MXFP8_BLOCK_SIZE = 32 - - -class MXFP8Config(QuantizationConfig): - """Config for online MXFP8 quantization on Ascend NPU (Diffusion).""" - - def __init__(self) -> None: - super().__init__() - - @classmethod - def get_name(cls) -> str: - return "mxfp8" - - @classmethod - def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return [torch.bfloat16, torch.float16] - - @classmethod - def get_min_capability(cls) -> int: - return 0 # NPU, not CUDA - - @classmethod - def get_config_filenames(cls) -> List[str]: - return [] - - @classmethod - def from_config(cls, config: Dict[str, Any]) -> MXFP8Config: - return cls() - - def get_quant_method( - self, layer: torch.nn.Module, prefix: str - ) -> Optional[QuantizeMethodBase]: - if isinstance(layer, LinearBase): - return NPUMXFP8DiffusionLinearMethod(self) - return None - - def get_scaled_act_names(self) -> List[str]: - return [] - - -class NPUMXFP8DiffusionLinearMethod(LinearMethodBase): - """Ascend NPU MXFP8 linear method for Diffusion models. - - Online mode: loads FP16/BF16 weights → quantises to MXFP8 at load time. - Inference: dynamic MXFP8 activation quant + MXFP8 matmul (block_size=32). - """ - - def __init__(self, quant_config: MXFP8Config): - self.quant_config = quant_config - - def create_weights( - self, - layer: torch.nn.Module, - input_size_per_partition: int, - output_partition_sizes: List[int], - input_size: int, - output_size: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - output_size_per_partition = sum(output_partition_sizes) - weight_loader = extra_weight_attrs.get("weight_loader") - - layer.logical_widths = output_partition_sizes - layer.input_size_per_partition = input_size_per_partition - layer.output_size_per_partition = output_size_per_partition - layer.orig_dtype = params_dtype - - # Load weights in original dtype; quantise later in process_weights_after_loading - weight = ModelWeightParameter( - data=torch.empty( - output_size_per_partition, - input_size_per_partition, - dtype=params_dtype, - ), - input_dim=1, - output_dim=0, - weight_loader=weight_loader, - ) - layer.register_parameter("weight", weight) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - - weight_fp = layer.weight.data - if weight_fp.dtype not in (torch.float16, torch.bfloat16): - weight_fp = weight_fp.to(torch.bfloat16) - - # Move weight to NPU if needed. We intentionally use a conditional - # move rather than an assert because `dit_cpu_offload` defaults to - # True in ServerArgs, which causes fsdp_load to move every parameter - # back to CPU after loading (even when the target device is NPU). - # npu_dynamic_mx_quant requires an NPU tensor, so we must transfer - # here. The quantized fp8 weights produced below will remain on NPU - # for inference; if the model still needs to be offloaded after - # quantization (e.g. very large model on a small NPU), a higher-level - # offload pass can move them back afterwards. - if not weight_fp.is_npu: - weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}") - - # Online MXFP8 quantisation of weights (block_size=32) - qw, w_scale = torch_npu.npu_dynamic_mx_quant( - weight_fp, dst_type=torch_npu.float8_e4m3fn - ) - layer.weight = Parameter(qw, requires_grad=False) - layer.weight_scale_inv = Parameter(w_scale, requires_grad=False) - - def apply( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - original_dtype = x.dtype - if original_dtype not in (torch.float16, torch.bfloat16): - x = x.to(torch.bfloat16) - original_dtype = torch.bfloat16 - - # Flatten to 2D [tokens, hidden] so npu_dynamic_mx_quant returns 3D scale - input_shape = x.shape - x_2d = x.reshape(-1, x.shape[-1]) - - # Dynamic MXFP8 activation quantisation - qx, input_scale = torch_npu.npu_dynamic_mx_quant( - x_2d, dst_type=torch_npu.float8_e4m3fn - ) - - # MXFP8 matmul - output = torch_npu.npu_quant_matmul( - qx, - layer.weight.transpose(0, 1), - layer.weight_scale_inv.transpose(0, 1), - scale_dtype=torch_npu.float8_e8m0fnu, - pertoken_scale=input_scale, - pertoken_scale_dtype=torch_npu.float8_e8m0fnu, - bias=bias.to(torch.float32) if bias is not None else None, - output_dtype=original_dtype, - group_sizes=[1, 1, MXFP8_BLOCK_SIZE], - ) - - # Restore original shape (replace last dim with output features) - output_shape = list(input_shape[:-1]) + [output.shape[-1]] - output = output.reshape(output_shape) - - return output diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index ca53f18cc..e1161acea 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -63,8 +63,7 @@ _DTYPE_MISMATCH_EXAMPLE_LIMIT = 3 def _is_bitsandbytes_quant_config(quant_config: Any | None) -> bool: if quant_config is None: return False - quant_name_getter = getattr(type(quant_config), "get_name", None) - return bool(callable(quant_name_getter) and quant_name_getter() == "bitsandbytes") + return quant_config.get_name() == "bitsandbytes" def _format_dtype_mismatch_summary( diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 477c95146..6e8d65482 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -75,8 +75,7 @@ _HF_SAFETENSORS_URL_RE = re.compile( def _get_quant_config_name(config: Optional[QuantizationConfig]) -> Optional[str]: if config is None: return None - quant_name_getter = getattr(type(config), "get_name", None) - return quant_name_getter() if callable(quant_name_getter) else None + return config.get_name() def _merge_modelopt_fp4_configs( @@ -176,7 +175,14 @@ class TransformerQuantLoadSpec: @property def uses_comfy_layer_markers(self) -> bool: - return self.is_comfy_fp8 or self.is_serialized_kitchen_int8 + return ( + self.is_comfy_fp8 + or self.is_serialized_kitchen_int8 + or ( + _get_quant_config_name(self.quant_config) == "mxfp8" + and self.quant_config.layer_markers is not None + ) + ) class _TransformerQuantAdapter: @@ -285,9 +291,7 @@ class _Flux2Nvfp4FallbackAdapter(_TransformerQuantAdapter): if cls_name != "Flux2Transformer2DModel" or quant_config is None: return - quant_name_getter = getattr(type(quant_config), "get_name", None) - quant_name = quant_name_getter() if callable(quant_name_getter) else None - if quant_name != "modelopt_fp4": + if _get_quant_config_name(quant_config) != "modelopt_fp4": return weights_path = os.path.basename(server_args.transformer_weights_path or "") @@ -370,10 +374,7 @@ class _ModelOptFp8OffloadAdapter(_TransformerQuantAdapter): if quant_config is None: return - quant_name_getter = getattr(type(quant_config), "get_name", None) - quant_name = quant_name_getter() if callable(quant_name_getter) else None - - if quant_name != "modelopt_fp8": + if _get_quant_config_name(quant_config) != "modelopt_fp8": return component_offload = _uses_component_offload( @@ -856,7 +857,7 @@ def _needs_device_weight_postprocess( ) -> bool: """Return whether post-load weight processing needs CUDA/NPU tensors.""" quant_name = _get_quant_config_name(quant_config) - if quant_name in ("modelopt_fp8", "comfy_fp8"): + if quant_name in ("modelopt_fp8", "comfy_fp8", "mxfp8"): return True if quant_name == "kitchen_int8": assert isinstance(quant_config, KitchenInt8Config) @@ -864,7 +865,6 @@ def _needs_device_weight_postprocess( serialized_flag_by_quant_name = { "fp8": "is_checkpoint_fp8_serialized", - "mxfp8": "is_checkpoint_fp8_serialized", "mxfp4": "is_checkpoint_mxfp4_serialized", "mxfp4_npu": "is_checkpoint_mxfp4_npu_serialized", } diff --git a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py index 80d415074..f1ba647bc 100644 --- a/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py +++ b/python/sglang/multimodal_gen/runtime/utils/quantization_utils.py @@ -16,6 +16,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ComfyFp8 from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( KitchenInt8Config, ) +from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.layers.modelopt_utils import canonicalize_modelopt_quant_algo from sglang.srt.model_loader.checkpoint_quantization import ( @@ -33,9 +34,13 @@ def inspect_comfy_quant_markers( checkpoint_meta: dict[str, tuple[str, tuple[int, ...]]] = {} raw_markers: dict[str, dict[str, Any]] = {} marked_dtype_weight_prefixes: set[str] = set() + global_quant_formats: set[str] = set() for path in safetensors_list: with safe_open(path, framework="pt", device="cpu") as checkpoint: + metadata = checkpoint.metadata() or {} + if quant_format := metadata.get("quant_format"): + global_quant_formats.add(quant_format.lower()) for key in checkpoint.keys(): tensor_slice = checkpoint.get_slice(key) checkpoint_meta[key] = ( @@ -67,6 +72,23 @@ def inspect_comfy_quant_markers( ) raw_markers[prefix] = marker + if global_quant_formats == {"mxfp8"}: + for prefix in marked_dtype_weight_prefixes: + weight_meta = checkpoint_meta[f"{prefix}.weight"] + scale_meta = checkpoint_meta.get(f"{prefix}.weight_scale") + weight_dtype, weight_shape = weight_meta + if ( + weight_dtype != "F8_E4M3" + or len(weight_shape) != 2 + or weight_shape[1] % 32 != 0 + or scale_meta != ("U8", (weight_shape[0], weight_shape[1] // 32)) + ): + raise ValueError( + f"MXFP8 layer {prefix!r} has incompatible weight/scale metadata: " + f"{weight_meta} and {scale_meta}" + ) + raw_markers.setdefault(prefix, {"format": "mxfp8"}) + missing_markers = marked_dtype_weight_prefixes - raw_markers.keys() if missing_markers: raise ValueError( @@ -124,6 +146,11 @@ def resolve_comfy_checkpoint_quantization( return KitchenInt8Config(layer_markers=layer_markers) if formats == ["float8_e4m3fn"]: return ComfyFp8Config(layer_markers) + if formats == ["mxfp8"]: + return MXFP8Config( + is_checkpoint_fp8_serialized=True, + layer_markers=layer_markers, + ) raise NotImplementedError( "Unsupported Comfy quantization format(s): " + ", ".join(formats) ) diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py index 07ddc6e6a..11ff7fa0c 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py @@ -10,8 +10,12 @@ from safetensors.torch import safe_open, save_file from torch import nn from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear +from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import ( + BitsAndBytesConfig, +) from sglang.multimodal_gen.runtime.loader import fsdp_load, rank_local_checkpoint from sglang.multimodal_gen.runtime.loader.weight_load_plan import WeightLoadPlan +from sglang.srt.layers.quantization.fp8 import Fp8Config class _UniformDtypeModel(nn.Module): @@ -45,6 +49,10 @@ class _CustomEntrypointModel(_UniformDtypeModel): class TestFSDPMixedPrecisionPolicy(unittest.TestCase): + def test_quant_config_detection_uses_the_runtime_instance(self): + self.assertTrue(fsdp_load._is_bitsandbytes_quant_config(BitsAndBytesConfig())) + self.assertFalse(fsdp_load._is_bitsandbytes_quant_config(Fp8Config())) + def _load_and_capture_policy( self, model_cls: type[nn.Module], diff --git a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py index b47aba6a8..48810e0e3 100644 --- a/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py +++ b/python/sglang/multimodal_gen/test/unit/test_transformer_quant.py @@ -69,6 +69,7 @@ from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import ( ModelOptFp8Config, _prepare_nvfp4_weight_bytes, ) +from sglang.multimodal_gen.runtime.layers.quantization.mxfp8 import MXFP8Config from sglang.multimodal_gen.runtime.loader.component_loaders import transformer_loader from sglang.multimodal_gen.runtime.loader.component_loaders.transformer_loader import ( TransformerLoader, @@ -101,6 +102,14 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import ( from sglang.multimodal_gen.tools.build_modelopt_nvfp4_transformer import ( _updated_quant_config, ) +from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( + NPUMXFP8LinearMethod, +) +from sglang.srt.layers.quantization.bitsandbytes import ( + BitsAndBytesConfig as SRTBitsAndBytesConfig, +) +from sglang.srt.layers.quantization.fp8 import Fp8Config as SRTFp8Config +from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod as SRTFp8LinearMethod class _FakeFluxTransformer: @@ -407,6 +416,57 @@ class TestTransformerQuantHelpers(unittest.TestCase): UnquantizedLinearMethod, ) + def test_minimax_h3_global_mxfp8_metadata_selects_srt_per_layer(self): + with tempfile.NamedTemporaryFile(suffix=".safetensors") as checkpoint: + save_file( + { + "blocks.0.attn.out_proj.weight": torch.ones( + (32, 64), dtype=torch.float8_e4m3fn + ), + "blocks.0.attn.out_proj.weight_scale": torch.ones( + (32, 2), dtype=torch.uint8 + ), + "token_refiner.blocks.0.attn.out_proj.weight": torch.ones( + (32, 64), dtype=torch.bfloat16 + ), + }, + checkpoint.name, + metadata={"quant_format": "mxfp8"}, + ) + _, layer_markers = inspect_minimax_h3_safetensors([checkpoint.name]) + config = resolve_minimax_h3_checkpoint_quantization(layer_markers) + + self.assertIsInstance(config, MXFP8Config) + self.assertIsInstance(config, SRTFp8Config) + layer = LinearBase(input_size=64, output_size=32) + self.assertIsInstance( + config.get_quant_method(layer, "blocks.0.attn.out_proj"), + SRTFp8LinearMethod, + ) + self.assertIsInstance( + config.get_quant_method(layer, "token_refiner.blocks.0.attn.out_proj"), + UnquantizedLinearMethod, + ) + + def test_mxfp8_npu_selects_srt_linear_method(self): + with ( + patch( + "sglang.multimodal_gen.runtime.layers.quantization.mxfp8.current_platform.is_mps", + return_value=False, + ), + patch( + "sglang.multimodal_gen.runtime.layers.quantization.mxfp8.current_platform.is_npu", + return_value=True, + ), + ): + config = MXFP8Config() + method = config.get_quant_method( + LinearBase(input_size=64, output_size=32), + "blocks.0.attn.out_proj", + ) + + self.assertIsInstance(method, NPUMXFP8LinearMethod) + def test_comfy_full_precision_fp8_dequantizes_before_linear(self): layer = torch.nn.Module() layer.weight = torch.nn.Parameter( @@ -626,7 +686,7 @@ class TestTransformerQuantHelpers(unittest.TestCase): ) ) self.assertTrue(_needs_device_weight_postprocess(_make_quant_config("mxfp8"))) - self.assertFalse( + self.assertTrue( _needs_device_weight_postprocess( _make_quant_config("mxfp8", is_checkpoint_fp8_serialized=True) ) @@ -788,6 +848,7 @@ class TestTransformerQuantHelpers(unittest.TestCase): ) self.assertEqual(config.get_name(), "bitsandbytes") + self.assertIsInstance(config, SRTBitsAndBytesConfig) self.assertTrue(config.load_in_4bit) self.assertEqual(config.bnb_4bit_quant_type, "nf4") @@ -805,6 +866,7 @@ class TestTransformerQuantHelpers(unittest.TestCase): ) self.assertIsInstance(config, Fp8Config) + self.assertIsInstance(config, SRTFp8Config) self.assertTrue(config.is_checkpoint_fp8_serialized) def test_bitsandbytes_quant_config_resolves_from_compression_config(self): diff --git a/python/sglang/srt/layers/quantization/bitsandbytes.py b/python/sglang/srt/layers/quantization/bitsandbytes.py index 85bbe8bc3..d1d35e777 100644 --- a/python/sglang/srt/layers/quantization/bitsandbytes.py +++ b/python/sglang/srt/layers/quantization/bitsandbytes.py @@ -26,6 +26,21 @@ if TYPE_CHECKING: ) +def require_bitsandbytes() -> None: + try: + import bitsandbytes + + if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"): + raise ImportError( + "bitsandbytes version is wrong. Please install bitsandbytes>=0.46.1." + ) + except ImportError as err: + raise ImportError( + "Please install bitsandbytes>=0.46.1 via " + "`pip install bitsandbytes>=0.46.1` to use bitsandbytes quantizer." + ) from err + + class BitsAndBytesConfig(QuantizationConfig): """Config class for BitsAndBytes Quantization. @@ -184,21 +199,7 @@ class BitsAndBytesLinearMethod(LinearMethodBase): """ def __init__(self, quant_config: BitsAndBytesConfig): - try: - import bitsandbytes - - if version.parse(bitsandbytes.__version__) < version.parse("0.46.1"): - raise ImportError( - "bitsandbytes version is wrong. Please " - "install bitsandbytes>=0.46.1." - ) - except ImportError as err: - raise ImportError( - "Please install bitsandbytes>=0.46.1 via " - "`pip install bitsandbytes>=0.46.1` to use " - "bitsandbytes quantizer." - ) from err - + require_bitsandbytes() self.quant_config = quant_config def create_weights(