[diffusion] quant: add support for svdquant and nunchaku (#18549)

Co-authored-by: AichenF <aichenf@nvidia.com>
Co-authored-by: jianyingzhu <53300651@qq.com>

python/sglang/multimodal_gen/configs/models/dits/qwenimage.py

@@ -28,6 +28,10 @@ class QwenImageArchConfig(DiTArchConfig):
         default_factory=lambda: {
             # LoRA mappings
             r"^(transformer_blocks\.\d+\.attn\..*\.lora_[AB])\.default$": r"\1",
+            # SVDquant mappings
+            r"(.*)\.add_qkv_proj\.(.+)$": r"\1.to_added_qkv.\2",
+            r"(transformer_blocks\.\d+\.(img_mlp|txt_mlp)\..*\.(smooth_factor_orig|wcscales))$": r"\1",
+            r".*\.wtscale$": r"",
         }
     )
 

python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py

@@ -186,6 +186,13 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         ] * batch_size
         txt_seq_lens = [prompt_embeds[0].shape[1]]
 
+        if rotary_emb is None:
+            return {
+                "img_shapes": img_shapes,
+                "txt_seq_lens": txt_seq_lens,
+                "freqs_cis": None,
+            }
+
         freqs_cis = self.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
         )
@@ -195,6 +202,7 @@ class QwenImagePipelineConfig(ImagePipelineConfig):
         return {
             "txt_seq_lens": txt_seq_lens,
             "freqs_cis": (img_cache, txt_cache),
+            "img_shapes": img_shapes,
         }
 
     def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype):
@@ -254,6 +262,14 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
             ],
         ] * batch_size
         txt_seq_lens = [prompt_embeds[0].shape[1]]
+
+        if rotary_emb is None:
+            return {
+                "img_shapes": img_shapes,
+                "txt_seq_lens": txt_seq_lens,
+                "freqs_cis": None,
+            }
+
         freqs_cis = QwenImagePipelineConfig.get_freqs_cis(
             img_shapes, txt_seq_lens, rotary_emb, device, dtype
         )
@@ -271,6 +287,7 @@ class QwenImageEditPipelineConfig(QwenImagePipelineConfig):
         return {
             "txt_seq_lens": txt_seq_lens,
             "freqs_cis": (img_cache, txt_cache),
+            "img_shapes": img_shapes,
         }
 
     def preprocess_condition_image(

python/sglang/multimodal_gen/configs/pipeline_configs/zimage.py

@@ -272,7 +272,6 @@ class ZImagePipelineConfig(ImagePipelineConfig):
             device=device,
         ).flatten(0, 2)
 
-        C = self.dit_config.num_channels_latents
         F = 1
         H = height // self.vae_config.arch_config.spatial_compression_ratio
         W = width // self.vae_config.arch_config.spatial_compression_ratio

python/sglang/multimodal_gen/configs/quantization.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+import re
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    is_nunchaku_available,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.utils import StoreBoolean
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class NunchakuSVDQuantArgs:
+    """CLI-facing configuration for Nunchaku (SVDQuant) inference.
+
+    This is intentionally lightweight and only contains arguments needed to
+    construct `runtime.layers.quantization.nunchaku_config.NunchakuConfig`.
+    """
+
+    enable_svdquant: bool = False
+    quantized_model_path: str | None = None
+    quantization_precision: str | None = None  # "int4" or "nvfp4"
+    quantization_rank: int | None = None
+    quantization_act_unsigned: bool = False
+
+    def _adjust_config(self) -> None:
+        """infer precision and rank from filename if not provided"""
+        if self.quantized_model_path and not self.enable_svdquant:
+            self.enable_svdquant = True
+
+        if not self.enable_svdquant or not self.quantized_model_path:
+            return
+
+        inferred_precision = None
+        inferred_rank = None
+
+        filename = os.path.basename(self.quantized_model_path)
+        # Expected pattern: svdq-{precision}_r{rank}-...
+        # e.g., svdq-int4_r32-qwen-image.safetensors
+        match = re.search(r"svdq-(int4|fp4)_r(\d+)", filename)
+
+        if match:
+            p_str, r_str = match.groups()
+            inferred_precision = "nvfp4" if p_str == "fp4" else "int4"
+            inferred_rank = int(r_str)
+
+        if self.quantization_precision is None:
+            self.quantization_precision = inferred_precision or "int4"
+            if inferred_precision:
+                logger.info(
+                    f"inferred --quantization-precision: {self.quantization_precision} "
+                    f"from --quantized-model-path: {self.quantized_model_path}"
+                )
+
+        if self.quantization_rank is None:
+            self.quantization_rank = inferred_rank or 32
+            if inferred_rank:
+                logger.info(
+                    f"inferred --quantization-rank: {self.quantization_rank} "
+                    f"from --quantized-model-path: {self.quantized_model_path}"
+                )
+
+    def validate(self) -> None:
+        # TODO: warn if the served model doesn't support nunchaku
+        self._adjust_config()
+
+        if not self.enable_svdquant:
+            return
+
+        if not current_platform.is_cuda():
+            raise ValueError(
+                "Nunchaku SVDQuant is only supported on NVIDIA CUDA GPUs "
+                "(Ampere SM8x or SM12x)."
+            )
+
+        device_count = torch.cuda.device_count()
+
+        unsupported: list[str] = []
+        for i in range(device_count):
+            major, minor = torch.cuda.get_device_capability(i)
+            if major == 9:
+                unsupported.append(f"cuda:{i} (SM{major}{minor}, Hopper)")
+            elif major not in (8, 12):
+                unsupported.append(f"cuda:{i} (SM{major}{minor})")
+
+        if unsupported:
+            raise ValueError(
+                "Nunchaku SVDQuant is currently only supported on Ampere (SM8x) or SM12x GPUs; "
+                "Hopper (SM90) is not supported. "
+                f"Unsupported devices: {', '.join(unsupported)}. "
+                "Disable it with --enable-svdquant false."
+            )
+
+        if not self.quantized_model_path:
+            raise ValueError(
+                "--enable-svdquant requires --quantized-model-path to be set"
+            )
+
+        if not is_nunchaku_available():
+            raise ValueError(
+                "Nunchaku is enabled, but not installed. Please refer to https://nunchaku.tech/docs/nunchaku/installation/installation.html for detailed installation methods."
+            )
+
+        if self.quantization_precision not in ("int4", "nvfp4"):
+            raise ValueError(
+                f"Invalid --quantization-precision: {self.quantization_precision}. "
+                "Must be one of: int4, nvfp4"
+            )
+
+        if self.quantization_rank <= 0:
+            raise ValueError(
+                f"Invalid --quantization-rank: {self.quantization_rank}. Must be > 0"
+            )
+
+    @staticmethod
+    def add_cli_args(parser) -> None:
+        parser.add_argument(
+            "--enable-svdquant",
+            action=StoreBoolean,
+            default=NunchakuSVDQuantArgs.enable_svdquant,
+            help="Enable Nunchaku SVDQuant (W4A4-style) inference.",
+        )
+        parser.add_argument(
+            "--quantized-model-path",
+            type=str,
+            default=NunchakuSVDQuantArgs.quantized_model_path,
+            help=(
+                "Path to pre-quantized Nunchaku weights. Can be a single .safetensors "
+                "file or a directory containing .safetensors."
+            ),
+        )
+        parser.add_argument(
+            "--quantization-precision",
+            type=str,
+            default=None,
+            help="Quantization precision: int4 or nvfp4. If not specified, inferred from model path or defaults to int4.",
+        )
+        parser.add_argument(
+            "--quantization-rank",
+            type=int,
+            default=None,
+            help="SVD low-rank dimension (e.g., 32). If not specified, inferred from model path or defaults to 32.",
+        )
+        parser.add_argument(
+            "--quantization-act-unsigned",
+            action=StoreBoolean,
+            default=NunchakuSVDQuantArgs.quantization_act_unsigned,
+            help="Use unsigned activation quantization (if supported).",
+        )
+
+    @classmethod
+    def from_dict(cls, kwargs: dict[str, Any]) -> "NunchakuSVDQuantArgs":
+        # Map CLI/config keys to dataclass fields (keep backwards compatibility).
+        return cls(
+            enable_svdquant=bool(kwargs.get("enable_svdquant", cls.enable_svdquant)),
+            quantized_model_path=kwargs.get(
+                "quantized_model_path", cls.quantized_model_path
+            ),
+            quantization_precision=kwargs.get("quantization_precision"),
+            quantization_rank=kwargs.get("quantization_rank"),
+            quantization_act_unsigned=bool(
+                kwargs.get("quantization_act_unsigned", cls.quantization_act_unsigned)
+            ),
+        )

