[diffusion] refactor: reuse srt quantization contracts and mxfp8 kernels (#36063)

2026-08-24 09:26:54 +08:00
parent 167c339c8e
commit 1c1c9d9b4e
16 changed files with 277 additions and 511 deletions
@@ -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 <repo/file.safetensors>` 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
@@ -176,6 +176,14 @@ backend.
      <td><code>comfy-kitchen</code></td>
      <td>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.</td>
    </tr>
    <tr>
      <td><code>mxfp8</code></td>
      <td>Self-describing serialized weights, or BF16/FP16 weights for online quantization</td>
      <td><code>--transformer-weights-path</code>, or <code>--quantization mxfp8</code> for online quantization</td>
      <td>Native diffusion linears; mixed MiniMax-H3 checkpoints are selected per layer</td>
      <td>SRT's platform MXFP8 backend</td>
      <td>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.</td>
    </tr>
    <tr>
      <td><code>qvg-kv</code></td>
      <td>Unquantized model with runtime causal KV-cache compression</td>
@@ -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`.
+Self-describing serialized MXFP8 transformers are inferred from their
-Weights are quantized at load time via `npu_dynamic_mx_quant`, and activations are quantized per-token
+checkpoint metadata. For online MXFP8 quantization, load the original
-during inference with `npu_quant_matmul` (block_size=32).
+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)
@@ -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",
@@ -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",
@@ -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:
+class BitsAndBytesConfig(SRTBitsAndBytesConfig, QuantizationConfig):
    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):
    """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__()
+        super().__init__(
-        self.load_in_8bit = load_in_8bit
+            load_in_8bit=load_in_8bit,
-        self.load_in_4bit = load_in_4bit
+            load_in_4bit=load_in_4bit,
-        self.bnb_4bit_compute_dtype = bnb_4bit_compute_dtype
+            bnb_4bit_compute_dtype=bnb_4bit_compute_dtype,
-        self.bnb_4bit_quant_storage = bnb_4bit_quant_storage
+            bnb_4bit_quant_storage=bnb_4bit_quant_storage,
-        self.bnb_4bit_quant_type = bnb_4bit_quant_type
+            bnb_4bit_quant_type=bnb_4bit_quant_type,
-        self.bnb_4bit_use_double_quant = bnb_4bit_use_double_quant
+            bnb_4bit_use_double_quant=bnb_4bit_use_double_quant,
-        self.llm_int8_enable_fp32_cpu_offload = llm_int8_enable_fp32_cpu_offload
+            llm_int8_enable_fp32_cpu_offload=llm_int8_enable_fp32_cpu_offload,
-        self.llm_int8_has_fp16_weight = llm_int8_has_fp16_weight
+            llm_int8_has_fp16_weight=llm_int8_has_fp16_weight,
-        self.llm_int8_skip_modules = llm_int8_skip_modules or []
+            llm_int8_skip_modules=llm_int8_skip_modules,
-        self.llm_int8_threshold = llm_int8_threshold
+            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(
@@ -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.srt.layers.quantization.base_config import (
-    from sglang.multimodal_gen.runtime.layers.quantization import QuantizationMethods
+    QuantizationConfig as SRTQuantizationConfig,
-else:
+)
-    QuantizationMethods = str
+from sglang.srt.layers.quantization.base_config import (
    QuantizeMethodBase as SRTQuantizeMethodBase,
 )
-class QuantizeMethodBase(ABC):
+class QuantizeMethodBase(SRTQuantizeMethodBase):
    """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
    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):
+class QuantizationConfig(SRTQuantizationConfig):
    """Base class for quantization configs."""
    # 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):
+    def get_scaled_act_names(self) -> list[str]:
-        super().__init__()
+        return []
        # 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 supports_input_partition(
        self, prefix: str, input_size_per_partition: int
@@ -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.
@@ -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 "<layer>.weight_scale".
        scale_dim = input_size_per_partition // MXFP8_BLOCK_SIZE
@@ -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"]
@@ -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
@@ -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 quant_config.get_name() == "bitsandbytes"
    return bool(callable(quant_name_getter) and quant_name_getter() == "bitsandbytes")
 def _format_dtype_mismatch_summary(
@@ -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 config.get_name()
    return quant_name_getter() if callable(quant_name_getter) else None
 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)
+        if _get_quant_config_name(quant_config) != "modelopt_fp4":
        quant_name = quant_name_getter() if callable(quant_name_getter) else None
        if quant_name != "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)
+        if _get_quant_config_name(quant_config) != "modelopt_fp8":
        quant_name = quant_name_getter() if callable(quant_name_getter) else None
        if quant_name != "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",
    }
@@ -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)
    )
@@ -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],
@@ -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):
@@ -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:
+        require_bitsandbytes()
            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
        self.quant_config = quant_config
    def create_weights(