chore: import upstream snapshot with attribution

python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py

@@ -0,0 +1,98 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+from typing import Literal, get_args
+
+from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import (
+    BitsAndBytesConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_fp8 import (
+    ModelOptFp8Config as ModelOptFp8DiffusionConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp4Config,
+    ModelOptFp8Config,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.modelslim import ModelSlimConfig
+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
+
+QuantizationMethods = Literal[
+    "fp8",
+    "modelopt",
+    "modelopt_fp8",
+    "modelopt_fp4",
+    "bitsandbytes",
+    "modelslim",
+    "mxfp8",
+    "mxfp4",
+    "mxfp4_npu",
+]
+
+QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods))
+
+# The customized quantization methods which will be added to this dict.
+_CUSTOMIZED_METHOD_TO_QUANT_CONFIG = {
+    "modelopt": ModelOptFp8DiffusionConfig,
+    "modelopt_fp8": ModelOptFp8Config,
+    "modelopt_fp4": ModelOptFp4Config,
+    "bitsandbytes": BitsAndBytesConfig,
+    "modelslim": ModelSlimConfig,
+    "fp8": Fp8Config,
+    "mxfp4": Mxfp4Config,
+    "mxfp8": MXFP8Config,
+    "mxfp4_npu": NPUMXFP4Config,
+}
+
+
+def register_quantization_config(quantization: str):
+    """Register a customized vllm quantization config.
+
+    When a quantization method is not supported by vllm, you can register a customized
+    quantization config to support it.
+
+    Args:
+        quantization (str): The quantization method name.
+
+
+    """  # noqa: E501
+
+    def _wrapper(quant_config_cls):
+        if quantization in QUANTIZATION_METHODS:
+            raise ValueError(
+                f"The quantization method `{quantization}` is already exists."
+            )
+        if not issubclass(quant_config_cls, QuantizationConfig):
+            raise ValueError(
+                "The quantization config must be a subclass of " "`QuantizationConfig`."
+            )
+        _CUSTOMIZED_METHOD_TO_QUANT_CONFIG[quantization] = quant_config_cls
+        QUANTIZATION_METHODS.append(quantization)
+        return quant_config_cls
+
+    return _wrapper
+
+
+def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
+    if quantization not in QUANTIZATION_METHODS:
+        raise ValueError(f"Invalid quantization method: {quantization}")
+
+    method_to_config: dict[str, type[QuantizationConfig]] = {}
+    # Update the `method_to_config` with customized quantization methods.
+    method_to_config.update(_CUSTOMIZED_METHOD_TO_QUANT_CONFIG)
+
+    return method_to_config[quantization]
+
+
+__all__ = [
+    "QuantizationMethods",
+    "QuantizationConfig",
+    "get_quantization_config",
+    "QUANTIZATION_METHODS",
+]