python/sglang/multimodal_gen/docs/quantization.md

@@ -0,0 +1,123 @@
+# Quantization
+
+This document introduces the model quantization schemes supported in SGLang and how to use them to reduce memory usage and accelerate inference.
+
+## Nunchaku (SVDQuant)
+
+### Introduction
+
+**SVDQuant** is a Post-Training Quantization (PTQ) technique for diffusion models that quantizes model weights and activations to 4-bit precision (W4A4) while maintaining high visual quality. This method uses Singular Value Decomposition (SVD) to decompose the weight matrix into low-rank components and residuals, effectively absorbing outliers in activations, making 4-bit quantization possible.
+
+**Nunchaku** is a high-performance inference engine that implements SVDQuant, optimized for low-bit neural networks. It is not Quantization-Aware Training (QAT), but directly quantizes pre-trained models.
+
+Paper: [SVDQuant: Absorbing Outliers by Low-Rank Components for 4-Bit Diffusion Models](https://arxiv.org/abs/2411.05007) (ICLR 2025 Spotlight)
+
+### Key Features
+
+SVDQuant significantly reduces memory usage and accelerates inference while maintaining visual quality:
+
+- **Memory Optimization**: Reduces memory usage by **3.6×** compared to BF16 models.
+- **Inference Acceleration**:
+  - **3.0×** faster than the NF4 (W4A16) baseline on desktop/laptop RTX 4090 GPUs.
+  - **8.7×** speedup on laptop RTX 4090 by eliminating CPU offloading compared to 16-bit models.
+  - **3.1×** faster than BF16 and NF4 models on RTX 5090 GPUs with NVFP4.
+
+### Supported Precisions
+
+Nunchaku supports two quantization precisions:
+
+- **INT4**: Standard INT4 quantization, supported on NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above).
+- **NVFP4**: FP4 quantization, providing better image quality on newer cards like the RTX 5090.
+
+### Usage
+
+#### 1. Install Nunchaku
+
+```bash
+pip install nunchaku
+```
+
+For more installation information, please refer to the [Nunchaku Official Documentation](https://nunchaku.tech/docs/nunchaku/installation/installation.html).
+
+#### 2. Download Quantized Models
+
+Nunchaku provides pre-quantized model weights available on Hugging Face:
+
+- [nunchaku-ai/nunchaku-qwen-image](https://huggingface.co/nunchaku-ai/nunchaku-qwen-image)
+- [nunchaku-ai/nunchaku-flux](https://huggingface.co/nunchaku-ai/nunchaku-flux)
+
+Taking Qwen-Image as an example, several quantized models with different configurations are provided:
+
+| Filename | Precision | Rank | Usage |
+|----------|-----------|------|-------|
+| `svdq-int4_r32-qwen-image.safetensors` | INT4 | 32 | Standard Version |
+| `svdq-int4_r128-qwen-image.safetensors` | INT4 | 128 | High-Quality Version |
+| `svdq-fp4_r32-qwen-image.safetensors` | NVFP4 | 32 | RTX 5090 Standard Version |
+| `svdq-fp4_r128-qwen-image.safetensors` | NVFP4 | 128 | RTX 5090 High-Quality Version |
+| `svdq-int4_r32-qwen-image-lightningv1.0-4steps.safetensors` | INT4 | 32 | Lightning 4-Step Version |
+| `svdq-int4_r128-qwen-image-lightningv1.1-8steps.safetensors` | INT4 | 128 | Lightning 8-Step Version |
+
+> **Note**: Higher Rank usually means better image quality, but with slightly increased memory usage and computation.
+
+#### 3. Run Quantized Models
+
+SGLang features **smart auto-detection** for Nunchaku models. In most cases, you only need to provide the path to the quantized weights, and the precision and rank will be automatically inferred from the filename.
+
+**Simplified Command (Recommended):**
+
+```bash
+sglang generate \
+  --model-path Qwen/Qwen-Image \
+  --prompt "change the raccoon to a cute cat" \
+  --save-output \
+  --quantized-model-path /path/to/svdq-int4_r32-qwen-image.safetensors
+```
+
+**Manual Override (If needed):**
+
+If your filename doesn't follow the standard naming convention, or you want to force specific settings:
+
+- `--enable-svdquant`: Manually enable SVDQuant.
+- `--quantization-precision`: Set to `int4` or `nvfp4`.
+- `--quantization-rank`: Set the SVD rank (e.g., 32, 128).
+- `--quantization-act-unsigned` (Optional): Use unsigned activation quantization.
+
+Example with manual overrides:
+
+```bash
+sglang generate \
+  --model-path Qwen/Qwen-Image \
+  --prompt "a beautiful sunset" \
+  --enable-svdquant \
+  --quantized-model-path /path/to/custom_model.safetensors \
+  --quantization-precision int4 \
+  --quantization-rank 128
+```
+
+#### 4. Configuration Recommendations
+
+Choose the appropriate configuration based on your hardware and requirements:
+
+| Scenario | Recommended Config | Description |
+|----------|-------------------|-------------|
+| Standard Use (20/30/40 Series GPU) | INT4 + Rank 32 | Balanced performance and quality |
+| Quality Focus (Sufficient VRAM) | INT4 + Rank 128 | Better image quality |
+| RTX 5090 Standard Use | NVFP4 + Rank 32 | Utilizes FP4 hardware acceleration |
+| RTX 5090 Quality Focus | NVFP4 + Rank 128 | Best image quality |
+| Fast Prototyping/Preview | Lightning 4-Step Version | Extremely fast generation, slightly reduced quality |
+
+### Notes
+
+1.  Model Path Correspondence: `--model-path` should point to the original non-quantized model (for loading config and tokenizer, etc.), while `--quantized-model-path` points to the quantized weight file.
+
+2.  Auto-Detection Requirements: For auto-detection to work, the filename must contain the pattern `svdq-{precision}_r{rank}` (e.g., `svdq-int4_r32`).
+
+3.  GPU Compatibility:
+    -   INT4: Supports NVIDIA GPUs with Compute Capability 7.0+ (RTX 20 series and above).
+    -   NVFP4: Optimized mainly for newer cards like the RTX 50 series that support FP4.
+
+4.  Lightning Models: When using Lightning versions, adjust `--num-inference-steps` accordingly (usually 4 or 8 steps).
+
+### Custom Model Quantization
+
+If you want to quantize your own models, you can use the [DeepCompressor](https://github.com/mit-han-lab/deepcompressor) tool. For detailed instructions, please refer to the Nunchaku official documentation.

python/sglang/multimodal_gen/runtime/layers/linear.py

@@ -17,7 +17,7 @@ from sglang.multimodal_gen.runtime.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
-from sglang.multimodal_gen.runtime.layers.quantization.base_config import (
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
@@ -165,7 +165,6 @@ class LinearBase(torch.nn.Module):
     Args:
         input_size: input dimension of the linear layer.
         output_size: output dimension of the linear layer.
-        bias: If true, add bias.
         skip_bias_add: If true, skip adding bias but instead return it.
         params_dtype: Data type for the parameters.
         quant_config: Quantization configure.

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

@@ -2,7 +2,7 @@
 
 from typing import Literal, get_args
 
-from sglang.multimodal_gen.runtime.layers.quantization.base_config import (
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
 )
 
@@ -23,17 +23,7 @@ def register_quantization_config(quantization: str):
     Args:
         quantization (str): The quantization method name.
 
-    Examples:
-        >>> from sglang.multimodal_gen.runtime.layers.quantization import register_quantization_config
-        >>> from sglang.multimodal_gen.runtime.layers.quantization import get_quantization_config
-        >>> from sglang.multimodal_gen.runtime.layers.quantization.base_config import QuantizationConfig
-        >>>
-        >>> @register_quantization_config("my_quant")
-        ... class MyQuantConfig(QuantizationConfig):
-        ...     pass
-        >>>
-        >>> get_quantization_config("my_quant")
-        <class 'MyQuantConfig'>
+
     """  # noqa: E501
 
     def _wrapper(quant_config_cls):
@@ -63,7 +53,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
     return method_to_config[quantization]
 
 
-all = [
+__all__ = [
     "QuantizationMethods",
     "QuantizationConfig",
     "get_quantization_config",

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

@@ -65,6 +65,9 @@ def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) ->
 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

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

@@ -0,0 +1,289 @@
+# 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__)
+
+SVDQ_W4A4_LAYER_PATTERNS = [
+    "attn.to_qkv",
+    "attn.to_out",
+    "attn.add_qkv_proj",
+    "attn.to_add_out",
+    "img_mlp",
+    "txt_mlp",
+]
+
+AWQ_W4A16_LAYER_PATTERNS = [
+    "img_mod",
+    "txt_mod",
+]
+
+SKIP_QUANTIZATION_PATTERNS = [
+    "norm",
+    "embed",
+    "rotary",
+    "pos_embed",
+]
+
+
+@lru_cache(maxsize=1)
+def is_nunchaku_available() -> bool:
+    try:
+        import nunchaku  # noqa
+
+        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
+        quantized_model_path: Path to pre-quantized model weights (.safetensors)
+    """
+
+    precision: str = "int4"  # "int4" or "nvfp4"
+    rank: int = 32
+    group_size: Optional[int] = None
+    act_unsigned: bool = False
+    quantized_model_path: Optional[str] = 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)),
+            quantized_model_path=config.get("quantized_model_path"),
+        )
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[QuantizeMethodBase]:
+
+        if not isinstance(layer, LinearBase):
+            return None
+
+        for pattern in SKIP_QUANTIZATION_PATTERNS:
+            if pattern in prefix.lower():
+                return None
+
+        for pattern in SVDQ_W4A4_LAYER_PATTERNS:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+                return NunchakuSVDQLinearMethod(
+                    precision=self.precision,
+                    rank=self.rank,
+                    act_unsigned=self.act_unsigned,
+                )
+
+        for pattern in AWQ_W4A16_LAYER_PATTERNS:
+            if pattern in prefix:
+                from ..nunchaku_linear import NunchakuAWQLinearMethod
+
+                return NunchakuAWQLinearMethod(
+                    group_size=64,
+                )
+
+        from ..nunchaku_linear import NunchakuSVDQLinearMethod
+
+        return NunchakuSVDQLinearMethod(
+            precision=self.precision,
+            rank=self.rank,
+            act_unsigned=self.act_unsigned,
+        )
+
+    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,
+            "quantized_model_path": self.quantized_model_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/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.models.utils import set_weight_attrs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+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/vocab_parallel_embedding.py

@@ -15,7 +15,7 @@ from sglang.multimodal_gen.runtime.distributed import (
     get_tp_group,
     tensor_model_parallel_all_reduce,
 )
-from sglang.multimodal_gen.runtime.layers.quantization.base_config import (
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
     method_has_implemented_embedding,

python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py

@@ -1,9 +1,16 @@
+import inspect
+import json
+import logging
 import os
 from copy import deepcopy
+from typing import Any
 
 import torch
 
 from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    _patch_nunchaku_scales,
+)
 from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import (
     ComponentLoader,
 )
@@ -16,8 +23,9 @@ from sglang.multimodal_gen.runtime.models.registry import ModelRegistry
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
     get_diffusers_component_config,
+    get_metadata_from_safetensors_file,
 )
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.logging_utils import get_log_level, init_logger
 from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
 
 logger = init_logger(__name__)
@@ -29,6 +37,34 @@ class TransformerLoader(ComponentLoader):
     component_names = ["transformer", "audio_dit", "video_dit"]
     expected_library = "diffusers"
 
+    def get_list_of_safetensors_to_load(
+        self, server_args: ServerArgs, component_model_path: str
+    ) -> list[str]:
+        """
+        get list of safetensors to load.
+
+        For some quantization framework, if --quantized-model-path is provided, load from this path instead of main model
+        """
+        nunchaku_config = server_args.nunchaku_config
+
+        if nunchaku_config is not None and nunchaku_config.quantized_model_path:
+            # load from quantized_model_path if applicable
+            weights_path = nunchaku_config.quantized_model_path
+
+            logger.info("Using quantized model weights from: %s", weights_path)
+            if os.path.isfile(weights_path) and weights_path.endswith(".safetensors"):
+                safetensors_list = [weights_path]
+            else:
+                safetensors_list = _list_safetensors_files(weights_path)
+        else:
+            weights_path = component_model_path
+            safetensors_list = _list_safetensors_files(weights_path)
+
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {weights_path}")
+
+        return safetensors_list
+
     def load_customized(
         self, component_model_path: str, server_args: ServerArgs, component_name: str
     ):
@@ -57,46 +93,50 @@ class TransformerLoader(ComponentLoader):
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
-        # Find all safetensors files
-        safetensors_list = _list_safetensors_files(component_model_path)
-        if not safetensors_list:
-            raise ValueError(f"No safetensors files found in {component_model_path}")
+        nunchaku_config = server_args.nunchaku_config
 
-        # Check if we should use custom initialization weights
-        custom_weights_path = getattr(
-            server_args, "init_weights_from_safetensors", None
+        if nunchaku_config is not None:
+            # respect dtype from checkpoint
+            # TODO: improve the condition
+            param_dtype = None
+
+            # check if the specified nunchaku quantized model path matches with the specified model path
+            original_dit_cls_name = json.loads(
+                get_metadata_from_safetensors_file(
+                    nunchaku_config.quantized_model_path
+                ).get("config")
+            )["_class_name"]
+            specified_dit_cls_name = str(model_cls.__name__)
+            if original_dit_cls_name != specified_dit_cls_name:
+                raise Exception(
+                    f"Class name of DiT specified in nunchaku quantized model_path: {original_dit_cls_name} does not match that of specified DiT name: {specified_dit_cls_name}"
+                )
+        else:
+            param_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        safetensors_list = self.get_list_of_safetensors_to_load(
+            server_args, component_model_path
         )
-        use_custom_weights = False
-
-        if use_custom_weights:
-            logger.info(
-                "Using custom initialization weights from: %s", custom_weights_path
-            )
-            assert (
-                custom_weights_path is not None
-            ), "Custom initialization weights must be provided"
-            if os.path.isdir(custom_weights_path):
-                safetensors_list = _list_safetensors_files(custom_weights_path)
-            else:
-                assert custom_weights_path.endswith(
-                    ".safetensors"
-                ), "Custom initialization weights must be a safetensors file"
-                safetensors_list = [custom_weights_path]
-
-        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
 
         logger.info(
-            "Loading %s from %s safetensors files, default_dtype: %s",
+            "Loading %s from %s safetensors files %s, param_dtype: %s",
             cls_name,
             len(safetensors_list),
-            default_dtype,
+            f": {safetensors_list}" if get_log_level() == logging.DEBUG else "",
+            param_dtype,
         )
 
+        init_params: dict[str, Any] = {"config": dit_config, "hf_config": hf_config}
+        if (
+            nunchaku_config is not None
+            and "quant_config" in inspect.signature(model_cls.__init__).parameters
+        ):
+            init_params["quant_config"] = nunchaku_config
+
         # Load the model using FSDP loader
-        assert server_args.hsdp_shard_dim is not None
         model = maybe_load_fsdp_model(
             model_cls=model_cls,
-            init_params={"config": dit_config, "hf_config": hf_config},
+            init_params=init_params,
             weight_dir_list=safetensors_list,
             device=get_local_torch_device(),
             hsdp_replicate_dim=server_args.hsdp_replicate_dim,
@@ -105,17 +145,22 @@ class TransformerLoader(ComponentLoader):
             pin_cpu_memory=server_args.pin_cpu_memory,
             fsdp_inference=server_args.use_fsdp_inference,
             # TODO(will): make these configurable
-            param_dtype=torch.bfloat16,
+            param_dtype=param_dtype,
             reduce_dtype=torch.float32,
             output_dtype=None,
             strict=False,
         )
 
+        if nunchaku_config is not None:
+            _patch_nunchaku_scales(model, safetensors_list)
+
         total_params = sum(p.numel() for p in model.parameters())
         logger.info("Loaded model with %.2fB parameters", total_params / 1e9)
 
-        assert (
-            next(model.parameters()).dtype == default_dtype
-        ), "Model dtype does not match default dtype"
+        # considering the existent of mixed-precision models (e.g., nunchaku)
+        if next(model.parameters()).dtype != param_dtype:
+            logger.warning(
+                f"Model dtype does not match expected param dtype, {next(model.parameters()).dtype} vs {param_dtype}"
+            )
 
         return model

python/sglang/multimodal_gen/runtime/loader/fsdp_load.py

@@ -30,6 +30,7 @@ from sglang.multimodal_gen.runtime.loader.utils import (
 from sglang.multimodal_gen.runtime.loader.weight_utils import (
     safetensors_weights_iterator,
 )
+from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import set_mixed_precision_policy
 
@@ -40,7 +41,12 @@ def _make_param_like(
     actual_param: torch.nn.Parameter, tensor: torch.Tensor
 ) -> torch.nn.Parameter:
     cls = actual_param.__class__
-    new_param = cls.__new__(cls, tensor)
+    # nn.Parameter defaults to requires_grad=True, which is illegal for non-floating/complex dtypes (e.g., int8/FP8
+    # quantized weights).
+    try:
+        new_param = cls.__new__(cls, tensor, requires_grad=False)
+    except TypeError:
+        new_param = cls.__new__(cls, tensor)
     new_param.__dict__.update(actual_param.__dict__)
     new_param.requires_grad = False
     return new_param
@@ -74,26 +80,25 @@ def maybe_load_fsdp_model(
     """
     # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
     # manually casting the inputs to the model
+    default_torch_dtype = param_dtype if param_dtype else torch.bfloat16
     mp_policy = MixedPrecisionPolicy(
-        param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
+        default_torch_dtype, reduce_dtype, output_dtype, cast_forward_inputs=False
     )
 
     set_mixed_precision_policy(
-        param_dtype=param_dtype,
+        param_dtype=default_torch_dtype,
         reduce_dtype=reduce_dtype,
         output_dtype=output_dtype,
         mp_policy=mp_policy,
     )
 
-    with set_default_torch_dtype(param_dtype), torch.device("meta"):
+    with set_default_torch_dtype(default_torch_dtype), torch.device("meta"):
         model = model_cls(**init_params)
 
     # Check if we should use FSDP
     use_fsdp = fsdp_inference
 
     # Disable FSDP for MPS as it's not compatible
-    from sglang.multimodal_gen.runtime.platforms import current_platform
-
     if current_platform.is_mps():
         use_fsdp = False
         logger.info("Disabling FSDP for MPS platform as it's not compatible")
@@ -191,7 +196,7 @@ def shard_model(
     # TODO(will): don't reshard after forward for the last layer to save on the
     # all-gather that will immediately happen Shard the model with FSDP,
     for n, m in reversed(list(model.named_modules())):
-        if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):
+        if any([shard_condition(n, m) for shard_condition in fsdp_shard_conditions]):  # type: ignore
             fully_shard(m, **fsdp_kwargs)
             num_layers_sharded += 1
 
@@ -209,7 +214,7 @@ def load_model_from_full_model_state_dict(
     model: FSDPModule | torch.nn.Module,
     full_sd_iterator: Generator[tuple[str, torch.Tensor], None, None],
     device: torch.device,
-    param_dtype: torch.dtype,
+    param_dtype: torch.dtype | None,
     strict: bool = False,
     cpu_offload: bool = False,
     param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
@@ -221,7 +226,7 @@ def load_model_from_full_model_state_dict(
         model (Union[FSDPModule, torch.nn.Module]): Model to generate fully qualified names for cpu_state_dict
         full_sd_iterator (Generator): an iterator yielding (param_name, tensor) pairs
         device (torch.device): device used to move full state dict tensors
-        param_dtype (torch.dtype): dtype used to move full state dict tensors
+        param_dtype (torch.dtype): dtype used to move full state dict tensors. If none, respect original dtype from checkpoint
         strict (bool): flag to check if to load the model in strict mode
         cpu_offload (bool): flag to check if FSDP offload is enabled
         param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
@@ -245,6 +250,9 @@ def load_model_from_full_model_state_dict(
     # sort parameter names to ensure all ranks process parameters in the same order
     sorted_param_names = sorted(custom_param_sd.keys())
 
+    requires_grad = False
+
+    # shard from loaded state_dict, custom_param_sd -> sharded_sd
     for target_param_name in sorted_param_names:
         full_tensor = custom_param_sd[target_param_name]
         meta_sharded_param = meta_sd.get(target_param_name)
@@ -259,8 +267,10 @@ def load_model_from_full_model_state_dict(
                     f"Parameter '{target_param_name}' from checkpoint not found in model; skipping. This is expected for optional parameters."
                 )
                 continue
+
+        target_dtype = param_dtype if param_dtype else full_tensor.dtype
         if not hasattr(meta_sharded_param, "device_mesh"):
-            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
+            full_tensor = full_tensor.to(device=device, dtype=target_dtype)
             actual_param = param_dict.get(target_param_name)
             weight_loader = (
                 getattr(actual_param, "weight_loader", None)
@@ -270,12 +280,20 @@ def load_model_from_full_model_state_dict(
             if weight_loader is not None:
                 assert actual_param is not None
                 sharded_tensor = torch.empty_like(
-                    meta_sharded_param, device=device, dtype=param_dtype
+                    meta_sharded_param, device=device, dtype=target_dtype
                 )
+                # Preserve requires_grad flag to avoid errors with non-floating dtypes
+                requires_grad = getattr(meta_sharded_param, "requires_grad", False)
                 temp_param = _make_param_like(actual_param, sharded_tensor)
+                if not (
+                    sharded_tensor.is_floating_point() or sharded_tensor.is_complex()
+                ):
+                    requires_grad = False
+                temp_param.requires_grad = requires_grad
                 weight_loader(temp_param, full_tensor)
                 sharded_tensor = temp_param.data
             else:
+                # In cases where parts of the model aren't sharded, some parameters will be plain tensors
                 sharded_tensor = full_tensor
 
             # Important: `cpu_offload` is intended for FSDP-managed parameter movement.
@@ -290,7 +308,7 @@ def load_model_from_full_model_state_dict(
             if cpu_offload and not is_fsdp_model:
                 sharded_tensor = sharded_tensor.cpu()
         else:
-            full_tensor = full_tensor.to(device=device, dtype=param_dtype)
+            full_tensor = full_tensor.to(device=device, dtype=target_dtype)
             sharded_tensor = distribute_tensor(
                 full_tensor,
                 meta_sharded_param.device_mesh,
@@ -298,16 +316,27 @@ def load_model_from_full_model_state_dict(
             )
             if cpu_offload:
                 sharded_tensor = sharded_tensor.to("cpu")
-        sharded_sd[target_param_name] = nn.Parameter(sharded_tensor)
+
+        requires_grad = False
+        sharded_sd[target_param_name] = nn.Parameter(
+            sharded_tensor, requires_grad=requires_grad
+        )
 
     model.reverse_param_names_mapping = reverse_param_names_mapping
+    # parameters in nn.Module that doesn't exist in safetensor files
     unused_keys = set(meta_sd.keys()) - set(sharded_sd.keys())
     if unused_keys:
         logger.warning("Found unloaded parameters in meta state dict: %s", unused_keys)
 
-    # List of allowed parameter name patterns
-    ALLOWED_NEW_PARAM_PATTERNS = ["gate_compress"]  # Can be extended as needed
+    # for nunchaku
+    ALLOWED_NEW_PARAM_PATTERNS = [
+        "gate_compress",
+        "wcscales",
+        "wtscale",
+        "bias",
+    ]
     for new_param_name in unused_keys:
+        # check unallowed missing params
         if not any(pattern in new_param_name for pattern in ALLOWED_NEW_PARAM_PATTERNS):
             logger.error(
                 "Unsupported new parameter: %s. Allowed patterns: %s",
@@ -318,17 +347,22 @@ def load_model_from_full_model_state_dict(
                 f"New parameter '{new_param_name}' is not supported. "
                 f"Currently only parameters containing {ALLOWED_NEW_PARAM_PATTERNS} are allowed."
             )
+
         meta_sharded_param = meta_sd.get(new_param_name)
+
+        if "wcscales" in new_param_name or "wtscale" in new_param_name:
+            init_like = torch.ones_like
+        else:
+            init_like = torch.zeros_like
+
         if not hasattr(meta_sharded_param, "device_mesh"):
-            # Initialize with zeros
-            sharded_tensor = torch.zeros_like(
+            sharded_tensor = init_like(
                 meta_sharded_param, device=device, dtype=param_dtype
             )
             if cpu_offload and not is_fsdp_model:
                 sharded_tensor = sharded_tensor.cpu()
         else:
-            # Initialize with zeros and distribute
-            full_tensor = torch.zeros_like(
+            full_tensor = init_like(
                 meta_sharded_param, device=device, dtype=param_dtype
             )
             sharded_tensor = distribute_tensor(

python/sglang/multimodal_gen/runtime/loader/utils.py

@@ -86,6 +86,8 @@ def hf_to_custom_state_dict(
         target_param_name, merge_index, num_params_to_merge = param_names_mapping(
             source_param_name
         )
+        if target_param_name == "" or target_param_name is None:  # type: ignore[comparison-overlap]
+            continue
         reverse_param_names_mapping[target_param_name] = (
             source_param_name,
             merge_index,

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -352,6 +352,7 @@ OOM detected. Possible solutions:
     2. Enable SP and/or TP
     3. Opt for a sparse-attention backend
     4. Enable FSDP by `--use-fsdp-inference` (in a multi-GPU setup)
+    5. Enable quantization (e.g. nunchaku)
   Or, open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose
 """
 

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -5,10 +5,12 @@
 import functools
 from typing import Any, Dict, List, Optional, Tuple, Union
 
+import diffusers
 import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from diffusers.models.attention import FeedForward
 from diffusers.models.embeddings import TimestepEmbedding, Timesteps
 from diffusers.models.modeling_outputs import Transformer2DModelOutput
 from diffusers.models.normalization import AdaLayerNormContinuous
@@ -26,8 +28,16 @@ from sglang.multimodal_gen.runtime.layers.layernorm import (
 )
 from sglang.multimodal_gen.runtime.layers.linear import (
     ColumnParallelLinear,
+    MergedColumnParallelLinear,
     RowParallelLinear,
 )
+from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
+    QuantizationConfig,
+)
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    NunchakuConfig,
+    is_nunchaku_available,
+)
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
 )
@@ -45,19 +55,36 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 logger = init_logger(__name__)  # pylint: disable=invalid-name
 _is_cuda = current_platform.is_cuda()
 
+try:
+    from nunchaku.models.attention import NunchakuFeedForward  # type: ignore[import]
+except Exception:
+    NunchakuFeedForward = None
+
 
 def _get_qkv_projections(
     attn: "QwenImageCrossAttention", hidden_states, encoder_hidden_states=None
 ):
-    img_query, _ = attn.to_q(hidden_states)
-    img_key, _ = attn.to_k(hidden_states)
-    img_value, _ = attn.to_v(hidden_states)
+    if attn.use_fused_qkv:
+        img_qkv, _ = attn.to_qkv(hidden_states)
+        img_query, img_key, img_value = [
+            x.contiguous() for x in img_qkv.chunk(3, dim=-1)
+        ]
+    else:
+        img_query, _ = attn.to_q(hidden_states)
+        img_key, _ = attn.to_k(hidden_states)
+        img_value, _ = attn.to_v(hidden_states)
 
     txt_query = txt_key = txt_value = None
     if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
-        txt_query, _ = attn.add_q_proj(encoder_hidden_states)
-        txt_key, _ = attn.add_k_proj(encoder_hidden_states)
-        txt_value, _ = attn.add_v_proj(encoder_hidden_states)
+        if attn.use_fused_added_qkv:
+            txt_qkv, _ = attn.to_added_qkv(encoder_hidden_states)
+            txt_query, txt_key, txt_value = [
+                x.contiguous() for x in txt_qkv.chunk(3, dim=-1)
+            ]
+        else:
+            txt_query, _ = attn.add_q_proj(encoder_hidden_states)
+            txt_key, _ = attn.add_k_proj(encoder_hidden_states)
+            txt_value, _ = attn.add_v_proj(encoder_hidden_states)
 
     return img_query, img_key, img_value, txt_query, txt_key, txt_value
 
@@ -71,14 +98,27 @@ class GELU(nn.Module):
         dim_out (`int`): The number of channels in the output.
         approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
         bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict.
     """
 
     def __init__(
-        self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True
+        self,
+        dim_in: int,
+        dim_out: int,
+        approximate: str = "none",
+        bias: bool = True,
+        quant_config=None,
+        prefix: str = "",
     ):
         super().__init__()
         self.proj = ColumnParallelLinear(
-            dim_in, dim_out, bias=bias, gather_output=False
+            dim_in,
+            dim_out,
+            bias=bias,
+            gather_output=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.proj" if prefix else "",
         )
         self.approximate = approximate
 
@@ -95,9 +135,7 @@ class FeedForward(nn.Module):
         dim (`int`): The number of channels in the input.
         dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
         mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
-        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
         activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
-        final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
         bias (`bool`, defaults to True): Whether to use a bias in the linear layer.
     """
 
@@ -194,14 +232,6 @@ class QwenEmbedRope(nn.Module):
             dim=1,
         )
 
