Add Nunchaku Lite single-file quantization

docs/source/en/_toctree.yml

@@ -174,6 +174,8 @@
     title: bitsandbytes
   - local: quantization/gguf
     title: gguf
+  - local: quantization/nunchaku
+    title: Nunchaku Lite
   - local: quantization/torchao
     title: torchao
   - local: quantization/quanto

docs/source/en/api/quantization.md

@@ -30,6 +30,10 @@ Quantization techniques reduce memory and computational costs by representing we
 
 [[autodoc]] quantizers.quantization_config.GGUFQuantizationConfig
 
+## NunchakuLiteQuantizationConfig
+
+[[autodoc]] quantizers.quantization_config.NunchakuLiteQuantizationConfig
+
 ## QuantoConfig
 
 [[autodoc]] quantizers.quantization_config.QuantoConfig

docs/source/en/quantization/nunchaku.md

@@ -0,0 +1,146 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+
+-->
+
+# Nunchaku Lite
+
+Nunchaku Lite is a quantization backend for loading prequantized checkpoints in Diffusers. Use
+[`diffuse-compressor`](https://github.com/rootonchair/diffuse-compressor) to create compatible checkpoints by choosing
+or adapting a target config for the model architecture, quantizing and exporting the transformer, and packaging the
+result as a Diffusers pipeline repository.
+
+Nunchaku Lite builds on the original [Nunchaku](https://github.com/nunchaku-ai/nunchaku) inference engine,
+[DeepCompressor](https://github.com/nunchaku-ai/deepcompressor) quantization library, and
+[SVDQuant paper](https://arxiv.org/abs/2411.05007).
+
+## Load a quantized pipeline
+
+Load the packaged prequantized pipeline with [`~DiffusionPipeline.from_pretrained`]. Diffusers reads the quantization
+config from `config.json`.
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+
+model_id = "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4"
+
+pipe = DiffusionPipeline.from_pretrained(
+    model_id, torch_dtype=torch.bfloat16,
+).to("cuda")
+
+prompt = "A modern red armchair in a quiet studio, soft window light, realistic product photography"
+image = pipe(
+    prompt=prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=8,
+    guidance_scale=1.0,
+    use_pe=False,
+).images[0]
+image.save("ernie-image-turbo-nunchaku-lite.png")
+```
+
+## Using Optimized CUDA Kernels with Nunchaku Lite
+
+Nunchaku Lite uses optimized CUDA kernels through the `kernels` package. The kernels are loaded automatically when you
+load a Nunchaku Lite checkpoint. Install the runtime dependency before loading a Nunchaku Lite checkpoint.
+
+```bash
+pip install -U kernels
+```
+
+> [!NOTE]
+> The exported state dict must match the target Diffusers model architecture exactly. For example, a checkpoint
+> quantized with fused QKV projections won't load into a model config that expects separate Q, K, and V projection
+> modules.
+
+## Supported Quantization Types
+
+Nunchaku Lite supports the following quantized linear layer formats.
+
+> [!TIP]
+> Use `nvfp4` on Blackwell GPUs. Running `int4` checkpoints on Blackwell can be slower than `nvfp4`.
+
+The CUDA kernels currently support the following NVIDIA GPU architectures:
+
+- `sm_75` (Turing, for example RTX 2080)
+- `sm_80` (Ampere, for example A100)
+- `sm_86` (Ampere, for example RTX 3090 and RTX A6000)
+- `sm_89` (Ada, for example RTX 4090)
+- `sm_120` (Blackwell, for example RTX 5090)
+
+> [!NOTE]
+> Hopper GPUs, such as `sm_90` H100 and H200, are not currently supported.
+
+`nvfp4` checkpoints require a Blackwell or newer NVIDIA GPU. On Blackwell GPUs, use PyTorch >= 2.7 with CUDA >= 12.8.
+`int4` checkpoints require a Turing or newer NVIDIA GPU.
+
+| Method | Precision | Group size | Notes |
+|---|---:|---:|---|
+| `svdq_w4a4` | `nvfp4` | 16 | Uses NVFP4 runtime kernels with SVDQ low-rank correction. |
+| `svdq_w4a4` | `int4` | 64 | Uses INT4 W4A4 kernels with SVDQ low-rank correction. |
+| `awq_w4a16` | `int4` | 64 | Uses INT4 weight-only AWQ-style kernels. |
+
+## NunchakuLiteQuantizationConfig
+
+The `config.json` file must include a compact [`NunchakuLiteQuantizationConfig`]. It defines the runtime
+`compute_dtype` and the target modules for each Nunchaku Lite quantization method.
+
+- `compute_dtype`: runtime dtype for floating-point buffers in quantized modules, typically `torch.bfloat16`.
+- `svdq_w4a4`: SVDQ W4A4 target config with `precision`, `group_size`, `rank`, and `targets`.
+- `awq_w4a16`: AWQ W4A16 target config with `precision`, `group_size`, and `targets`.
+
+Each entry in `targets` must point to a linear layer. Diffusers replaces the linear layers listed under `svdq_w4a4`
+with SVDQ W4A4 layers and the linear layers listed under `awq_w4a16` with AWQ W4A16 layers. The example below shows the
+expected shape with shortened target lists.
+
+> [!NOTE]
+> A target module can only use one Nunchaku Lite quantization method. Do not list the same target under both
+> `svdq_w4a4` and `awq_w4a16`.
+
+```json
+{
+  "_class_name": "ErnieImageTransformer2DModel",
+  "quantization_config": {
+    "quant_method": "nunchaku_lite",
+    "compute_dtype": "bfloat16",
+    "svdq_w4a4": {
+      "precision": "nvfp4",
+      "group_size": 16,
+      "rank": 32,
+      "targets": ["layers.0.self_attention.to_q"]
+    },
+    "awq_w4a16": {
+      "precision": "int4",
+      "group_size": 64,
+      "targets": ["final_linear"]
+    }
+  }
+}
+```
+
+## torch.compile
+
+Nunchaku Lite kernels and quantized linear layers are compatible with [`torch.compile`](../optimization/fp16#torchcompile).
+Compile the quantized transformer after loading the pipeline for faster inference.
+
+```python
+pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=True)
+```
+
+An ERNIE-Image-Turbo benchmark on an RTX PRO 6000 reported that Nunchaku Lite NVFP4 with `torch.compile` reduced the full pipeline latency from 2.271s to 1.675s. Compared to the original BF16 pipeline, the compiled
+Nunchaku Lite NVFP4 pipeline reached a 1.8x speedup.
+
+## Resources
+
+- [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor)
+- [Nunchaku installation requirements](https://nunchaku.tech/docs/nunchaku/installation/installation.html)

docs/source/en/quantization/overview.md

@@ -28,7 +28,7 @@ There are two ways to use [`~quantizers.PipelineQuantizationConfig`] depending o
 
 Initialize [`~quantizers.PipelineQuantizationConfig`] with the following parameters.
 
-- `quant_backend` specifies which quantization backend to use. Currently supported backends include: `bitsandbytes_4bit`, `bitsandbytes_8bit`, `gguf`, `quanto`, and `torchao`.
+- `quant_backend` specifies which quantization backend to use. Currently supported backends include: `bitsandbytes_4bit`, `bitsandbytes_8bit`, `gguf`, `nunchaku_lite`, `quanto`, and `torchao`.
 - `quant_kwargs` specifies the quantization arguments to use.
 
 > [!TIP]

src/diffusers/__init__.py

@@ -124,6 +124,18 @@
 else:
     _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig")
 
+try:
+    if not is_torch_available():
+        raise OptionalDependencyNotAvailable()
+except OptionalDependencyNotAvailable:
+    from .utils import dummy_nunchaku_lite_objects
+
+    _import_structure["utils.dummy_nunchaku_lite_objects"] = [
+        name for name in dir(dummy_nunchaku_lite_objects) if not name.startswith("_")
+    ]
+else:
+    _import_structure["quantizers.quantization_config"].append("NunchakuLiteQuantizationConfig")
+
 try:
     if not is_auto_round_available():
         raise OptionalDependencyNotAvailable()
@@ -1012,6 +1024,14 @@
     else:
         from .quantizers.quantization_config import NVIDIAModelOptConfig
 
+    try:
+        if not is_torch_available():
+            raise OptionalDependencyNotAvailable()
+    except OptionalDependencyNotAvailable:
+        from .utils.dummy_nunchaku_lite_objects import *
+    else:
+        from .quantizers.quantization_config import NunchakuLiteQuantizationConfig
+
     try:
         if not is_auto_round_available():
             raise OptionalDependencyNotAvailable()

src/diffusers/quantizers/auto.py

@@ -22,10 +22,12 @@
 from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer
 from .gguf import GGUFQuantizer
 from .modelopt import NVIDIAModelOptQuantizer
+from .nunchaku import NunchakuLiteQuantizer
 from .quantization_config import (
     AutoRoundConfig,
     BitsAndBytesConfig,
     GGUFQuantizationConfig,
+    NunchakuLiteQuantizationConfig,
     NVIDIAModelOptConfig,
     QuantizationConfigMixin,
     QuantizationMethod,
@@ -44,6 +46,7 @@
     "torchao": TorchAoHfQuantizer,
     "modelopt": NVIDIAModelOptQuantizer,
     "auto-round": AutoRoundQuantizer,
+    "nunchaku_lite": NunchakuLiteQuantizer,
 }
 
 AUTO_QUANTIZATION_CONFIG_MAPPING = {
@@ -54,6 +57,7 @@
     "torchao": TorchAoConfig,
     "modelopt": NVIDIAModelOptConfig,
     "auto-round": AutoRoundConfig,
+    "nunchaku_lite": NunchakuLiteQuantizationConfig,
 }
 
 

src/diffusers/quantizers/nunchaku/__init__.py

@@ -0,0 +1 @@
+from .nunchaku_quantizer import NunchakuLiteQuantizer

src/diffusers/quantizers/nunchaku/nunchaku_quantizer.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from ..base import DiffusersQuantizer
+
+
+if TYPE_CHECKING:
+    from ...models.modeling_utils import ModelMixin
+
+
+from ...utils import is_kernels_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class NunchakuLiteQuantizer(DiffusersQuantizer):
+    def __init__(self, quantization_config, **kwargs):
+        super().__init__(quantization_config, **kwargs)
+        self.compute_dtype = quantization_config.compute_dtype
+        self.pre_quantized = quantization_config.pre_quantized
+
+    def validate_environment(self, *args, **kwargs):
+        if not is_kernels_available():
+            raise ImportError(
+                "Loading Nunchaku checkpoints requires the Hugging Face `kernels` package. "
+                "Install it with `pip install kernels`."
+            )
+        import torch
+
+        cuda_available = torch.cuda.is_available()
+        if not cuda_available:
+            raise ValueError("Loading Nunchaku checkpoints requires a CUDA-capable NVIDIA GPU.")
+
+        device_capability = torch.cuda.get_device_capability()
+
+        if device_capability[0] == 9:
+            raise ValueError("Loading Nunchaku checkpoints is not supported on Hopper NVIDIA GPUs.")
+
+        has_nvfp4_config = (
+            self.quantization_config.svdq_w4a4 is not None
+            and self.quantization_config.svdq_w4a4["precision"] == "nvfp4"
+        )
+        has_int4_config = any(
+            config is not None and config["precision"] == "int4"
+            for config in (self.quantization_config.svdq_w4a4, self.quantization_config.awq_w4a16)
+        )
+        if has_nvfp4_config and device_capability < (10, 0):
+            raise ValueError("Loading Nunchaku NVFP4 checkpoints requires a Blackwell or newer NVIDIA GPU.")
+        if has_int4_config and device_capability < (7, 5):
+            raise ValueError("Loading Nunchaku INT4 checkpoints on CUDA requires a Turing or newer NVIDIA GPU.")
+
+    def update_torch_dtype(self, torch_dtype):
+        if torch_dtype is None:
+            torch_dtype = self.compute_dtype
+        else:
+            self.compute_dtype = torch_dtype
+        return torch_dtype
+
+    def _process_model_before_weight_loading(
+        self,
+        model: "ModelMixin",
+        state_dict: dict[str, Any] | None = None,
+        **kwargs,
+    ):
+        from .utils import check_strict_state_dict_match, replace_with_nunchaku_linear
+
+        quantization_config = self.quantization_config.to_dict()
+        num_replaced = replace_with_nunchaku_linear(model, quantization_config, self.compute_dtype)
+
+        if state_dict is not None:
+            check_strict_state_dict_match(model, state_dict)
+        logger.info(f"Applied Nunchaku quantization config with {num_replaced} targets.")
+
+    def _process_model_after_weight_loading(self, model: "ModelMixin", **kwargs):
+        return model
+
+    @property
+    def is_serializable(self):
+        return False
+
+    @property
+    def is_trainable(self) -> bool:
+        return False
+
+    @property
+    def is_compileable(self) -> bool:
+        return True