python/sglang/multimodal_gen/runtime/layers/quantization/bitsandbytes.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+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,
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+
+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):
+    """Config class for pre-quantized bitsandbytes 4-bit checkpoints."""
+
+    def __init__(
+        self,
+        load_in_8bit: bool = False,
+        load_in_4bit: bool = True,
+        bnb_4bit_compute_dtype: str = "float32",
+        bnb_4bit_quant_storage: str = "uint8",
+        bnb_4bit_quant_type: str = "fp4",
+        bnb_4bit_use_double_quant: bool = False,
+        llm_int8_enable_fp32_cpu_offload: bool = False,
+        llm_int8_has_fp16_weight: bool = False,
+        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
+
+        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):
+                return UnquantizedLinearMethod()
+            return BitsAndBytesLinearMethod(self)
+        return None
+
+
+class BitsAndBytesLinearMethod(LinearMethodBase):
+    """Linear method for pre-quantized bitsandbytes 4-bit weights."""
+
+    def __init__(self, quant_config: BitsAndBytesConfig):
+        _require_bitsandbytes()
+        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,
+    ) -> None:
+        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:
+            raise ValueError(
+                "The input size is not aligned with the quantized weight shape."
+            )
+
+        qweight = nn.Parameter(
+            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            qweight,
+            {
+                "input_dim": 0,
+                "output_dim": 0,
+                "pack_factor": quant_ratio,
+                "use_bitsandbytes_4bit": True,
+                "bnb_full_shape": (output_size, input_size),
+                "bnb_local_shape": (
+                    output_size_per_partition,
+                    input_size_per_partition,
+                ),
+                "bnb_output_shard_start": getattr(layer, "tp_rank", 0)
+                * output_size_per_partition,
+                "bnb_input_shard_start": (
+                    0
+                    if input_size_per_partition == input_size
+                    else getattr(layer, "tp_rank", 0) * input_size_per_partition
+                ),
+            },
+        )
+        layer.register_parameter("weight", qweight)
+        set_weight_attrs(qweight, extra_weight_attrs)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        original_type = x.dtype
+        original_shape = x.shape
+        if x.ndim > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        out_dim = sum(
+            quant_state.shape[0]
+            for quant_state in layer.weight.bnb_quant_state.values()
+        )
+        out = torch.empty(x.shape[0], out_dim, dtype=torch.bfloat16, device=x.device)
+        apply_bnb_4bit(x.to(torch.bfloat16), layer.weight, out)
+        out = out.to(original_type)
+
+        if len(original_shape) > 2:
+            out = out.view(*original_shape[:-1], out.size(-1))
+
+        if bias is not None:
+            out = out + bias
+        return out
+
+
+def apply_bnb_4bit(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    out: torch.Tensor,
+) -> None:
+    from bitsandbytes import matmul_4bit
+
+    offsets = weight.bnb_shard_offsets
+    quant_states = weight.bnb_quant_state
+    current_index = 0
+    for i in range(len(quant_states)):
+        output_size = quant_states[i].shape[0]
+        out[:, current_index : current_index + output_size] = matmul_4bit(
+            x,
+            weight[offsets[i] : offsets[i + 1]].t(),
+            quant_states[i],
+        )
+        current_index += output_size
+
+
+class BitsAndBytes4BitLinear(nn.Module):
+    """Storage-only bitsandbytes 4-bit linear for nn.Linear-based encoders."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+    ) -> None:
+        super().__init__()
+        _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())
+        total_size = in_features * out_features
+        if total_size % quant_ratio != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized weight shape."
+            )
+
+        self.weight = nn.Parameter(
+            torch.empty(total_size // quant_ratio, 1, dtype=torch.uint8),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "pack_factor": quant_ratio,
+                "use_bitsandbytes_4bit": True,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        original_type = x.dtype
+        original_shape = x.shape
+        if x.ndim > 2:
+            x = x.reshape(-1, x.size(-1))
+
+        out = torch.empty(
+            x.shape[0], self.out_features, dtype=torch.bfloat16, device=x.device
+        )
+        apply_bnb_4bit(x.to(torch.bfloat16), self.weight, out)
+        out = out.to(original_type)
+
+        if len(original_shape) > 2:
+            out = out.view(*original_shape[:-1], out.size(-1))
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+def swap_linears_to_bitsandbytes_4bit(module: nn.Module) -> None:
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear):
+            replacement = BitsAndBytes4BitLinear(
+                child.in_features,
+                child.out_features,
+                bias=child.bias is not None,
+                compute_dtype=child.weight.dtype,
+            )
+            setattr(module, name, replacement)
+        else:
+            swap_linears_to_bitsandbytes_4bit(child)
+
+
+_BNB_4BIT_STATE_SUFFIXES = {
+    "absmax",
+    "quant_map",
+    "nested_absmax",
+    "nested_quant_map",
+    "bitsandbytes",
+}
+
+
+def is_bitsandbytes_4bit_state_name(weight_name: str) -> bool:
+    suffix = weight_name.split(".")[-1]
+    return any(state_suffix in suffix for state_suffix in _BNB_4BIT_STATE_SUFFIXES)
+
+
+def split_bitsandbytes_4bit_state(
+    weights: Any,
+) -> tuple[list[tuple[str, torch.Tensor]], dict[str, torch.Tensor]]:
+    normal_weights: list[tuple[str, torch.Tensor]] = []
+    quant_state_dict: dict[str, torch.Tensor] = {}
+    for name, tensor in weights:
+        if is_bitsandbytes_4bit_state_name(name):
+            if "quant_state.bitsandbytes" in name:
+                tensor = tensor.cpu().data
+            quant_state_dict[name] = tensor
+            continue
+        normal_weights.append((name, tensor))
+    return normal_weights, quant_state_dict
+
+
+def build_bitsandbytes_4bit_quant_states(
+    normal_weight_names: list[str],
+    quant_state_dict: dict[str, torch.Tensor],
+    device: torch.device,
+    param_names_mapping=None,
+) -> dict[str, Any]:
+    from bitsandbytes.functional import QuantState
+
+    quant_states: dict[str, Any] = {}
+    device_str = str(device)
+    for source_name in normal_weight_names:
+        if (
+            f"{source_name}.quant_state.bitsandbytes__nf4" not in quant_state_dict
+            and f"{source_name}.quant_state.bitsandbytes__fp4" not in quant_state_dict
+        ):
+            continue
+        target_name = source_name
+        if param_names_mapping is not None:
+            target_name, _, _ = param_names_mapping(source_name)
+        state_tensors = {
+            name: tensor
+            for name, tensor in quant_state_dict.items()
+            if name.startswith(f"{source_name}.")
+        }
+        quant_states[target_name] = QuantState.from_dict(
+            state_tensors, device=device_str
+        )
+    return quant_states
+
+
+def attach_bitsandbytes_4bit_quant_states(
+    params_dict: dict[str, torch.nn.Parameter],
+    quant_states: dict[str, Any],
+) -> None:
+    for param_name, quant_state in quant_states.items():
+        param = params_dict.get(param_name)
+        if param is None:
+            raise ValueError(f"Parameter {param_name} not found in the model.")
+
+        quant_state = _maybe_shard_bitsandbytes_4bit_quant_state(param, quant_state)
+        state_by_shard = {0: quant_state}
+        set_weight_attrs(param, {"bnb_quant_state": state_by_shard})
+        offsets = torch.tensor([0, param.numel()]).cpu()
+        set_weight_attrs(param, {"bnb_shard_offsets": offsets})
+
+
+def _maybe_shard_bitsandbytes_4bit_quant_state(
+    param: torch.nn.Parameter,
+    quant_state: Any,
+) -> Any:
+    full_shape = tuple(getattr(param, "bnb_full_shape", tuple(quant_state.shape or ())))
+    local_shape = tuple(getattr(param, "bnb_local_shape", full_shape))
+    if not full_shape or local_shape == full_shape:
+        return quant_state
+
+    output_start = getattr(param, "bnb_output_shard_start", 0)
+    input_start = getattr(param, "bnb_input_shard_start", 0)
+    if input_start != 0 or local_shape[1] != full_shape[1]:
+        raise NotImplementedError(
+            "bitsandbytes 4-bit TP only supports column-parallel output shards."
+        )
+    if getattr(quant_state, "nested", False):
+        raise NotImplementedError(
+            "bitsandbytes 4-bit TP does not support nested quant states."
+        )
+
+    blocksize = quant_state.blocksize
+    start_elem = output_start * full_shape[1]
+    local_numel = local_shape[0] * local_shape[1]
+    if start_elem % blocksize != 0 or local_numel % blocksize != 0:
+        raise ValueError(
+            "bitsandbytes 4-bit TP shard is not aligned to quantization blocks."
+        )
+    start_block = start_elem // blocksize
+    num_blocks = local_numel // blocksize
+    return type(quant_state)(
+        absmax=quant_state.absmax.narrow(0, start_block, num_blocks).contiguous(),
+        shape=torch.Size(local_shape),
+        code=quant_state.code,
+        blocksize=quant_state.blocksize,
+        quant_type=quant_state.quant_type,
+        dtype=quant_state.dtype,
+        offset=None,
+        state2=None,
+    )

python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py

@@ -0,0 +1,155 @@
+# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
+
+# SPDX-License-Identifier: Apache-2.0
+# 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
+
+
+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
+    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:
+    """
+    Not all quant methods have embedding implemented, so we need to check that
+    it exists for our given method. We check this by making sure the function
+    has been changed from the base implementation.
+    """
+    base_embedding = inspect.getattr_static(QuantizeMethodBase, "embedding", None)
+    class_embedding = inspect.getattr_static(method_class, "embedding", None)
+
+    return class_embedding is not None and class_embedding is not base_embedding
+
+
+class QuantizationConfig(ABC):
+    """Base class for quantization configs."""
+
+    # for quantization frameworks with a separate quantized model provided, e.g. Nunchaku
+    quantized_model_path: str | None = None
+
+    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

python/sglang/multimodal_gen/runtime/layers/quantization/configs/nunchaku_config.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+import json
+import os
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Any, Optional
+
+import torch
+from safetensors.torch import load_file as safetensors_load_file
+from torch import nn
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+from .base_config import QuantizationConfig, QuantizeMethodBase
+
+logger = init_logger(__name__)
+
+
+@lru_cache(maxsize=1)
+def is_nunchaku_available() -> bool:
+    try:
+        import nunchaku  # noqa
+
+        logger.debug("Nunchaku package detected")
+        return True
+    except Exception:
+        return False
+
+
+@dataclass
+class NunchakuConfig(QuantizationConfig):
+    """
+    Configuration for Nunchaku (SVDQuant) W4A4-style quantization.
+
+    Attributes:
+        precision: Quantization precision type. Options:
+            - "int4": Standard INT4 quantization
+            - "nvfp4": FP4 quantization
+        rank: SVD low-rank dimension for absorbing outliers
+        group_size: Quantization group size (automatically set based on precision)
+        act_unsigned: Use unsigned activation quantization
+        transformer_weights_path: Path to pre-quantized transformer weights (.safetensors)
+        model_cls: DiT model class that provides quantization rules via get_nunchaku_quant_rules()
+    """
+
+    precision: str = "int4"
+    rank: int = 32
+    group_size: Optional[int] = None
+    act_unsigned: bool = False
+    transformer_weights_path: Optional[str] = None
+    model_cls: Optional[type] = None
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "svdquant"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 70
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return ["quantization_config.json", "quant_config.json"]
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "NunchakuConfig":
+
+        return cls(
+            precision=config.get("precision", "int4"),
+            rank=int(config.get("rank", 32)),
+            group_size=config.get("group_size"),
+            act_unsigned=bool(config.get("act_unsigned", False)),
+            transformer_weights_path=config.get("transformer_weights_path"),
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if not isinstance(layer, LinearBase):
+            return None
+
+        # get quantization rules from model class
+        quant_rules = self._get_quant_rules()
+
+        # priority: skip > awq_w4a16 > svdq_w4a4 > default
+        skip_patterns = quant_rules.get("skip", [])
+        for pattern in skip_patterns:
+            if pattern in prefix.lower():
+                return None
+
+        awq_patterns = quant_rules.get("awq_w4a16", [])
+        for pattern in awq_patterns:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuAWQLinearMethod
+
+                return NunchakuAWQLinearMethod(group_size=64)
+
+        svdq_patterns = quant_rules.get("svdq_w4a4", [])
+        for pattern in svdq_patterns:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+                return NunchakuSVDQLinearMethod(
+                    precision=self.precision,
+                    rank=self.rank,
+                    act_unsigned=self.act_unsigned,
+                )
+
+        # default: apply svdq_w4a4 to all remaining linear layers
+        from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+        return NunchakuSVDQLinearMethod(
+            precision=self.precision,
+            rank=self.rank,
+            act_unsigned=self.act_unsigned,
+        )
+
+    def _get_quant_rules(self) -> dict[str, list[str]]:
+        if self.model_cls is not None and hasattr(
+            self.model_cls, "get_nunchaku_quant_rules"
+        ):
+            return self.model_cls.get_nunchaku_quant_rules()
+        return {}
+
+    def __post_init__(self):
+        if self.group_size is None:
+            if self.precision == "nvfp4":
+                self.group_size = 16
+            elif self.precision == "int4":
+                self.group_size = 64
+            else:
+                raise ValueError(
+                    f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'"
+                )
+
+        if self.precision not in ["int4", "nvfp4"]:
+            raise ValueError(
+                f"Invalid precision: {self.precision}. Must be 'int4' or 'nvfp4'"
+            )
+
+        if self.rank <= 0:
+            raise ValueError(f"Rank must be positive, got {self.rank}")
+
+    @classmethod
+    def from_dict(cls, config_dict: dict) -> "NunchakuConfig":
+        """Create configuration from dictionary."""
+        return cls(**config_dict)
+
+    def to_dict(self) -> dict:
+        """Convert configuration to dictionary."""
+        return {
+            "precision": self.precision,
+            "rank": self.rank,
+            "group_size": self.group_size,
+            "act_unsigned": self.act_unsigned,
+            "transformer_weights_path": self.transformer_weights_path,
+        }
+
+    @classmethod
+    def from_pretrained(cls, model_path: str) -> Optional["NunchakuConfig"]:
+        for filename in cls.get_config_filenames():
+            config_path = os.path.join(model_path, filename)
+            if os.path.exists(config_path):
+                with open(config_path, "r") as f:
+                    config_dict = json.load(f)
+                if config_dict.get("quant_method") == cls.get_name():
+                    return cls.from_config(config_dict)
+        return None
+
+
+def _patch_native_svdq_linear(
+    module: nn.Module, tensor: Any, svdq_linear_cls: type
+) -> bool:
+    if (
+        isinstance(module, svdq_linear_cls)
+        and getattr(module, "wtscale", None) is not None
+    ):
+        module.wtscale = tensor
+        return True
+    return False
+
+
+def _patch_sglang_svdq_linear(
+    module: nn.Module, tensor: Any, svdq_method_cls: type
+) -> bool:
+    quant_method = getattr(module, "quant_method", None)
+    if not isinstance(quant_method, svdq_method_cls):
+        return False
+
+    existing = getattr(module, "wtscale", None)
+    if isinstance(existing, nn.Parameter):
+        with torch.no_grad():
+            existing.data.copy_(tensor.to(existing.data.dtype))
+    else:
+        module.wtscale = tensor
+
+    # Keep alpha in sync (kernel reads `layer._nunchaku_alpha`)
+    try:
+        module._nunchaku_alpha = float(tensor.detach().cpu().item())
+    except Exception:
+        module._nunchaku_alpha = None
+    return True
+
+
+def _patch_sglang_svdq_wcscales(
+    module: nn.Module, tensor: Any, svdq_method_cls: type
+) -> bool:
+    quant_method = getattr(module, "quant_method", None)
+    if not isinstance(quant_method, svdq_method_cls):
+        return False
+
+    existing = getattr(module, "wcscales", None)
+    if isinstance(existing, nn.Parameter):
+        with torch.no_grad():
+            existing.data.copy_(tensor.to(existing.data.dtype))
+    else:
+        module.wcscales = tensor
+    return True
+
+
+def _patch_nunchaku_scales(
+    model: nn.Module,
+    safetensors_list: list[str],
+) -> None:
+    """Patch transformer module with Nunchaku scale tensors from safetensors weights.
+
+    For NVFP4 checkpoints, correctness depends on `wtscale` and attention
+    `wcscales`. The FSDP loader may skip some of these metadata tensors.
+    """
+
+    if not safetensors_list:
+        return
+
+    if len(safetensors_list) != 1:
+        logger.warning(
+            "Nunchaku scale patch expects a single safetensors file, "
+            "but got %d files. Skipping.",
+            len(safetensors_list),
+        )
+        return
+
+    from nunchaku.models.linear import SVDQW4A4Linear  # type: ignore[import]
+
+    state_dict = safetensors_load_file(safetensors_list[0])
+    if state_dict is None:
+        return
+
+    num_wtscale = 0
+    num_wcscales = 0
+
+    from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+    for name, module in model.named_modules():
+        wt = state_dict.get(f"{name}.wtscale")
+        if wt is not None:
+            if _patch_native_svdq_linear(module, wt, SVDQW4A4Linear):
+                num_wtscale += 1
+            elif _patch_sglang_svdq_linear(module, wt, NunchakuSVDQLinearMethod):
+                num_wtscale += 1
+
+        wc = state_dict.get(f"{name}.wcscales")
+        if wc is not None:
+            # Some modules may have wcscales as a direct attribute/Parameter.
+            existing = getattr(module, "wcscales", None)
+            if isinstance(existing, nn.Parameter):
+                with torch.no_grad():
+                    existing.data.copy_(wc.to(existing.data.dtype))
+                num_wcscales += 1
+            elif existing is not None:
+                setattr(module, "wcscales", wc)
+                num_wcscales += 1
+            elif _patch_sglang_svdq_wcscales(module, wc, NunchakuSVDQLinearMethod):
+                num_wcscales += 1
+
+    if num_wtscale > 0:
+        logger.info("Patched wtscale for %d layers", num_wtscale)
+    if num_wcscales > 0:
+        logger.info("Patched wcscales for %d layers", num_wcscales)

python/sglang/multimodal_gen/runtime/layers/quantization/fp8.py

@@ -0,0 +1,508 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_tensor_model_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    BlockQuantScaleParameter,
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.platforms.aiter import USE_AITER
+from sglang.multimodal_gen.runtime.utils.common import (
+    cpu_has_amx_support,
+    get_bool_env_var,
+    use_intel_amx_backend,
+)
+from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
+from sglang.srt.layers.quantization.fp8_kernel import (
+    is_fp8_fnuz,
+    per_token_group_quant_fp8,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    can_auto_enable_marlin_fp8,
+    cutlass_fp8_supported,
+    dispatch_w8a8_block_fp8_linear,
+    input_to_float8,
+    normalize_e4m3fn_to_e4m3fnuz,
+    requant_weight_ue8m0_inplace,
+)
+from sglang.srt.layers.quantization.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear,
+    prepare_fp8_layer_for_marlin,
+)
+from sglang.srt.layers.quantization.utils import (
+    convert_to_channelwise,
+    is_layer_skipped,
+    requantize_with_max_scale,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
+
+_is_hip = current_platform.is_hip()
+_is_cuda = current_platform.is_cuda()
+_is_npu = current_platform.is_npu()
+_is_cpu_amx_available = cpu_has_amx_support()
+_is_cpu = current_platform.is_cpu()
+_is_fp8_fnuz = is_fp8_fnuz()
+_use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip
+
+if USE_AITER or _use_hip_int4:
+    pass
+
+
+ACTIVATION_SCHEMES = ["static", "dynamic"]
+
+logger = logging.getLogger(__name__)
+
+
+class Fp8Config(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,
+                self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+            return Fp8LinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class Fp8LinearMethod(LinearMethodBase):
+    """Linear method for FP8.
+    Supports loading FP8 checkpoints with static weight scale and
+    dynamic/static activation scale.
+
+    Also supports loading quantized FP16/BF16 model checkpoints with dynamic
+    activation scaling. The weight scaling factor will be initialized after
+    the model weights are loaded.
+
+    Limitations:
+    1. Only support per-tensor quantization due to torch._scaled_mm support.
+    2. Only support float8_e4m3fn data type due to the limitation of
+       torch._scaled_mm (https://github.com/pytorch/pytorch/blob/2e48b39603411a41c5025efbe52f89560b827825/aten/src/ATen/native/cuda/Blas.cpp#L854-L856)
+
+    Args:
+        quant_config: The quantization config.
+    """
+
+    def __init__(self, quant_config: Union[Fp8Config, W4AFp8Config]):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+        # For GPUs that lack FP8 hardware support, we can leverage the Marlin
+        # kernel for fast weight-only FP8 quantization
+        self.use_marlin = False
+        if _is_cuda:
+            force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN")
+            auto_enable = can_auto_enable_marlin_fp8()
+            self.use_marlin = force_marlin or auto_enable
+
+        self.block_quant = self.quant_config.weight_block_size is not None
+
+        self.w8a8_block_fp8_linear = dispatch_w8a8_block_fp8_linear()
+
+    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")
+
+        tp_size = get_tensor_model_parallel_world_size()
+        if self.block_quant:
+            block_n, block_k = (
+                self.quant_config.weight_block_size[0],
+                self.quant_config.weight_block_size[1],
+            )
+            # Required by row parallel
+            if tp_size > 1 and input_size // input_size_per_partition == tp_size:
+                if input_size_per_partition % block_k != 0:
+                    raise ValueError(
+                        f"Weight input_size_per_partition = "
+                        f"{input_size_per_partition} is not divisible by "
+                        f"weight quantization block_k = {block_k}."
+                    )
+            # Required by column parallel or enabling merged weights
+            if (
+                tp_size > 1 and output_size // output_size_per_partition == tp_size
+            ) or len(output_partition_sizes) > 1:
+                for output_partition_size in output_partition_sizes:
+                    if output_partition_size % block_n != 0:
+                        raise ValueError(
+                            f"Weight output_partition_size = "
+                            f"{output_partition_size} is not divisible by "
+                            f"weight quantization block_n = {block_n}."
+                        )
+
+        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
+
+        # WEIGHT
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_fp8_serialized
+            else params_dtype
+        )
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition, input_size_per_partition, dtype=weight_dtype
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # If checkpoint is serialized fp8, load them.
+        # Otherwise, wait until process_weights_after_loading.
+        if self.quant_config.is_checkpoint_fp8_serialized:
+            # WEIGHT SCALE
+            if self.block_quant:
+                if hasattr(self.quant_config, "activation_scheme"):
+                    assert self.quant_config.activation_scheme == "dynamic"
+                elif hasattr(self.quant_config, "linear_activation_scheme"):
+                    assert self.quant_config.linear_activation_scheme == "dynamic"
+                scale = BlockQuantScaleParameter(
+                    data=torch.empty(
+                        (output_size_per_partition + block_n - 1) // block_n,
+                        (input_size_per_partition + block_k - 1) // block_k,
+                        dtype=torch.float32,
+                    ),
+                    input_dim=1,
+                    output_dim=0,
+                    weight_loader=weight_loader,
+                )
+                scale.format_ue8m0 = False
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale_inv", scale)
+            else:
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("weight_scale", scale)
+
+            # INPUT ACTIVATION SCALE
+            if (
+                hasattr(self.quant_config, "activation_scheme")
+                and self.quant_config.activation_scheme == "static"
+            ) or (
+                hasattr(self.quant_config, "linear_activation_scheme")
+                and self.quant_config.linear_activation_scheme == "static"
+            ):
+                scale = PerTensorScaleParameter(
+                    data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+                    weight_loader=weight_loader,
+                )
+
+                scale[:] = torch.finfo(torch.float32).min
+                layer.register_parameter("input_scale", scale)
+            else:
+                layer.register_parameter("input_scale", None)
+
+    def process_weights_after_loading(self, layer: Module) -> None:
+        if self.block_quant:
+            # If ROCm, normalize the weights and scales to e4m3fnuz
+            if _is_fp8_fnuz:
+                # activation_scheme: dynamic
+                weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
+                    weight=layer.weight,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=None,
+                )
+                layer.input_scale = None
+            elif _is_cpu:
+                assert (
+                    _is_cpu_amx_available
+                ), "Fp8LinearMethod on CPU requires that CPU has AMX support"
+                _amx_process_weight_after_loading(layer, ["weight"])
+                layer.weight_scale_inv = torch.nn.Parameter(
+                    layer.weight_scale_inv.data, requires_grad=False
+                )
+                return
+            else:
+                # For fp8 linear weights run with deepgemm, the weights and scales need be requantized to ue8m0
+                from sglang.srt.layers.quantization.fp8_utils import (
+                    deepgemm_w8a8_block_fp8_linear_with_fallback,
+                )
+                from sglang.srt.model_loader.utils import (
+                    should_deepgemm_weight_requant_ue8m0,
+                )
+
+                if (
+                    should_deepgemm_weight_requant_ue8m0(
+                        weight_block_size=getattr(
+                            self.quant_config, "weight_block_size", None
+                        ),
+                    )
+                    and (
+                        self.w8a8_block_fp8_linear
+                        is deepgemm_w8a8_block_fp8_linear_with_fallback
+                    )
+                    and (not layer.weight_scale_inv.format_ue8m0)
+                ):
+                    requant_weight_ue8m0_inplace(
+                        layer.weight,
+                        layer.weight_scale_inv,
+                        self.quant_config.weight_block_size,
+                    )
+                    layer.weight_scale_inv.format_ue8m0 = True
+                weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data
+
+            layer.weight.data = weight.data
+            layer.weight_scale_inv.data = weight_scale.data
+        else:
+            layer.weight = Parameter(layer.weight.data, requires_grad=False)
+
+            # If checkpoint not serialized fp8, quantize the weights.
+            if not self.quant_config.is_checkpoint_fp8_serialized:
+                if self.cutlass_fp8_supported or self.use_marlin:
+                    # apply per-channel quantization default as
+                    # cutlass sgl-kernel and marlin only support per-channel scale
+                    qweight, weight_scale = per_token_group_quant_fp8(
+                        layer.weight, layer.weight.shape[-1]
+                    )
+                    weight_scale = weight_scale.t().contiguous()
+                else:
+                    # per-tensor quantization
+                    qweight, weight_scale = input_to_float8(layer.weight)
+
+                # Update the layer with the new values.
+                layer.weight = Parameter(qweight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                layer.input_scale = None
+
+            # If checkpoint is fp8, handle that there are N scales for N
+            # shards in a fused module
+            else:
+                layer.weight_scale = Parameter(
+                    layer.weight_scale.data, requires_grad=False
+                )
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.data, requires_grad=False
+                    )
+
+                # cutlass sgl-kernel and marlin only support per-channel scale
+                if self.cutlass_fp8_supported or self.use_marlin:
+                    weight = layer.weight
+                    weight_scale = convert_to_channelwise(
+                        layer.weight_scale, layer.logical_widths
+                    )
+                else:
+                    # Dequant -> Quant with max scale so we can run per tensor.
+                    weight = layer.weight
+                    weight_scale = layer.weight_scale
+                    # If ROCm, normalize the weights and scales to e4m3fnuz
+                    if _is_fp8_fnuz:
+                        weight, weight_scale, input_scale = (
+                            normalize_e4m3fn_to_e4m3fnuz(
+                                weight=weight,
+                                weight_scale=weight_scale,
+                                input_scale=layer.input_scale,
+                            )
+                        )
+                        if input_scale is not None:
+                            layer.input_scale = Parameter(
+                                input_scale, requires_grad=False
+                            )
+
+                    weight_scale, weight = requantize_with_max_scale(
+                        weight=weight,
+                        weight_scale=weight_scale,
+                        logical_widths=layer.logical_widths,
+                    )
+
+                # Update layer with new values.
+                layer.weight = Parameter(weight.t(), requires_grad=False)
+                layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+                if (
+                    hasattr(self.quant_config, "activation_scheme")
+                    and self.quant_config.activation_scheme == "static"
+                ) or (
+                    hasattr(self.quant_config, "linear_activation_scheme")
+                    and self.quant_config.linear_activation_scheme == "static"
+                ):
+                    layer.input_scale = Parameter(
+                        layer.input_scale.max(), requires_grad=False
+                    )
+
+        if self.use_marlin:
+            if self.block_quant:
+                layer.weight_block_size = self.quant_config.weight_block_size
+            prepare_fp8_layer_for_marlin(layer, not self.block_quant)
+            # Activations not quantized for marlin.
+            del layer.input_scale
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_marlin:
+            return apply_fp8_marlin_linear(
+                input=x,
+                weight=layer.weight,
+                weight_scale=layer.weight_scale,
+                workspace=layer.workspace,
+                size_n=layer.output_size_per_partition,
+                size_k=layer.input_size_per_partition,
+                bias=bias,
+            )
+
+        if self.block_quant:
+            if use_intel_amx_backend(layer):
+                return torch.ops.sgl_kernel.fp8_scaled_mm_cpu(
+                    x,
+                    layer.weight,
+                    layer.weight_scale_inv,
+                    self.quant_config.weight_block_size,
+                    bias,
+                    x.dtype,
+                    True,  # is_vnni
+                )
+
+            if isinstance(x, tuple):
+                return self.w8a8_block_fp8_linear(
+                    input=x[0],
+                    weight=layer.weight,
+                    block_size=self.quant_config.weight_block_size,
+                    weight_scale=layer.weight_scale_inv,
+                    input_scale=x[1],
+                    bias=bias,
+                )
+
+            return self.w8a8_block_fp8_linear(
+                input=x,
+                weight=layer.weight,
+                block_size=self.quant_config.weight_block_size,
+                weight_scale=layer.weight_scale_inv,
+                input_scale=None,
+                bias=bias,
+            )
+
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+            use_per_token_if_dynamic=False,
+        )

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_fp8.py

@@ -0,0 +1,210 @@
+"""ModelOpt FP8 quantization support for diffusion models.
+
+Handles checkpoints produced by NVIDIA Model Optimizer (ModelOpt) with
+``quant_algo: "FP8"`` and ``quant_method: "modelopt"``.
+
+Per quantized linear layer the checkpoint contains:
+    .weight         float8_e4m3fn  [out, in]   FP8 quantized weight
+    .weight_scale   float32        scalar       per-tensor weight scale
+    .input_scale    float32        scalar       per-tensor static activation scale
+    .bias           bfloat16       [out]        bias (unquantized)
+    ._amax          (ignored)                   calibration artifact
+
+Layers listed in the ``ignore`` field of the quantization config remain in
+bfloat16 and use the standard unquantized linear method.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import logging
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    cutlass_fp8_supported,
+)
+from sglang.srt.layers.quantization.utils import convert_to_channelwise
+
+logger = logging.getLogger(__name__)
+
+
+class ModelOptFp8Config(QuantizationConfig):
+    """Config for ModelOpt static per-tensor FP8 quantization."""
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = True,
+        ignore: Optional[List[str]] = None,
+    ) -> None:
+        super().__init__()
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        self.ignore = ignore or []
+
+    # -- QuantizationConfig interface ----------------------------------------
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @staticmethod
+    def get_config_filenames() -> list[str]:
+        return []
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
+        quant_algo = config.get("quant_algo")
+        if quant_algo is None:
+            raise ValueError(
+                "ModelOptFp8Config requires 'quant_algo' in the quantization config."
+            )
+        if "FP8" not in quant_algo:
+            raise ValueError(
+                f"ModelOptFp8Config only supports FP8, got quant_algo={quant_algo!r}."
+            )
+        ignore = config.get("ignore", [])
+        if ignore_remap and ignore:
+            ignore = [ignore_remap.get(pattern, pattern) for pattern in ignore]
+        return cls(is_checkpoint_fp8_serialized=True, ignore=ignore)
+
+    def _is_layer_ignored(self, prefix: str) -> bool:
+        """Check whether *prefix* matches any pattern in the ignore list.
+
+        ModelOpt ignore patterns are matched against the full prefix as a glob
+        (e.g. ``"norm_out*"`` matches ``"norm_out.linear"``) **and** against the
+        first path component (e.g. ``"proj_out"`` matches only the top-level
+        ``proj_out``, not ``single_transformer_blocks.0.proj_out``).
+        """
+        first_component = prefix.split(".")[0]
+        for pattern in self.ignore:
+            if fnmatch.fnmatch(prefix, pattern):
+                return True
+            if fnmatch.fnmatch(first_component, pattern):
+                return True
+        return False
+
+    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 self._is_layer_ignored(prefix):
+                return UnquantizedLinearMethod()
+            return ModelOptFp8LinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> list[str]:
+        return []
+
+
+class ModelOptFp8LinearMethod(LinearMethodBase):
+    """Linear method for ModelOpt static per-tensor FP8 quantization.
+
+    Uses ``torch._scaled_mm`` (or CUTLASS FP8 GEMM when available) for
+    the FP8 matrix multiply - the same kernels used by the LLM runtime.
+    """
+
+    def __init__(self, quant_config: ModelOptFp8Config):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+    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,
+    ) -> None:
+        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
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        for scale_name in ("weight_scale", "input_scale"):
+            scale = PerTensorScaleParameter(
+                data=torch.full(
+                    (len(output_partition_sizes),),
+                    torch.finfo(torch.float32).min,
+                    dtype=torch.float32,
+                ),
+                weight_loader=weight_loader,
+            )
+            layer.register_parameter(scale_name, scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Diffusion models use single-partition layers (no TP, no fused QKV),
+        # so we just take the max scale directly without the
+        # dequantize-requantize round-trip that the LLM path does (which
+        # requires CUDA kernels that are unavailable during CPU-phase loading).
+        max_w_scale = layer.weight_scale.max()
+
+        # Transpose weight to [in, out] column-major layout for
+        # apply_fp8_linear / CUTLASS fp8_scaled_mm. Do not call .contiguous();
+        # the kernel requires column-major stride.
+        layer.weight = torch.nn.Parameter(layer.weight.data.t(), requires_grad=False)
+
+        if self.cutlass_fp8_supported:
+            max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths)
+        layer.weight_scale = torch.nn.Parameter(max_w_scale, requires_grad=False)
+        layer.input_scale = torch.nn.Parameter(
+            layer.input_scale.max(), requires_grad=False
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+        )