-        # self.rope = NDRotaryEmbedding(
-        #     rope_dim_list=axes_dim,
-        #     rope_theta=theta,
-        #     use_real=False,
-        #     repeat_interleave_real=False,
-        #     dtype=torch.float32 if current_platform.is_mps() or current_platform.is_musa() else torch.float64,
-        # )
-
         # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
         self.scale_rope = scale_rope
 
@@ -546,6 +576,8 @@ class QwenImageCrossAttention(nn.Module):
         context_pre_only: bool = False,
         parallel_attention=False,
         out_dim: int = None,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
     ) -> None:
         assert dim % num_heads == 0
         super().__init__()
@@ -557,38 +589,92 @@ class QwenImageCrossAttention(nn.Module):
         self.eps = eps
         self.parallel_attention = parallel_attention
         self.added_kv_proj_dim = added_kv_proj_dim
+        self.prefix = prefix
+
+        self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)
 
-        # Use separate Q/K/V projections
         self.inner_dim = out_dim if out_dim is not None else head_dim * num_heads
         self.inner_kv_dim = self.inner_dim
-        self.to_q = ColumnParallelLinear(
-            dim, self.inner_dim, bias=True, gather_output=True
-        )
-        self.to_k = ColumnParallelLinear(
-            dim, self.inner_dim, bias=True, gather_output=True
-        )
-        self.to_v = ColumnParallelLinear(
-            dim, self.inner_dim, bias=True, gather_output=True
-        )
+
+        if self.use_fused_qkv:
+            # Use fused QKV projection for nunchaku quantization
+            self.to_qkv = MergedColumnParallelLinear(
+                dim,
+                [self.inner_dim] * 3,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_qkv",
+            )
+        else:
+            # Use separate Q/K/V projections for non-quantized models
+            self.to_q = ColumnParallelLinear(
+                dim,
+                self.inner_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_q",
+            )
+            self.to_k = ColumnParallelLinear(
+                dim,
+                self.inner_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_k",
+            )
+            self.to_v = ColumnParallelLinear(
+                dim,
+                self.inner_dim,
+                bias=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_v",
+            )
 
         if self.qk_norm:
             self.norm_q = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
             self.norm_k = RMSNorm(head_dim, eps=eps) if qk_norm else nn.Identity()
 
         if added_kv_proj_dim is not None:
-            self.add_q_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
-            )
-            self.add_k_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
-            )
-            self.add_v_proj = ColumnParallelLinear(
-                added_kv_proj_dim, self.inner_dim, bias=True, gather_output=True
-            )
+            self.use_fused_added_qkv = isinstance(quant_config, NunchakuConfig)
+            if self.use_fused_added_qkv:
+                self.to_added_qkv = MergedColumnParallelLinear(
+                    added_kv_proj_dim,
+                    [self.inner_dim] * 3,
+                    bias=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_added_qkv",
+                )
+            else:
+                # Use separate Q/K/V projections for non-quantized models
+                self.add_q_proj = ColumnParallelLinear(
+                    added_kv_proj_dim,
+                    self.inner_dim,
+                    bias=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.add_q_proj",
+                )
+                self.add_k_proj = ColumnParallelLinear(
+                    added_kv_proj_dim,
+                    self.inner_dim,
+                    bias=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.add_k_proj",
+                )
+                self.add_v_proj = ColumnParallelLinear(
+                    added_kv_proj_dim,
+                    self.inner_dim,
+                    bias=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.add_v_proj",
+                )
 
         if context_pre_only is not None and not context_pre_only:
             self.to_add_out = ColumnParallelLinear(
-                self.inner_dim, self.dim, bias=out_bias, gather_output=True
+                self.inner_dim,
+                self.dim,
+                bias=out_bias,
+                gather_output=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.to_add_out",
             )
         else:
             self.to_add_out = None
