sgl-project--sglang/python/sglang/multimodal_gen/runtime/layers/quantization/modelslim.py

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)