python/sglang/multimodal_gen/runtime/layers/quantization/modelopt_quant.py

@@ -0,0 +1,683 @@
+# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/modelopt_quant.py
+from __future__ import annotations
+
+import logging
+import re
+from functools import lru_cache
+from typing import Any, Dict, List, Optional
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+from sglang.srt.layers.quantization.fp8_utils import (
+    apply_fp8_linear,
+    cutlass_fp8_supported,
+)
+from sglang.srt.layers.quantization.modelopt_quant import (
+    pad_nvfp4_activation_for_cutlass,
+    pad_nvfp4_weight,
+    slice_nvfp4_output,
+)
+from sglang.srt.layers.quantization.utils import (
+    convert_to_channelwise,
+    is_layer_skipped,
+    requantize_with_max_scale,
+)
+from sglang.srt.layers.utils.common import copy_or_rebind_param
+from sglang.srt.utils.common import is_flashinfer_available, round_up
+
+logger = logging.getLogger(__name__)
+
+if is_flashinfer_available():
+    import flashinfer
+else:
+    flashinfer = None
+
+
+@lru_cache(maxsize=1)
+def _get_fp4_quantize_op():
+    return current_platform.get_modelopt_fp4_quantize_op()
+
+
+@lru_cache(maxsize=1)
+def _get_fp4_gemm_op():
+    return current_platform.get_modelopt_fp4_gemm_op()
+
+
+def _prepare_nvfp4_weight_bytes(
+    weight: torch.Tensor, *, swap_weight_nibbles: bool
+) -> torch.Tensor:
+    """Normalize serialized NVFP4 bytes before padding for the runtime kernel."""
+    if not swap_weight_nibbles:
+        return weight.contiguous()
+    return ((weight >> 4) | (weight << 4)).contiguous()
+
+
+def _swizzled_nvfp4_scales_to_linear(scales: torch.Tensor) -> torch.Tensor:
+    """Convert FlashInfer/CUTLASS-swizzled FP4 scales back to row-major layout."""
+    scale_ndim = scales.ndim
+    if scale_ndim == 2:
+        scales = scales.unsqueeze(0)
+    assert scales.ndim == 3
+
+    B, M, K = scales.shape
+    M_padded = round_up(M, 128)
+    K_padded = round_up(K, 4)
+    if M != M_padded or K != K_padded:
+        padded = torch.zeros(
+            (B, M_padded, K_padded), dtype=scales.dtype, device=scales.device
+        )
+        padded[:B, :M, :K] = scales
+        scales = padded
+
+    linear = scales.reshape(B, M_padded // 128, K_padded // 4, 32, 4, 4)
+    linear = linear.permute(0, 1, 4, 3, 2, 5).contiguous()
+    linear = linear.reshape(B, M_padded, K_padded)[:, :M, :K]
+    return linear.squeeze(0) if scale_ndim == 2 else linear
+
+
+def _require_flashinfer():
+    if flashinfer is None:
+        raise RuntimeError(
+            "flashinfer is required for the diffusion NVFP4 FlashInfer path."
+        )
+    return flashinfer
+
+
+class ModelOptQuantConfig(QuantizationConfig):
+    def __init__(
+        self,
+        exclude_modules: Optional[List[str]],
+        packed_modules_mapping: Optional[Dict[str, List[str]]],
+    ):
+        super().__init__()
+        self.packed_modules_mapping = packed_modules_mapping or {}
+        self.exclude_modules = exclude_modules or []
+
+    def _get_quant_method(
+        self,
+        layer: torch.nn.Module,
+        prefix: str,
+        *,
+        Linear: type[LinearMethodBase],
+    ) -> Optional[QuantizeMethodBase]:
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if self.is_layer_excluded(prefix) or (
+                self.packed_modules_mapping
+                and is_layer_skipped(prefix, [], self.packed_modules_mapping)
+            ):
+                return UnquantizedLinearMethod()
+            return Linear(self)
+        return None
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        return ["hf_quant_config.json"]
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+    @classmethod
+    def override_quantization_method(cls, hf_quant_config, user_quant) -> Optional[str]:
+        if hf_quant_config is None:
+            return None
+
+        quant_algo = (
+            hf_quant_config.get("quant_algo")
+            or hf_quant_config.get("quantization", {}).get("quant_algo")
+            or ""
+        ).upper()
+        if user_quant in {"modelopt", "modelopt_fp8"} and "FP8" in quant_algo:
+            return "modelopt_fp8"
+        if user_quant in {"modelopt", "modelopt_fp4"} and (
+            "NVFP4" in quant_algo or "FP4" in quant_algo
+        ):
+            return "modelopt_fp4"
+        return None
+
+    def is_layer_excluded(self, prefix: str) -> bool:
+        for pattern in self.exclude_modules:
+            regex_str = re.escape(pattern).replace(r"\*", r".*")
+            if re.fullmatch(regex_str, prefix):
+                return True
+        return False
+
+
+class ModelOptFp8Config(ModelOptQuantConfig):
+    """Config class for ModelOpt FP8 diffusion checkpoints."""
+
+    def __init__(
+        self,
+        is_checkpoint_fp8_serialized: bool = False,
+        exclude_modules: Optional[List[str]] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+    ) -> None:
+        super().__init__(exclude_modules, packed_modules_mapping)
+        self.is_checkpoint_fp8_serialized = is_checkpoint_fp8_serialized
+        if is_checkpoint_fp8_serialized:
+            logger.warning(
+                "Detected ModelOpt FP8 checkpoint. The format is experimental and subject to change."
+            )
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt_fp8"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 89
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Dict[str, Any],
+        ignore_remap: Optional[Dict[str, str]] = None,
+    ) -> ModelOptFp8Config:
+        quant_method = config.get("quant_algo")
+        exclude_modules = config.get("ignore")
+        if quant_method is None:
+            try:
+                quantization_section = cls.get_from_keys(config, ["quantization"])
+                quant_method = quantization_section.get("quant_algo")
+                exclude_modules = quantization_section.get("exclude_modules")
+            except ValueError as exc:
+                raise ValueError(
+                    "Cannot find 'quant_algo' in the model's quantization config."
+                ) from exc
+
+        if quant_method is None or "FP8" not in quant_method:
+            raise ValueError(
+                "ModelOptFp8Config only supports static FP8 quantization in SGLang diffusion."
+            )
+
+        if ignore_remap and exclude_modules:
+            exclude_modules = [ignore_remap.get(p, p) for p in exclude_modules]
+
+        return cls(
+            is_checkpoint_fp8_serialized=True,
+            exclude_modules=exclude_modules,
+            packed_modules_mapping=config.get("packed_modules_mapping"),
+        )
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        return self._get_quant_method(layer, prefix, Linear=ModelOptFp8LinearMethod)
+
+
+class ModelOptFp4Config(ModelOptQuantConfig):
+    """Config class for NVFP4."""
+
+    def __init__(
+        self,
+        is_checkpoint_nvfp4_serialized: bool = False,
+        group_size: int = None,
+        exclude_modules: List[str] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+        checkpoint_uses_packed_qkv: bool = False,
+        swap_weight_nibbles: bool = False,
+        checkpoint_weight_scale_layout: str = "linear",
+    ) -> None:
+        super().__init__(exclude_modules, packed_modules_mapping)
+        self.is_checkpoint_nvfp4_serialized = is_checkpoint_nvfp4_serialized
+        if is_checkpoint_nvfp4_serialized:
+            logger.warning(
+                "Detected nvfp4 checkpoint. Please note that the "
+                "format is experimental and subject to change."
+            )
+        self.group_size = group_size
+        self.checkpoint_uses_packed_qkv = checkpoint_uses_packed_qkv
+        self.swap_weight_nibbles = swap_weight_nibbles
+        self.checkpoint_weight_scale_layout = checkpoint_weight_scale_layout
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelopt_fp4"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.bfloat16, torch.half, torch.float8_e4m3fn]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 100
+
+    @staticmethod
+    def common_group_size(cfg: dict) -> int:
+        """Return the unique group_size across the config; raise if missing/mismatched."""
+        sizes = set()
+
+        def _add_group_size_from_dict(config: dict):
+            group_size = config.get("group_size")
+            if isinstance(group_size, int):
+                sizes.add(group_size)
+
+        # Top-level and 'quantization' block
+        _add_group_size_from_dict(cfg)
+        quantization = cfg.get("quantization")
+        if isinstance(quantization, dict):
+            _add_group_size_from_dict(quantization)
+
+        # config_groups: accept group-level or nested dicts (e.g., weights/input_activations)
+        for config_groups in (cfg.get("config_groups") or {}).values():
+            if isinstance(config_groups, dict):
+                _add_group_size_from_dict(config_groups)
+                for config_group in config_groups.values():
+                    if isinstance(config_group, dict):
+                        _add_group_size_from_dict(config_group)
+
+        if not sizes:
+            raise ValueError("No group_size found in config.")
+        if len(sizes) > 1:
+            raise ValueError(f"Inconsistent group_size values: {sorted(sizes)}")
+        return next(iter(sizes))
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> ModelOptFp4Config:
+        group_size = None
+        exclude_modules = []
+        swap_weight_nibbles = False
+
+        # Flat format (config.json quantization_config)
+        quant_method = config.get("quant_algo")
+        if quant_method is not None:
+            group_size = config.get("group_size")
+            if group_size is None:
+                config_groups = config.get("config_groups", {})
+                if config_groups:
+                    first_group = next(iter(config_groups.values()), {})
+                    group_size = first_group.get("weights", {}).get("group_size")
+            exclude_modules = config.get("ignore", [])
+            swap_weight_nibbles = config.get(
+                "swap_weight_nibbles",
+                config.get("checkpoint_uses_packed_qkv", False),
+            )
+        else:
+            # Nested format (hf_quant_config.json)
+            try:
+                quant_config = cls.get_from_keys(config, ["quantization"])
+                quant_method = quant_config["quant_algo"]
+                group_size = ModelOptFp4Config.common_group_size(config)
+                exclude_modules = quant_config.get("exclude_modules", [])
+                swap_weight_nibbles = quant_config.get(
+                    "swap_weight_nibbles",
+                    config.get(
+                        "swap_weight_nibbles",
+                        config.get("checkpoint_uses_packed_qkv", False),
+                    ),
+                )
+            except (ValueError, KeyError):
+                raise ValueError("Cannot find 'quant_algo' in quantization config.")
+
+        if quant_method not in ["NVFP4"]:
+            raise ValueError(
+                f"Only NVFP4 quantization is supported for diffusion, got '{quant_method}'."
+            )
+
+        if group_size is None or exclude_modules is None:
+            raise ValueError(
+                "NVFP4 quantization requires group_size and exclude_modules "
+                "in the quantization config"
+            )
+        return cls(
+            is_checkpoint_nvfp4_serialized=True,
+            group_size=group_size,
+            exclude_modules=exclude_modules,
+            packed_modules_mapping=config.get("packed_modules_mapping"),
+            checkpoint_uses_packed_qkv=config.get("checkpoint_uses_packed_qkv", False),
+            swap_weight_nibbles=swap_weight_nibbles,
+            checkpoint_weight_scale_layout=config.get(
+                "checkpoint_weight_scale_layout", "linear"
+            ),
+        )
+
+    def get_quant_method(self, layer: torch.nn.Module, prefix: str):
+        return self._get_quant_method(layer, prefix, Linear=ModelOptFp4LinearMethod)
+
+
+class ModelOptFp8LinearMethod(LinearMethodBase):
+    """Linear method for ModelOpt static FP8 checkpoints."""
+
+    def __init__(self, quant_config: ModelOptFp8Config):
+        self.quant_config = quant_config
+        self.cutlass_fp8_supported = cutlass_fp8_supported()
+
+    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,
+    ):
+        del input_size, output_size
+        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
+
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_fp8_serialized
+            else params_dtype
+        )
+        layer.register_parameter(
+            "weight",
+            ModelWeightParameter(
+                data=torch.empty(
+                    output_size_per_partition,
+                    input_size_per_partition,
+                    dtype=weight_dtype,
+                ),
+                input_dim=1,
+                output_dim=0,
+                weight_loader=weight_loader,
+            ),
+        )
+
+        if self.quant_config.is_checkpoint_fp8_serialized:
+            for scale_name in ["weight_scale", "input_scale"]:
+                layer.register_parameter(
+                    scale_name,
+                    PerTensorScaleParameter(
+                        data=torch.full(
+                            (len(output_partition_sizes),),
+                            torch.finfo(torch.float32).min,
+                            dtype=torch.float32,
+                        ),
+                        weight_loader=weight_loader,
+                    ),
+                )
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        max_w_scale, quantized_weight = requantize_with_max_scale(
+            layer.weight, layer.weight_scale, layer.logical_widths
+        )
+        # Preserve the parameter subclass metadata while rebinding to the
+        # transposed FP8 view expected by the runtime.
+        layer.weight.data = quantized_weight.t().detach()
+        layer.weight.requires_grad_(False)
+        if self.cutlass_fp8_supported:
+            max_w_scale = convert_to_channelwise(max_w_scale, layer.logical_widths)
+        copy_or_rebind_param(layer, "weight_scale", max_w_scale)
+        copy_or_rebind_param(layer, "input_scale", layer.input_scale.max())
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return apply_fp8_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            input_scale=layer.input_scale,
+            bias=bias,
+            cutlass_fp8_supported=self.cutlass_fp8_supported,
+        )
+
+
+class ModelOptFp4LinearMethod(LinearMethodBase):
+    """NVFP4 linear method using the selected FP4 GEMM backend."""
+
+    def __init__(self, quant_config: ModelOptFp4Config):
+        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,
+    ):
+        del input_size, output_size
+        if not self.quant_config.is_checkpoint_nvfp4_serialized:
+            raise ValueError(
+                "NVFP4 quantization was selected, "
+                " dynamic quantization is not supported."
+            )
+        if input_size_per_partition % 16 != 0:
+            raise ValueError(
+                f"Unsupported model when input features size is {input_size_per_partition}, not multiple of 16, for NVFP4 quantization."
+            )
+
+        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
+
+        weight_dtype = (
+            torch.float8_e4m3fn
+            if self.quant_config.is_checkpoint_nvfp4_serialized
+            else params_dtype
+        )
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // 2,
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        input_scale = PerTensorScaleParameter(
+            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(input_scale, {"missing_param_init": "ones"})
+        layer.register_parameter("input_scale", input_scale)
+
+        weight_scale_2 = PerTensorScaleParameter(
+            data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(weight_scale_2, {"missing_param_init": "ones"})
+        layer.register_parameter("weight_scale_2", weight_scale_2)
+
+        weight_scale = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // self.quant_config.group_size,
+                dtype=weight_dtype,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        set_weight_attrs(weight_scale, {"missing_param_init": "ones"})
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        input_scale_2 = layer.input_scale.max().to(torch.float32)
+        weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
+
+        copy_or_rebind_param(
+            layer, "alpha", (input_scale_2 * weight_scale_2).to(torch.float32)
+        )
+        copy_or_rebind_param(
+            layer, "input_scale_inv", (1 / input_scale_2).to(torch.float32)
+        )
+
+        layer.output_size_per_partition = layer.weight.shape[0]
+
+        w = layer.weight.data
+        w_swapped = _prepare_nvfp4_weight_bytes(
+            w,
+            swap_weight_nibbles=getattr(
+                self.quant_config, "swap_weight_nibbles", False
+            ),
+        )
+        scales = layer.weight_scale
+        if (
+            getattr(self.quant_config, "checkpoint_weight_scale_layout", "linear")
+            == "swizzled"
+        ):
+            scales = _swizzled_nvfp4_scales_to_linear(scales)
+
+        _, flashinfer_backend = _get_fp4_gemm_op()
+        if flashinfer_backend == "trtllm":
+            flashinfer_ops = _require_flashinfer()
+
+            weight, _ = pad_nvfp4_weight(w_swapped, n_alignment=128, k_alignment=0)
+            if scales.shape[0] != weight.shape[0]:
+                pad_n = weight.shape[0] - scales.shape[0]
+                scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_n))
+
+            scale_k = scales.shape[1]
+            weights_padding_cols = 0
+            if scale_k % 4 != 0:
+                padded_scale_k = round_up(scale_k, 4)
+                pad_scale_k = padded_scale_k - scale_k
+                scales = torch.nn.functional.pad(scales, (0, pad_scale_k, 0, 0))
+                pad_weight_k = pad_scale_k * 8
+                weight = torch.nn.functional.pad(weight, (0, pad_weight_k, 0, 0))
+                weights_padding_cols = pad_weight_k
+
+            epilogue_tile_m = 128
+            shuffled_scale_shape = scales.shape
+            if not weight.is_cuda:
+                weight = weight.cuda()
+            if scales.device != weight.device:
+                scales = scales.to(device=weight.device)
+            weight = flashinfer_ops.shuffle_matrix_a(
+                weight.view(torch.uint8), epilogue_tile_m
+            )
+            scales = (
+                flashinfer_ops.shuffle_matrix_sf_a(
+                    scales.view(torch.uint8), epilogue_tile_m
+                )
+                .reshape(shuffled_scale_shape)
+                .view(torch.float8_e4m3fn)
+            )
+
+            layer.weights_padding_cols = weights_padding_cols
+            copy_or_rebind_param(layer, "weight", weight)
+            copy_or_rebind_param(layer, "weight_scale_interleaved", scales)
+            return
+        weight, weights_padding_cols = pad_nvfp4_weight(w_swapped)
+        layer.weights_padding_cols = weights_padding_cols
+        copy_or_rebind_param(layer, "weight", weight)
+
+        scale_ndim = scales.ndim
+        if scale_ndim == 2:
+            scales = scales.unsqueeze(0)
+        assert scales.ndim == 3
+        B, M, K = scales.shape
+        M_padded = round_up(M, 128)
+        K_padded = round_up(K, 4)
+        padded_scales = torch.zeros((B, M_padded, K_padded), dtype=scales.dtype)
+        padded_scales[:B, :M, :K] = scales
+
+        _, flashinfer_backend = _get_fp4_gemm_op()
+        uses_flux1_scale_layout = not getattr(
+            self.quant_config, "checkpoint_uses_packed_qkv", False
+        ) and getattr(layer, "prefix", "").startswith(
+            ("transformer_blocks.", "single_transformer_blocks.")
+        )
+        if flashinfer_backend is None or uses_flux1_scale_layout:
+            # CUTLASS and FLUX.1 CUDNN paths need the TMA scale layout.
+            padded_scales = padded_scales.reshape(
+                B, M_padded // 128, 4, 32, K_padded // 4, 4
+            )
+            padded_scales = padded_scales.permute(0, 1, 4, 3, 2, 5)
+
+        padded_scales = padded_scales.contiguous().cuda()
+        padded_scales = (
+            padded_scales.reshape(M_padded, K_padded)
+            if scale_ndim == 2
+            else padded_scales.reshape(B, M_padded, K_padded)
+        )
+        copy_or_rebind_param(layer, "weight_scale_interleaved", padded_scales)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        output_dtype = x.dtype
+        input_shape = x.shape
+        x = x.view(-1, input_shape[-1])
+
+        output_size = layer.output_size_per_partition
+        output_shape = list(input_shape[:-1]) + [output_size]
+
+        fp4_quantize = _get_fp4_quantize_op()
+        if fp4_quantize is None:
+            raise RuntimeError(
+                "No FP4 quantization kernel available. Install flashinfer or sgl_kernel."
+            )
+
+        x_fp4, x_scale_interleaved = fp4_quantize(x, layer.input_scale_inv)
+        weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
+        x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
+
+        w = layer.weight
+        w_scale_interleaved = layer.weight_scale_interleaved
+
+        if x_scale_interleaved.dtype == torch.uint8:
+            x_scale_interleaved = x_scale_interleaved.view(torch.float8_e4m3fn)
+        if w_scale_interleaved.dtype == torch.uint8:
+            w_scale_interleaved = w_scale_interleaved.view(torch.float8_e4m3fn)
+        fp4_gemm, flashinfer_backend = _get_fp4_gemm_op()
+        if flashinfer_backend is not None:
+            out = fp4_gemm(
+                x_fp4,
+                w.T,
+                x_scale_interleaved,
+                w_scale_interleaved.T,
+                layer.alpha,
+                output_dtype,
+                backend=flashinfer_backend,
+            )
+        elif fp4_gemm is not None:
+            out = fp4_gemm(
+                x_fp4,
+                w,
+                x_scale_interleaved,
+                w_scale_interleaved,
+                layer.alpha,
+                output_dtype,
+            )
+        else:
+            raise RuntimeError(
+                "No FP4 GEMM kernel available. Install flashinfer or sgl_kernel."
+            )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py