@@ -597,7 +683,12 @@ class QwenImageCrossAttention(nn.Module):
             self.to_out = nn.ModuleList([])
             self.to_out.append(
                 ColumnParallelLinear(
-                    self.inner_dim, self.dim, bias=out_bias, gather_output=True
+                    self.inner_dim,
+                    self.dim,
+                    bias=out_bias,
+                    gather_output=True,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.to_out.0",
                 )
             )
         else:
@@ -719,20 +810,29 @@ class QwenImageTransformerBlock(nn.Module):
         attention_head_dim: int,
         qk_norm: str = "rms_norm",
         eps: float = 1e-6,
+        quant_config: Optional[QuantizationConfig] | NunchakuConfig = None,
+        prefix: str = "",
         zero_cond_t: bool = False,
     ):
         super().__init__()
+        self.prefix = prefix
 
         self.dim = dim
         self.num_attention_heads = num_attention_heads
         self.attention_head_dim = attention_head_dim
+        self.quant_config = quant_config
         self.zero_cond_t = zero_cond_t
 
         # Image processing modules
         self.img_mod = nn.Sequential(
             nn.SiLU(),
             ColumnParallelLinear(
-                dim, 6 * dim, bias=True, gather_output=True
+                dim,
+                6 * dim,
+                bias=True,
+                gather_output=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.img_mod",
             ),  # For scale, shift, gate for norm1 and norm2
         )
         self.img_norm1 = LayerNormScaleShift(
@@ -745,19 +845,23 @@ class QwenImageTransformerBlock(nn.Module):
             added_kv_proj_dim=dim,
             context_pre_only=False,
             head_dim=attention_head_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.attn",
         )
         self.img_norm2 = ScaleResidualLayerNormScaleShift(
-            hidden_size=dim, eps=eps, elementwise_affine=False
-        )
-        self.img_mlp = FeedForward(
-            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
+            dim, eps=eps, elementwise_affine=False
         )
 
         # Text processing modules
         self.txt_mod = nn.Sequential(
             nn.SiLU(),
             ColumnParallelLinear(
-                dim, 6 * dim, bias=True, gather_output=True
+                dim,
+                6 * dim,
+                bias=True,
+                gather_output=True,
+                quant_config=quant_config,
+                prefix=f"{prefix}.txt_mod",
             ),  # For scale, shift, gate for norm1 and norm2
         )
         self.txt_norm1 = LayerNormScaleShift(
@@ -767,12 +871,38 @@ class QwenImageTransformerBlock(nn.Module):
         self.txt_norm2 = ScaleResidualLayerNormScaleShift(
             hidden_size=dim, eps=eps, elementwise_affine=False
         )
-        self.txt_mlp = FeedForward(
-            dim=dim, dim_out=dim, activation_fn="gelu-approximate"
-        )
         # Utils
         self.fuse_mul_add = MulAdd()
 
+        nunchaku_enabled = (
+            quant_config is not None
+            and hasattr(quant_config, "get_name")
+            and quant_config.get_name() == "svdquant"
+            and is_nunchaku_available()
+        )
+        ff_class = (
+            diffusers.models.attention.FeedForward if nunchaku_enabled else FeedForward
+        )
+        self.img_mlp = ff_class(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+        self.txt_mlp = ff_class(
+            dim=dim,
+            dim_out=dim,
+            activation_fn="gelu-approximate",
+        )
+
+        if nunchaku_enabled:
+            nunchaku_kwargs = {
+                "precision": quant_config.precision,
+                "rank": quant_config.rank,
+                "act_unsigned": quant_config.act_unsigned,
+            }
+            self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs)
+            self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs)
+
     def _modulate(
         self,
         x: torch.Tensor,
@@ -880,6 +1010,24 @@ class QwenImageTransformerBlock(nn.Module):
         img_mod_params = self.img_mod[1](temb_img_silu)[0]  # [B, 6*dim]
         txt_mod_params = self.txt_mod[1](temb_txt_silu)[0]  # [B, 6*dim]
 
+        if (
+            self.quant_config is not None
+            and hasattr(self.quant_config, "get_name")
+            and self.quant_config.get_name() == "svdquant"
+        ):
+            # When NOT using nunchaku, reshape mod_params from [B, 6*dim] to [B, dim*6]
+            # When using nunchaku (svdquant), keep original format
+            img_mod_params = (
+                img_mod_params.view(img_mod_params.shape[0], -1, 6)
+                .transpose(1, 2)
+                .reshape(img_mod_params.shape[0], -1)
+            )
+            txt_mod_params = (
+                txt_mod_params.view(txt_mod_params.shape[0], -1, 6)
+                .transpose(1, 2)
+                .reshape(txt_mod_params.shape[0], -1)
+            )
+
         # Split modulation parameters for norm1 and norm2
         img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
         txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
@@ -923,8 +1071,11 @@ class QwenImageTransformerBlock(nn.Module):
             gate_x=img_gate1,
             residual_x=hidden_states,
         )
