chore: import upstream snapshot with attribution

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)