@@ -0,0 +1,253 @@
+from __future__ import annotations
+
+import logging
+from types import MappingProxyType
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, cast
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+    QuantizeMethodBase,
+)
+from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
+from sglang.srt.layers.quantization.modelslim.schemes import (
+    ModelSlimW4A4Int4,
+    ModelSlimW8A8Int8,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.modelslim.schemes import (
+        ModelSlimLinearScheme,
+    )
+
+from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
+
+logger = logging.getLogger(__name__)
+
+
+class ModelSlimConfig(QuantizationConfig):
+    """
+    Config class for ModelSlim Quantization of Diffusion models https://gitcode.com/Ascend/msmodelslim, a NPU-specific quantization type.
+    The quantization method (W8A8, W4A4, etc.) will be automatically parsed from the `quant_model_description.json` config.
+
+    ModelSlim for Diffusion models includes support for various quantization schemes, such as:
+    - W4A4 dynamic linear
+    - W8A8 static linear
+    - W8A8 dynamic linear
+    """
+
+    def __init__(
+        self,
+        quant_config: Dict[str, Any] = {},
+        reverse_param_names_mapping: dict = None,
+    ):
+        super().__init__()
+        self.quant_description = quant_config
+        ignore = cast(List[str], quant_config.get("ignore", []))
+        self.ignore = ignore
+        packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
+        self.packed_modules_mapping = (
+            packed_modules_mapping if packed_modules_mapping is not None else {}
+        )
+        self._name_mapper = (
+            get_param_names_mapping(reverse_param_names_mapping)
+            if reverse_param_names_mapping is not None
+            else None
+        )
+
+    def get_linear_method(self) -> ModelSlimLinearMethod:
+        return ModelSlimLinearMethod(self)
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.int8, torch.float16, torch.bfloat16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 0
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "modelslim"
+
+    @classmethod
+    def get_config_filenames(cls) -> List[str]:
+        filenames = ["quant_model_description.json"]
+        return filenames
+
+    @classmethod
+    def from_config(
+        cls, config: Dict[str, Any], reverse_param_names_mapping: dict = None
+    ) -> ModelSlimConfig:
+        return cls(config, reverse_param_names_mapping)
+
+    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 should_ignore_layer(
+                prefix,
+                ignore=self.ignore,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                return UnquantizedLinearMethod()
+            key = "model"
+            packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
+            prefix_in_quant_config = prefix
+            proj_name = prefix.split(".")[-1]
+            if proj_name in packed_modules_mapping_subset:
+                prefix_in_quant_config = prefix.replace(
+                    proj_name, packed_modules_mapping_subset[proj_name][0]
+                )
+
+            if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
+                return UnquantizedLinearMethod()
+            scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config)
+            layer.scheme = scheme
+            return ModelSlimLinearMethod(self)
+        else:
+            return None
+
+    def _get_scheme_from_parts(
+        self,
+        layer_name: str,
+    ) -> ModelSlimLinearScheme:
+        full_weight_name = layer_name + ".weight"
+        if self._name_mapper is not None:
+            mapped_name, _, _ = self._name_mapper(full_weight_name)
+        else:
+            mapped_name = full_weight_name
+
+        quant_type = self.quant_description.get(mapped_name, "")
+        prefix = mapped_name.removesuffix(".weight")
+        if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8":
+            return ModelSlimW8A8Int8(quant_config=self.quant_description, prefix=prefix)
+        elif quant_type == "W4A4_DYNAMIC":
+            return ModelSlimW4A4Int4(quant_config=self.quant_description, prefix=prefix)
+        elif quant_type == "W8A8_MXFP8":
+            from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp8_scheme import (
+                ModelSlimMXFP8Scheme,
+            )
+
+            return ModelSlimMXFP8Scheme()
+        elif quant_type in ("W4A4_MXFP4", "W4A4_MXFP4_DUALSCALE"):
+            from sglang.multimodal_gen.runtime.layers.quantization.modelslim_mxfp4_scheme import (
+                ModelSlimMXFP4Scheme,
+            )
+
+            return ModelSlimMXFP4Scheme()
+        raise NotImplementedError(
+            f"No modelslim compatible scheme was found for layer '{layer_name}'. "
+            f"quant_description['{layer_name}.weight'] = '{quant_type}'"
+        )
+
+    def get_scheme(
+        self, layer: torch.nn.Module, layer_name: Optional[str] = None
+    ) -> Optional[ModelSlimLinearScheme]:
+        """
+        get_scheme method adjusted for modelslim, taken from
+        python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+        """
+        scheme = self._get_scheme_from_parts(
+            layer_name=layer_name,
+        )
+
+        # Ascend doesn't support device capability
+        logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)
+        return scheme
+
+    def is_layer_skipped(
+        self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
+    ):
+        # adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
+        proj_name = prefix.split(".")[-1]
+        if proj_name in fused_mapping:
+            shard_prefixes = [
+                prefix.replace(proj_name, shard_proj_name)
+                for shard_proj_name in fused_mapping[proj_name]
+            ]
+
+            is_skipped = None
+            for shard_prefix in shard_prefixes:
+                is_shard_skipped = (
+                    self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT"
+                )
+
+                if is_skipped is None:
+                    is_skipped = is_shard_skipped
+                elif is_shard_skipped != is_skipped:
+                    raise ValueError(
+                        f"Detected some but not all shards of {prefix} "
+                        "are quantized. All shards of fused layers "
+                        "to have the same precision."
+                    )
+        else:
+            is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT"
+
+        assert is_skipped is not None
+        return is_skipped
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class ModelSlimLinearMethod(LinearMethodBase):
+
+    def __init__(self, quantization_config: ModelSlimConfig):
+        self.quantization_config = quantization_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.scheme.process_weights_after_loading(layer)
+
+    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,
+    ):
+        """
+        Use the ModelSlimLinearScheme associated with each layer to create
+        the necessary parameters for the layer. See LinearMethodBase for param
+        details
+        """
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        layer.scheme.create_weights(
+            layer=layer,
+            input_size=input_size,
+            input_size_per_partition=input_size_per_partition,
+            output_partition_sizes=output_partition_sizes,
+            output_size=output_size,
+            params_dtype=params_dtype,
+            weight_loader=weight_loader,
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ):
+        """
+        Use the output of create_weights and the CompressedTensorsScheme
+        associated with the layer to apply the forward pass with the
+        layer input.  See LinearMethodBase for param details
+
+        """
+
+        scheme = layer.scheme
+        if scheme is None:
+            raise ValueError("A scheme must be defined for each layer")
+        return scheme.apply_weights(layer, x, bias=bias)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp4_scheme.py