-        img_mlp_output = self.img_mlp(img_modulated2)
-        hidden_states = self.fuse_mul_add(img_mlp_output[0], img_gate2, hidden_states)
+        img_mlp_output = self.img_mlp(img_modulated2)[0]
+
+        if img_mlp_output.dim() == 2:
+            img_mlp_output = img_mlp_output.unsqueeze(0)
+        hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)
 
         # Process text stream - norm2 + MLP
         txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
@@ -936,9 +1087,12 @@ class QwenImageTransformerBlock(nn.Module):
             scale=txt_scale2,
         )
         txt_gate2 = txt_gate2_raw.unsqueeze(1)
-        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)[0]
+
+        if txt_mlp_output.dim() == 2:
+            txt_mlp_output = txt_mlp_output.unsqueeze(0)
         encoder_hidden_states = self.fuse_mul_add(
-            txt_mlp_output[0], txt_gate2, encoder_hidden_states
+            txt_mlp_output, txt_gate2, encoder_hidden_states
         )
 
         # Clip to prevent overflow for fp16
@@ -973,6 +1127,7 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
         self,
         config: QwenImageDitConfig,
         hf_config: dict[str, Any],
+        quant_config: Optional[QuantizationConfig] = None,
     ):
         super().__init__(config=config, hf_config=hf_config)
         patch_size = config.arch_config.patch_size