@@ -0,0 +1,197 @@
+"""ModelSlim MXFP4 scheme for pre-quantized weight inference on Ascend NPU.
+
+Loads weights pre-quantized by msmodelslim and runs MXFP4 dual-level
+matmul at inference via npu_dual_level_quant_matmul.
+
+Checkpoint tensor formats (verified from msmodelslim export):
+  weight:           [out, in]           float8_e4m3fn  (FP4 data in fp8 container)
+  weight_scale:     [out, in/32]        uint8          (L1 block scales, e8m0+127)
+  weight_dual_scale:[out, in/512, 1]    float32        (L0 coarse scales)
+  mul_scale:        [in]                float32        (smooth quant activation scale)
+
+Reference: MindIE-SD W4A4MXFP4DualQuantLinear
+(MindIE-SD/mindiesd/quantization/layer.py)
+"""
+
+from typing import List, Optional
+
+import torch
+
+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.models.parameter import (
+    BasevLLMParameter,
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
+
+MXFP4_BLOCK_SIZE = 32
+# L1 (dual) scale groups this many L0 blocks together.
+# L1 block covers 16 * 32 = 512 elements.
+MXFP4_DUAL_LEVEL_RATIO = 16
+
+
+class ModelSlimMXFP4Scheme(ModelSlimLinearScheme):
+
+    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,
+    ):
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        output_size_per_partition = sum(output_partition_sizes)
+
+        # msmodelslim exports weight as float8_e4m3fn, shape [out, in].
+        # Each byte is a float8 container for FP4 data; the actual FP4 packing
+        # (npu_dtype_cast → float4_e2m1fn_x2) happens in process_weights_after_loading.
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                (output_size_per_partition, input_size_per_partition),
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # L1 block scale: uint8 [out, in/32], e8m0 scale with +127 offset.
+        scale_dim = input_size_per_partition // MXFP4_BLOCK_SIZE
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, scale_dim),
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+        # L0 (coarse) scale for dual-level quantization matmul.
+        # Each L0 block covers MXFP4_DUAL_LEVEL_RATIO L1 blocks = 16 * 32 = 512 elements.
+        dual_scale_dim = scale_dim // MXFP4_DUAL_LEVEL_RATIO  # in/32 / 16 = in/512
+        weight_dual_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, dual_scale_dim, 1),
+                dtype=torch.float32,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_dual_scale", weight_dual_scale)
+
+        # Smooth quant activation scale (mul_scale) from NonFusionSmoothQuantWrapper.
+        # msmodelslim exports this as `<prefix>.div.mul_scale` with shape [in].
+        # After repack, it becomes `<prefix>.mul_scale`.
+        # This is CRITICAL: the offline-quantized weights were calibrated with
+        # x * mul_scale applied to the activation. Omitting it causes mosaic output.
+        # Ref: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul lines 385-386.
+        mul_scale = BasevLLMParameter(
+            data=torch.empty(
+                (input_size_per_partition,),
+                dtype=torch.float32,
+            ),
+            weight_loader=weight_loader,
+        )
+        # If mul_scale is not in the checkpoint (e.g. non-smooth-quant model
+        # or old repack without .div. handling), initialize to ones so that
+        # x * 1.0 = x (no-op). fsdp_load.py checks this attribute.
+        mul_scale.missing_param_init = "ones"
+        layer.register_parameter("mul_scale", mul_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # Cast weight from fp8 container to FP4 packed format
+        weight = layer.weight.data
+        if not weight.is_npu:
+            weight = weight.to(f"npu:{torch.npu.current_device()}")
+        weight = torch_npu.npu_dtype_cast(weight, torch_npu.float4_e2m1fn_x2)
+        # npu_dual_level_quant_matmul requires x2 in FRACTAL_NZ format (format 29).
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
+        weight = torch_npu.npu_format_cast(
+            weight.view(torch.int8), 29, customize_dtype=torch.int8
+        )
+        layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+        # Reshape weight_scale: [out, in/32] -> [out, in/64, 2]
+        # The dual-level matmul API expects L1 scales in this 3D format
+        weight_scale = layer.weight_scale.data
+        if not weight_scale.is_npu:
+            weight_scale = weight_scale.to(f"npu:{torch.npu.current_device()}")
+        weight_scale = weight_scale.reshape(weight_scale.shape[0], -1, 2)
+        layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+
+        # Transform weight_dual_scale: [out, in/512, 1] -> [in/512, out]
+        weight_dual_scale = layer.weight_dual_scale.data
+        if not weight_dual_scale.is_npu:
+            weight_dual_scale = weight_dual_scale.to(
+                f"npu:{torch.npu.current_device()}"
+            )
+        weight_dual_scale = weight_dual_scale.squeeze(-1).transpose(0, 1).contiguous()
+        layer.weight_dual_scale = torch.nn.Parameter(
+            weight_dual_scale, requires_grad=False
+        )
+
+        # Move mul_scale to NPU if present and not already there
+        mul_scale = layer.mul_scale.data
+        if not mul_scale.is_npu:
+            mul_scale = mul_scale.to(f"npu:{torch.npu.current_device()}")
+        layer.mul_scale = torch.nn.Parameter(mul_scale, requires_grad=False)
+        layer.use_mul_scale = not torch.all(mul_scale == 1.0).item()
+
+    def apply_weights(
+        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 for npu_dynamic_dual_level_mx_quant
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Apply smooth quant scale before activation quantization.
+        # The offline-quantized weights were calibrated under x * mul_scale,
+        # so we MUST apply it here for scale alignment.
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear.quant_matmul
+        mul_scale = layer.mul_scale
+        if getattr(layer, "use_mul_scale", True):
+            x_2d = x_2d * mul_scale.to(x_2d.dtype)
+
+        # Dual-level MXFP4 activation quantization
+        x1, l0_scale, l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            x_2d, smooth_scale=None
+        )
+
+        # Dual-level MXFP4 matmul
+        output = torch_npu.npu_dual_level_quant_matmul(
+            x1,
+            layer.weight,
+            l0_scale,
+            layer.weight_dual_scale,
+            l1_scale,
+            layer.weight_scale,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+        )
+
+        # Restore original shape
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        return output.reshape(output_shape)

python/sglang/multimodal_gen/runtime/layers/quantization/modelslim_mxfp8_scheme.py

@@ -0,0 +1,124 @@
+"""ModelSlim MXFP8 scheme for pre-quantized weight inference on Ascend NPU.
+
+Loads weights pre-quantized by msmodelslim (float8_e4m3fn weights,
+uint8 scales) and runs MXFP8 matmul at inference.
+"""
+
+from typing import List, Optional
+
+import torch
+
+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.models.parameter import (
+    GroupQuantScaleParameter,
+    ModelWeightParameter,
+)
+from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme
+
+MXFP8_BLOCK_SIZE = 32
+
+
+class ModelSlimMXFP8Scheme(ModelSlimLinearScheme):
+
+    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,
+    ):
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        output_size_per_partition = sum(output_partition_sizes)
+
+        # msmodelslim exports weight as float8_e4m3fn, shape [out, in]
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                (output_size_per_partition, input_size_per_partition),
+                dtype=torch.float8_e4m3fn,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight", weight)
+
+        # 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
+        # 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
+        weight_scale = GroupQuantScaleParameter(
+            data=torch.empty(
+                (output_size_per_partition, scale_dim),
+                dtype=torch.uint8,
+            ),
+            input_dim=1,
+            output_dim=0,
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # weight is already float8_e4m3fn, no cast needed
+        weight = layer.weight.data
+        layer.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+        # Reshape weight_scale: [out, in/32] -> [out, in/32//2, 2]
+        weight_scale = layer.weight_scale.data
+        weight_scale = weight_scale.reshape(weight_scale.shape[0], -1, 2)
+        layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
+
+    def apply_weights(
+        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):
+            # npu_dynamic_mx_quant only accepts fp16/bf16 activations
+            x = x.to(torch.bfloat16)
+            original_dtype = torch.bfloat16
+
+        # npu_dynamic_mx_quant requires a 2D input [tokens, hidden_size].
+        # Diffusion transformer inputs are typically 3D [batch, seq, hidden] or
+        # higher. Flattening to 2D merges all leading dimensions into a single
+        # token axis so the NPU kernel can compute per-token MXFP8 scales, then
+        # we restore the original shape from the output.
+        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.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
+        output_shape = list(input_shape[:-1]) + [output.shape[-1]]
+        output = output.reshape(output_shape)
+
+        return output

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4.py

@@ -0,0 +1,238 @@
+import logging
+from typing import Dict, List, Optional
+
+import torch
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.layers.linear import (
+    LinearMethodBase,
+    UnquantizedLinearMethod,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.models.parameter import (
+    ModelWeightParameter,
+    PerTensorScaleParameter,
+)
+from sglang.srt.layers.quantization.utils import is_layer_skipped
+from sglang.srt.utils import is_hip, mxfp_supported
+
+logger = logging.getLogger(__name__)
+_is_hip = is_hip()
+
+if _is_hip:
+    try:
+        import aiter
+        from aiter.ops.gemm_op_a4w4 import gemm_a4w4
+        from aiter.ops.shuffle import shuffle_weight
+        from aiter.utility.fp4_utils import dynamic_mxfp4_quant
+    except ImportError as e:
+        logger.warning(f"aiter MXFP4 kernels not available: {e}")
+        aiter = None
+        shuffle_weight = None
+        dynamic_mxfp4_quant = None
+        gemm_a4w4 = None
+
+# The gemm_a4w4 ASM kernel has degraded precision when the output
+# dimension (N) is smaller than its minimum tile size.
+# Layers with output_size falls below this threshold will stay unquantized
+_MXFP4_MIN_OUTPUT_DIM = 256
+
+
+class Mxfp4Config(QuantizationConfig):
+    """
+    MXFP4 quantization config for diffusion models.
+
+    Supports online quantization from unquantized BF16/FP16 checkpoints;
+    no-arg ``Mxfp4Config()`` selects that online (post-load) path.
+    Note: MXFP4 requires ROCm and MI350+ (gfx95x).
+    """
+
+    def __init__(
+        self,
+        is_checkpoint_mxfp4_serialized: bool = False,
+        ignored_layers: Optional[List[str]] = None,
+        packed_modules_mapping: Optional[Dict[str, List[str]]] = None,
+    ):
+        super().__init__()
+        self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized
+        self.ignored_layers = ignored_layers or []
+        self.packed_modules_mapping = packed_modules_mapping or {}
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "mxfp4"
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+        return [torch.bfloat16, torch.float16]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 95  # gfx95x, Note: mxfp_supported() is a better check
+
+    @classmethod
+    def get_config_filenames(cls) -> list[str]:
+        return []  # No config file needed for online quantization
+
+    @classmethod
+    def from_config(cls, config: dict) -> "Mxfp4Config":
+        """Create from model config (for pre-quantized checkpoints)."""
+        is_serialized = config.get("quant_method") == "mxfp4"
+        return cls(is_checkpoint_mxfp4_serialized=is_serialized)
+
+    def get_quant_method(self, layer, prefix: str):
+        from sglang.multimodal_gen.runtime.layers.linear import LinearBase
+
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(
+                prefix,
+                self.ignored_layers,
+                fused_mapping=self.packed_modules_mapping,
+            ):
+                logger.debug(
+                    f"MXFP4: Keeping layer {prefix} unquantized (in ignored_layers)"
+                )
+                return UnquantizedLinearMethod()
+            # Skip layers whose output dims are too small, see ASM kernel comment above
+            output_size = getattr(layer, "output_size", None)
+            if output_size is not None and output_size < _MXFP4_MIN_OUTPUT_DIM:
+                logger.info(
+                    f"MXFP4: Keeping layer {prefix} unquantized "
+                    f"(output_size={output_size} < {_MXFP4_MIN_OUTPUT_DIM})"
+                )
+                return UnquantizedLinearMethod()
+            logger.debug(f"MXFP4: Replacing layer {prefix} with MXFP4 linear method")
+            return Mxfp4LinearMethod(self)
+        else:
+            logger.debug(f"MXFP4: Skipping layer {prefix} (not a LinearBase)")
+        return None
+
+
+class Mxfp4LinearMethod(LinearMethodBase):
+    """
+    MXFP4 online quantization method for linear layers.
+
+    Quantizes unquantized BF16/FP16 weights to MXFP4 format during
+    process_weights_after_loading().
+    """
+
+    def __init__(self, quant_config: Mxfp4Config):
+        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,
+    ):
+        """
+        Creates BF16/FP16 parameters that will be
+        quantized to MXFP4 in process_weights_after_loading().
+        """
+        output_size_per_partition = sum(output_partition_sizes)
+        weight_loader = extra_weight_attrs.get("weight_loader")
+
+        weight = ModelWeightParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition,
+                dtype=params_dtype,
+            ),
+            weight_loader=weight_loader,
+            input_dim=1,
+            output_dim=0,
+        )
+        layer.register_parameter("weight", weight)
+
+        # Placeholder scale (will be created during quantization)
+        weight_scale = PerTensorScaleParameter(
+            data=torch.empty(1, dtype=torch.float32),
+            weight_loader=weight_loader,
+        )
+        layer.register_parameter("weight_scale", weight_scale)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        """
+        Quantize BF16/FP16 weights to MXFP4 after loading from checkpoint.
+
+        Converts weights from unquantized format to:
+        - Packed uint8 (2 FP4 values per byte)
+        - E8M0 scales (one per 32-element block)
+        """
+        if not mxfp_supported():
+            platform = "unknown"
+            if _is_hip:
+                try:
+                    platform = torch.cuda.get_device_properties(0).gcnArchName
+                except:
+                    platform = "ROCm (unknown arch)"
+            raise RuntimeError(
+                f"MXFP4 quantization requires ROCm and MI350+ (gfx95x). "
+                f"Current platform: {platform}."
+            )
+
+        # Check if weights are already quantized
+        if layer.weight.dtype not in [torch.bfloat16, torch.float16]:
+            # Already quantized or unexpected dtype
+            logger.info("Weights are quantized or unexpected dtype")
+            return
+
+        if any(fn is None for fn in (dynamic_mxfp4_quant, shuffle_weight, gemm_a4w4)):
+            raise RuntimeError(
+                "aiter MXFP4 kernels not available. "
+                "Install aiter with MXFP4 support."
+            )
+
+        weight_data = layer.weight.data
+        was_on_cpu = weight_data.device.type == "cpu"
+        if was_on_cpu:
+            weight_data = weight_data.cuda()
+
+        w_quant, mx_scales = dynamic_mxfp4_quant(weight_data, shuffle=True)
+
+        w_quant_shuffled = shuffle_weight(w_quant)
+
+        if was_on_cpu:
+            w_quant_shuffled = w_quant_shuffled.cpu()
+            mx_scales = mx_scales.cpu()
+
+        layer.weight = Parameter(w_quant_shuffled, requires_grad=False)
+        layer.weight_scale = Parameter(mx_scales, requires_grad=False)
+
+        logger.debug(
+            f"MXFP4: Quantized layer weights - weight {layer.weight.shape} {layer.weight.dtype}, "
+            f"scale {layer.weight_scale.shape}"
+        )
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if not mxfp_supported():
+            raise RuntimeError(
+                "MXFP4 inference requires ROCm and MI350+ (gfx95x). "
+                "Current platform not supported."
+            )
+
+        # Handle 3D input tensors [batch, seq, hidden]
+        original_shape = x.shape
+        if x.dim() == 3:
+            x = x.view(-1, x.shape[-1])
+
+        x_fp4, x_scale = dynamic_mxfp4_quant(x, shuffle=True)
+
+        y = gemm_a4w4(x_fp4, layer.weight, x_scale, layer.weight_scale)
+
+        if bias is not None:
+            y = y + bias
+
+        return y.view(*original_shape[:-1], layer.weight.shape[0])

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp4_npu.py

@@ -0,0 +1,201 @@
+"""Online MXFP4 quantization for Diffusion models on Ascend NPU.
+
+Provides ``NPUMXFP4Config`` (registered as ``"mxfp4_npu"``) and
+``NPUMXFP4DiffusionLinearMethod`` which quantises FP16/BF16 weights to MXFP4
+at load time using dual-level MX quantization and uses
+``npu_dynamic_dual_level_mx_quant`` + ``npu_dual_level_quant_matmul`` for
+inference.
+
+The ``"mxfp4_npu"`` key is distinct from upstream's ROCm ``"mxfp4"``
+(``Mxfp4Config`` in ``mxfp4.py``) which targets AMD MI350+ via aiter kernels.
+
+NOTE: Online weight quantization via ``npu_dynamic_dual_level_mx_quant`` is
+experimental. MindIE-SD only uses an offline (pre-quantized) path for MXFP4
+weights. The online path quantizes FP16/BF16 weights at load time, which may
+produce different numerical results than the offline calibrated path.
+"""
+
+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__)
+
+
+class NPUMXFP4Config(QuantizationConfig):
+    """Config for online MXFP4 quantization on Ascend NPU (Diffusion)."""
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "mxfp4_npu"
+
+    @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]) -> NPUMXFP4Config:
+        return cls()
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+        if isinstance(layer, LinearBase):
+            return NPUMXFP4DiffusionLinearMethod(self)
+        return None
+
+    def get_scaled_act_names(self) -> List[str]:
+        return []
+
+
+class NPUMXFP4DiffusionLinearMethod(LinearMethodBase):
+    """Ascend NPU MXFP4 linear method for Diffusion models (dual-level).
+
+    Online mode: loads FP16/BF16 weights → quantises to MXFP4 at load time
+    via ``npu_dynamic_dual_level_mx_quant``.
+    Inference: dynamic dual-level MXFP4 activation quant + dual-level matmul.
+
+    Reference: MindIE-SD ``W4A4MXFP4DualQuantLinear`` (offline path only).
+    """
+
+    def __init__(self, quant_config: NPUMXFP4Config):
+        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. dit_cpu_offload defaults to True in
+        # ServerArgs, which causes fsdp_load to move parameters back to CPU
+        # after loading. npu_dynamic_dual_level_mx_quant requires an NPU tensor.
+        if not weight_fp.is_npu:
+            weight_fp = weight_fp.to(f"npu:{torch.npu.current_device()}")
+
+        # Online dual-level MXFP4 weight quantisation.
+        # NOTE: This is experimental — MindIE-SD only has an offline path for
+        # MXFP4 weights. We assume npu_dynamic_dual_level_mx_quant can also
+        # quantise weights (not just activations).
+        # Returns: (qw, w_dual_scale, w_scale)
+        #   qw          — quantized weight in float4_e2m1fn_x2 (2 FP4 packed/byte)
+        #   w_dual_scale — L0-level scale (goes to pos 3 in npu_dual_level_quant_matmul)
+        #   w_scale      — L1-level scale (goes to pos 5 in npu_dual_level_quant_matmul)
+        qw, w_dual_scale, w_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            weight_fp, smooth_scale=None
+        )
+
+        # npu_dual_level_quant_matmul requires x2 (weight) in FRACTAL_NZ format.
+        # Reference: MindIE-SD W4A4MXFP4DualQuantLinear._init_dynamic_quant_param
+        qw = torch_npu.npu_format_cast(
+            qw.view(torch.int8), 29, customize_dtype=torch.int8
+        )
+
+        # x2Level0Scale must be [in/level0_block_size, out] — transpose from
+        # the [out, in/level0_block_size] shape returned by the quant op.
+        # Reference: MindIE-SD layer.py:409
+        w_dual_scale = w_dual_scale.squeeze(-1).transpose(0, 1).contiguous()
+
+        layer.weight = Parameter(qw, requires_grad=False)
+        layer.weight_dual_scale = Parameter(w_dual_scale, requires_grad=False)
+        layer.weight_scale = 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] for the quantization operators
+        input_shape = x.shape
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        # Dynamic dual-level MXFP4 activation quantisation
+        qx, act_l0_scale, act_l1_scale = torch_npu.npu_dynamic_dual_level_mx_quant(
+            x_2d, smooth_scale=None
+        )
+
+        # Dual-level MXFP4 matmul
+        # Arg order: act_quant, weight_quant, act_l0_scale, weight_dual_scale,
+        #            act_l1_scale, weight_scale, bias=, output_dtype=
+        # NOTE: weight is NOT transposed (unlike MXFP8's npu_quant_matmul).
+        output = torch_npu.npu_dual_level_quant_matmul(
+            qx,
+            layer.weight,
+            act_l0_scale,
+            layer.weight_dual_scale,
+            act_l1_scale,
+            layer.weight_scale,
+            bias=bias.to(torch.float32) if bias is not None else None,
+            output_dtype=original_dtype,
+        )
+
+        # 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