@@ -1019,9 +1174,11 @@ class QwenImageTransformer2DModel(CachableDiT, OffloadableDiTMixin):
                     dim=self.inner_dim,
                     num_attention_heads=num_attention_heads,
                     attention_head_dim=attention_head_dim,
+                    quant_config=quant_config,
+                    prefix=f"transformer_blocks.{layer_idx}",
                     zero_cond_t=self.zero_cond_t,
                 )
-                for _ in range(num_layers)
+                for layer_idx in range(num_layers)
             ]
         )
 

python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py

@@ -74,8 +74,11 @@ class ImageEncodingStage(PipelineStage):
             self.move_to_device("cpu")
 
     def move_to_device(self, device):
+        if self.server_args.use_fsdp_inference:
+            return
         fields = [
             "image_processor",
+            "image_encoder",
         ]
         for field in fields:
             processor = getattr(self, field, None)

python/sglang/multimodal_gen/runtime/server_args.py

@@ -21,6 +21,10 @@ import yaml
 
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.pipeline_configs.base import PipelineConfig
+from sglang.multimodal_gen.configs.quantization import NunchakuSVDQuantArgs
+from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
+    NunchakuConfig,
+)
 from sglang.multimodal_gen.runtime.platforms import (
     AttentionBackendEnum,
     current_platform,
@@ -309,6 +313,11 @@ class ServerArgs:
 
     disable_autocast: bool | None = None
 
+    # Quantization / Nunchaku SVDQuant configuration
+    nunchaku_config: NunchakuSVDQuantArgs | NunchakuConfig | None = field(
+        default_factory=NunchakuSVDQuantArgs, repr=False
+    )
+
     # Master port for distributed inference
     # TODO: do not hard code
     master_port: int | None = None
@@ -362,6 +371,23 @@ class ServerArgs:
         """
         return self.host is None or self.port is None
 
+    def adjust_quant_config(self):
+        """validate and adjust"""
+
+        # nunchaku
+        ncfg = self.nunchaku_config
+        ncfg.validate()
+        if not ncfg.enable_svdquant or not ncfg.quantized_model_path:
+            # if nunchaku is not applied
+            self.nunchaku_config = None
+        else:
+            self.nunchaku_config = NunchakuConfig(
+                precision=self.nunchaku_config.quantization_precision,
+                rank=self.nunchaku_config.quantization_rank,
+                act_unsigned=self.nunchaku_config.quantization_act_unsigned,
+                quantized_model_path=self.nunchaku_config.quantized_model_path,
+            )
+
     def adjust_offload(self):
         if self.pipeline_config.task_type.is_image_gen():
             logger.info(
@@ -429,6 +455,7 @@ class ServerArgs:
         configure_logger(server_args=self)
 
         self.adjust_offload()
+        self.adjust_quant_config()
 
         if self.attention_backend in ["fa3", "fa4"]:
             self.attention_backend = "fa"
@@ -684,6 +711,9 @@ class ServerArgs:
             help="Disable autocast for denoising loop and vae decoding in pipeline sampling",
         )
 
+        # Nunchaku SVDQuant quantization parameters
+        NunchakuSVDQuantArgs.add_cli_args(parser)
+
         # Master port for distributed inference
         parser.add_argument(
             "--master-port",
@@ -846,6 +876,9 @@ class ServerArgs:
                 pipeline_config = PipelineConfig.from_kwargs(kwargs)
                 logger.debug(f"Using PipelineConfig: {type(pipeline_config)}")
                 server_args_kwargs["pipeline_config"] = pipeline_config
+            elif attr == "nunchaku_config":
+                nunchaku_config = NunchakuSVDQuantArgs.from_dict(kwargs)
+                server_args_kwargs["nunchaku_config"] = nunchaku_config
             elif attr in kwargs:
                 server_args_kwargs[attr] = kwargs[attr]
 
@@ -977,6 +1010,7 @@ class ServerArgs:
 
     def check_server_args(self) -> None:
         """Validate inference arguments for consistency"""
+
         # layerwise offload
         if current_platform.is_mps():
             self.use_fsdp_inference = False

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -38,6 +38,7 @@ from huggingface_hub.errors import (
 )
 from requests.exceptions import ConnectionError as RequestsConnectionError
 from requests.exceptions import RequestException
+from safetensors import safe_open
 from transformers import AutoConfig, PretrainedConfig
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 
@@ -836,3 +837,12 @@ def snapshot_download(
         }
         hf_kwargs.update(kwargs)
         return _hf_snapshot_download(**hf_kwargs)
+
+
+def get_metadata_from_safetensors_file(file_path: str):
+    try:
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            metadata = f.metadata()
+            return metadata
+    except Exception as e:
+        logger.warning(e)

python/sglang/multimodal_gen/runtime/utils/logging_utils.py

@@ -375,6 +375,12 @@ def configure_logger(server_args, prefix: str = ""):
     set_uvicorn_logging_configs()
 
 
+@lru_cache(maxsize=1)
+def get_log_level() -> int:
+    root = logging.getLogger()
+    return root.level
+
+
 def suppress_loggers(loggers_to_suppress: list[str], level: int = logging.WARNING):
     original_levels = {}