python/sglang/multimodal_gen/runtime/layers/quantization/mxfp8_npu.py

@@ -0,0 +1,176 @@
+"""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

python/sglang/multimodal_gen/runtime/layers/quantization/nunchaku_linear.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import List, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.parameter import Parameter
+
+from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+logger = init_logger(__name__)
+
+try:
+    from nunchaku.ops.gemm import svdq_gemm_w4a4_cuda
+    from nunchaku.ops.gemv import awq_gemv_w4a16_cuda
+    from nunchaku.ops.quantize import svdq_quantize_w4a4_act_fuse_lora_cuda
+except ImportError:
+    svdq_gemm_w4a4_cuda = None
+    awq_gemv_w4a16_cuda = None
+    svdq_quantize_w4a4_act_fuse_lora_cuda = None
+
+
+class NunchakuSVDQLinearMethod(LinearMethodBase):
+    def __init__(
+        self,
+        precision: str = "int4",
+        rank: int = 32,
+        act_unsigned: bool = False,
+    ):
+        self.precision = precision
+        self.rank = rank
+        self.act_unsigned = act_unsigned
+
+        if precision == "nvfp4":
+            self.group_size = 16
+        else:
+            self.group_size = 64
+
+    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,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+
+        qweight = Parameter(
+            torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // 2,
+                dtype=torch.int8,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0})
+
+        num_groups = input_size_per_partition // self.group_size
+        if self.precision == "nvfp4":
+            scale_dtype = torch.float8_e4m3fn
+        else:
+            scale_dtype = params_dtype
+        wscales = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=scale_dtype),
+            requires_grad=False,
+        )
+
+        smooth_factor = Parameter(
+            torch.empty(input_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        smooth_factor_orig = Parameter(
+            torch.empty(input_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        proj_down = Parameter(
+            torch.empty(input_size_per_partition, self.rank, dtype=params_dtype),
+            requires_grad=False,
+        )
+        proj_up = Parameter(
+            torch.empty(output_size_per_partition, self.rank, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        if self.precision == "nvfp4":
+            wcscales = Parameter(
+                torch.empty(
+                    output_size_per_partition,
+                    dtype=params_dtype,
+                ),
+                requires_grad=False,
+            )
+            wtscale = Parameter(
+                torch.empty(1, dtype=params_dtype),
+                requires_grad=False,
+            )
+        else:
+            wcscales = None
+            wtscale = None
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("wscales", wscales)
+        layer.register_parameter("smooth_factor", smooth_factor)
+        layer.register_parameter("smooth_factor_orig", smooth_factor_orig)
+        layer.register_parameter("proj_down", proj_down)
+        layer.register_parameter("proj_up", proj_up)
+        if wcscales is not None:
+            layer.register_parameter("wcscales", wcscales)
+        if wtscale is not None:
+            layer.register_parameter("wtscale", wtscale)
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.precision = self.precision
+        layer.rank = self.rank
+        layer.group_size = self.group_size
+        layer.act_unsigned = self.act_unsigned
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if weight_loader is not None:
+            set_weight_attrs(qweight, {"weight_loader": weight_loader})
+            set_weight_attrs(wscales, {"weight_loader": weight_loader})
+            set_weight_attrs(smooth_factor, {"weight_loader": weight_loader})
+            set_weight_attrs(smooth_factor_orig, {"weight_loader": weight_loader})
+            set_weight_attrs(proj_down, {"weight_loader": weight_loader})
+            set_weight_attrs(proj_up, {"weight_loader": weight_loader})
+            if wcscales is not None:
+                set_weight_attrs(wcscales, {"weight_loader": weight_loader})
+            if wtscale is not None:
+                set_weight_attrs(wtscale, {"weight_loader": weight_loader})
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.wscales = Parameter(layer.wscales.data, requires_grad=False)
+        layer.smooth_factor = Parameter(layer.smooth_factor.data, requires_grad=False)
+        layer.smooth_factor_orig = Parameter(
+            layer.smooth_factor_orig.data, requires_grad=False
+        )
+        layer.proj_down = Parameter(layer.proj_down.data, requires_grad=False)
+        layer.proj_up = Parameter(layer.proj_up.data, requires_grad=False)
+        if hasattr(layer, "wcscales") and layer.wcscales is not None:
+            layer.wcscales = Parameter(layer.wcscales.data, requires_grad=False)
+        if hasattr(layer, "wtscale") and layer.wtscale is not None:
+            layer.wtscale = Parameter(layer.wtscale.data, requires_grad=False)
+
+        alpha: float | None = None
+        wtscale = getattr(layer, "wtscale", None)
+        if wtscale is not None:
+            if isinstance(wtscale, Parameter):
+                wtscale = wtscale.data
+            if isinstance(wtscale, torch.Tensor):
+                alpha = float(wtscale.detach().cpu().item())
+            else:
+                alpha = float(wtscale)
+        layer._nunchaku_alpha = alpha
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, orig_shape[-1])
+        quantized_x, ascales, lora_act_out = svdq_quantize_w4a4_act_fuse_lora_cuda(
+            x_2d,
+            lora_down=layer.proj_down,
+            smooth=layer.smooth_factor,
+            fp4=layer.precision == "nvfp4",
+            pad_size=256,
+        )
+        out_2d = torch.empty(
+            x_2d.shape[0],
+            layer.output_size_per_partition,
+            dtype=x_2d.dtype,
+            device=x_2d.device,
+        )
+        alpha: float | None = getattr(layer, "_nunchaku_alpha", None)
+        wcscales = getattr(layer, "wcscales", None)
+
+        svdq_gemm_w4a4_cuda(
+            act=quantized_x,
+            wgt=layer.qweight,
+            out=out_2d,
+            ascales=ascales,
+            wscales=layer.wscales,
+            lora_act_in=lora_act_out,
+            lora_up=layer.proj_up,
+            bias=bias,
+            fp4=layer.precision == "nvfp4",
+            alpha=alpha,
+            wcscales=wcscales,
+            act_unsigned=getattr(layer, "act_unsigned", False),
+        )
+        out = out_2d.reshape(*orig_shape[:-1], layer.output_size_per_partition)
+        return out
+
+
+class NunchakuAWQLinearMethod(LinearMethodBase):
+    def __init__(self, group_size: int = 64):
+        self.group_size = group_size
+        self.pack_factor = 8
+
+    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,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+
+        qweight = Parameter(
+            torch.empty(
+                output_size_per_partition // 4,
+                input_size_per_partition // 2,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0})
+
+        num_groups = input_size_per_partition // self.group_size
+        wscales = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        wzeros = Parameter(
+            torch.empty(num_groups, output_size_per_partition, dtype=params_dtype),
+            requires_grad=False,
+        )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("wscales", wscales)
+        layer.register_parameter("wzeros", wzeros)
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.group_size = self.group_size
+        layer.pack_factor = self.pack_factor
+
+        weight_loader = extra_weight_attrs.get("weight_loader")
+        if weight_loader is not None:
+            set_weight_attrs(qweight, {"weight_loader": weight_loader})
+            set_weight_attrs(wscales, {"weight_loader": weight_loader})
+            set_weight_attrs(wzeros, {"weight_loader": weight_loader})
+
+    def process_weights_after_loading(self, layer: nn.Module) -> None:
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.wscales = Parameter(layer.wscales.data, requires_grad=False)
+        layer.wzeros = Parameter(layer.wzeros.data, requires_grad=False)
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        orig_shape = x.shape
+        x_2d = x.reshape(-1, orig_shape[-1])
+
+        in_features = layer.input_size_per_partition
+        out_features = layer.output_size_per_partition
+        out_2d = awq_gemv_w4a16_cuda(
+            in_feats=x_2d,
+            kernel=layer.qweight,
+            scaling_factors=layer.wscales,
+            zeros=layer.wzeros,
+            m=x_2d.shape[0],
+            n=out_features,
+            k=in_features,
+            group_size=layer.group_size,
+        )
+        if bias is not None:
+            view_shape = [1] * (out_2d.ndim - 1) + [-1]
+            out_2d.add_(bias.view(view_shape))
+
+        out = out_2d.reshape(*orig_shape[:-1], out_features)
+        return out

python/sglang/multimodal_gen/runtime/layers/quantization/weight_only_fp8.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import sglang.multimodal_gen.envs as envs
+from sglang.multimodal_gen.runtime.distributed import (
+    divide,
+    get_tp_group,
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_gather,
+    tensor_model_parallel_all_reduce,
+)
+from sglang.multimodal_gen.runtime.layers.utils import get_group_rank, get_group_size
+from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
+
+FP8_WEIGHT_DTYPE = torch.float8_e4m3fn
+W8A8_FP8_GEMM_ENV = "SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM"
+
+logger = logging.getLogger(__name__)
+_w8a8_fp8_gemm_warning_logged = False
+
+
+def _can_apply_fused_w8a8_fp8_linear(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    compute_dtype: torch.dtype,
+) -> bool:
+    return (
+        x.device.type == "cuda"
+        and weight.device.type == "cuda"
+        and weight_scale.device.type == "cuda"
+        and not x.is_meta
+        and not weight.is_meta
+        and not weight_scale.is_meta
+        and compute_dtype in (torch.float16, torch.bfloat16)
+    )
+
+
+def dequantize_rowwise_fp8_weight(
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    if weight.ndim != 2:
+        raise ValueError(f"FP8 linear weight must be 2-D, got shape {weight.shape}")
+    if weight_scale.ndim != 1 or weight_scale.shape[0] != weight.shape[0]:
+        raise ValueError(
+            "FP8 row-wise scale must have shape (out_features,), "
+            f"got weight={tuple(weight.shape)} scale={tuple(weight_scale.shape)}"
+        )
+    return weight.to(dtype) * weight_scale.to(dtype).unsqueeze(1)
+
+
+def _apply_srt_w8a8_fp8_linear(*args, **kwargs) -> torch.Tensor:
+    from sglang.srt.layers.quantization.fp8_utils import apply_fp8_linear
+
+    return apply_fp8_linear(*args, **kwargs)
+
+
+def _is_cutlass_fp8_supported() -> bool:
+    from sglang.srt.layers.quantization.fp8_utils import cutlass_fp8_supported
+
+    return cutlass_fp8_supported()
+
+
+def _apply_weight_only_fp8_linear(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    bias: torch.Tensor | None,
+    compute_dtype: torch.dtype,
+    enable_fused_w8a8: bool,
+) -> torch.Tensor:
+    x = x.to(compute_dtype)
+    bias = bias.to(compute_dtype) if bias is not None else None
+    if enable_fused_w8a8 and _can_apply_fused_w8a8_fp8_linear(
+        x, weight, weight_scale, compute_dtype
+    ):
+        try:
+            # The fused kernel uses W8A8 compute; fallback keeps BF16/FP16
+            # activations after dequantizing the FP8 weights.
+            output = _apply_srt_w8a8_fp8_linear(
+                input=x,
+                weight=weight.t(),
+                weight_scale=weight_scale,
+                input_scale=None,
+                bias=bias,
+                cutlass_fp8_supported=_is_cutlass_fp8_supported(),
+            )
+            _log_w8a8_fp8_gemm_warning_once()
+            return output
+        except (ImportError, NotImplementedError):
+            pass
+
+    dequant_weight = dequantize_rowwise_fp8_weight(weight, weight_scale, compute_dtype)
+    return F.linear(x, dequant_weight, bias)
+
+
+class WeightOnlyFP8Linear(nn.Module):
+    """Storage-only e4m3 FP8 linear with row-wise weight scales."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.weight = nn.Parameter(
+            torch.empty(out_features, in_features, dtype=FP8_WEIGHT_DTYPE),
+            requires_grad=False,
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(out_features, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(self.weight_scale, {"missing_param_init": "error"})
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        compute_dtype = self.compute_dtype or x.dtype
+        return _apply_weight_only_fp8_linear(
+            x,
+            self.weight,
+            self.weight_scale,
+            self.bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+
+
+class WeightOnlyFP8ColumnParallelLinear(nn.Module):
+    """Column-parallel storage-only e4m3 FP8 linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        gather_output: bool = True,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.gather_output = gather_output
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.tp_group = tp_group or get_tp_group()
+        self.tp_size = get_group_size(self.tp_group)
+        self.tp_rank = get_group_rank(self.tp_group)
+        self.out_features_per_partition = divide(out_features, self.tp_size)
+        self.weight = nn.Parameter(
+            torch.empty(
+                self.out_features_per_partition,
+                in_features,
+                dtype=FP8_WEIGHT_DTYPE,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(self.out_features_per_partition, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight_scale,
+            {
+                "missing_param_init": "error",
+                "output_dim": 0,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    self.out_features_per_partition,
+                    dtype=compute_dtype or torch.get_default_dtype(),
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(
+                self.bias,
+                {
+                    "output_dim": 0,
+                    "weight_loader": self.weight_loader,
+                },
+            )
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        output_dim = getattr(param, "output_dim", None)
+        if output_dim is not None:
+            shard_size = param.data.shape[output_dim]
+            loaded_weight = loaded_weight.narrow(
+                output_dim, self.tp_rank * shard_size, shard_size
+            )
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        compute_dtype = self.compute_dtype or x.dtype
+        output_parallel = _apply_weight_only_fp8_linear(
+            x,
+            self.weight,
+            self.weight_scale,
+            self.bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+        if self.gather_output:
+            return tensor_model_parallel_all_gather(
+                output_parallel, tp_group=self.tp_group
+            )
+        return output_parallel
+
+
+class WeightOnlyFP8MergedColumnParallelLinear(WeightOnlyFP8ColumnParallelLinear):
+    """Column-parallel storage-only FP8 packed linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        output_sizes: list[int],
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        gather_output: bool = False,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        self.output_sizes = output_sizes
+        super().__init__(
+            in_features,
+            sum(output_sizes),
+            bias=bias,
+            compute_dtype=compute_dtype,
+            gather_output=gather_output,
+            tp_group=tp_group,
+            enable_fused_w8a8=enable_fused_w8a8,
+        )
+        assert all(output_size % self.tp_size == 0 for output_size in output_sizes)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        output_dim = getattr(param, "output_dim", None)
+        if output_dim is not None:
+            shards = []
+            current_offset = 0
+            for output_size in self.output_sizes:
+                loaded_shard = loaded_weight.narrow(
+                    output_dim, current_offset, output_size
+                )
+                shard_size = output_size // self.tp_size
+                loaded_shard = loaded_shard.narrow(
+                    output_dim, self.tp_rank * shard_size, shard_size
+                )
+                shards.append(loaded_shard)
+                current_offset += output_size
+            loaded_weight = torch.cat(shards, dim=output_dim)
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+
+class WeightOnlyFP8RowParallelLinear(nn.Module):
+    """Row-parallel storage-only e4m3 FP8 linear."""
+
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = True,
+        compute_dtype: torch.dtype | None = None,
+        input_is_parallel: bool = True,
+        reduce_results: bool = True,
+        tp_group=None,
+        enable_fused_w8a8: bool | None = None,
+    ) -> None:
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.compute_dtype = compute_dtype
+        self.input_is_parallel = input_is_parallel
+        self.reduce_results = reduce_results
+        self.enable_fused_w8a8 = _resolve_enable_fused_w8a8(enable_fused_w8a8)
+        self.tp_group = tp_group or get_tp_group()
+        self.tp_size = get_group_size(self.tp_group)
+        self.tp_rank = get_group_rank(self.tp_group)
+        self.in_features_per_partition = divide(in_features, self.tp_size)
+        self.weight = nn.Parameter(
+            torch.empty(
+                out_features,
+                self.in_features_per_partition,
+                dtype=FP8_WEIGHT_DTYPE,
+            ),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight,
+            {
+                "input_dim": 1,
+                "weight_loader": self.weight_loader,
+            },
+        )
+        self.weight_scale = nn.Parameter(
+            torch.empty(out_features, dtype=torch.float32),
+            requires_grad=False,
+        )
+        set_weight_attrs(
+            self.weight_scale,
+            {
+                "missing_param_init": "error",
+                "weight_loader": self.weight_loader,
+            },
+        )
+        if bias:
+            self.bias = nn.Parameter(
+                torch.empty(
+                    out_features, dtype=compute_dtype or torch.get_default_dtype()
+                ),
+                requires_grad=False,
+            )
+            set_weight_attrs(self.bias, {"weight_loader": self.weight_loader})
+        else:
+            self.register_parameter("bias", None)
+
+    def weight_loader(
+        self, param: torch.nn.Parameter, loaded_weight: torch.Tensor
+    ) -> None:
+        input_dim = getattr(param, "input_dim", None)
+        if input_dim is not None:
+            shard_size = param.data.shape[input_dim]
+            loaded_weight = loaded_weight.narrow(
+                input_dim, self.tp_rank * shard_size, shard_size
+            )
+        if len(loaded_weight.shape) == 0:
+            loaded_weight = loaded_weight.reshape(1)
+        assert param.data.shape == loaded_weight.shape
+        param.data.copy_(loaded_weight)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.input_is_parallel:
+            input_parallel = x
+        else:
+            input_parallel = split_tensor_along_last_dim(
+                x, num_partitions=self.tp_size
+            )[self.tp_rank].contiguous()
+
+        compute_dtype = self.compute_dtype or x.dtype
+        bias = None if self.tp_rank > 0 else self.bias
+        output_parallel = _apply_weight_only_fp8_linear(
+            input_parallel,
+            self.weight,
+            self.weight_scale,
+            bias,
+            compute_dtype,
+            self.enable_fused_w8a8,
+        )
+        if self.reduce_results and self.tp_size > 1:
+            return tensor_model_parallel_all_reduce(
+                output_parallel, tp_group=self.tp_group
+            )
+        return output_parallel
+
+
+def _resolve_enable_fused_w8a8(value: bool | None) -> bool:
+    if value is not None:
+        return value
+    return envs.SGLANG_DIFFUSION_ENABLE_W8A8_FP8_GEMM
+
+
+def _log_w8a8_fp8_gemm_warning_once() -> None:
+    global _w8a8_fp8_gemm_warning_logged
+    if _w8a8_fp8_gemm_warning_logged:
+        return
+    logger.warning(
+        "%s=1 enables W8A8 FP8 GEMM for weight-only FP8 linears; activations "
+        "are dynamically quantized to FP8 and outputs may differ from the "
+        "official weight-only FP8 path.",
+        W8A8_FP8_GEMM_ENV,
+    )
+    _w8a8_fp8_gemm_warning_logged = True
+
+
+def swap_linears_to_weight_only_fp8(module: nn.Module) -> None:
+    """Recursively replace nn.Linear with WeightOnlyFP8Linear.
+
+    Ideogram FP8 checkpoints provide ``<linear>.weight_scale`` for every
+    quantized linear. Swapping before load lets strict state-dict checks verify
+    both the FP8 weight and its row-wise scale.
+    """
+
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear):
+            replacement = WeightOnlyFP8Linear(
+                child.in_features,
+                child.out_features,
+                bias=child.bias is not None,
+                compute_dtype=child.weight.dtype,
+            )
+            setattr(module, name, replacement)
+        else:
+            swap_linears_to_weight_only_fp8(child)