Nunchaku Lite

Nunchaku Lite 是用于在 Diffusers 中加载预量化检查点的量化后端。请使用 diffuse-compressor 创建兼容的检查点。它会量化并导出 transformer，然后将其打包为 Diffusers pipeline。

Nunchaku Lite 基于原始的 Nunchaku 推理引擎、DeepCompressor 量化库以及 SVDQuant 论文 构建。

安装 CUDA 内核

kernels 包提供经过优化的 CUDA 内核，它们会自动加载。请先安装该包。

pip install -U kernels

Nunchaku Lite 会从 rootonchair/nunchaku-lite-kernels 仓库加载其内核，而该仓库的发布者并非 Hub 上受信任的内核发布者。加载该内核会从 Hub 下载并执行代码，因此 Diffusers 要求您通过设置 DIFFUSERS_TRUST_REMOTE_KERNELS=true 来明确选择加入。有关详细信息，请参阅 信任远程内核。

加载量化流水线

使用 from_pretrained() 加载预量化流水线，该函数会从 config.json 中读取量化配置。

import torch
from diffusers import DiffusionPipeline

model_id = "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4"

pipe = DiffusionPipeline.from_pretrained(
    model_id, dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

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,
).images[0]
image.save("ernie-image-turbo-nunchaku-lite.png")

导出的 state dict 必须与目标 Diffusers 模型架构完全匹配。例如，使用融合 QKV 投影进行量化的检查点无法加载到期望使用独立 Q、K 和 V 投影模块的模型配置中。

支持的量化类型

Nunchaku Lite 支持以下量化线性层格式。

在 Blackwell GPU 上使用 nvfp4。在 Blackwell 上运行 int4 检查点可能比 nvfp4 更慢。

目前，CUDA 内核支持以下 NVIDIA GPU 架构：

• sm_75（Turing，例如 RTX 2080）
• sm_80（Ampere，例如 A100）
• sm_86（Ampere，例如 RTX 3090 和 RTX A6000）
• sm_89（例如 RTX 4090 的 Ada）
• sm_120（Blackwell，例如 RTX 5090）

Hopper 架构的 GPU（例如 sm_90 H100 和 H200）目前尚不支持。

nvfp4 检查点需要 Blackwell 或更新版本的 NVIDIA GPU。在 Blackwell GPU 上，请使用 PyTorch >= 2.7 和 CUDA >= 12.8。int4 检查点需要 Turing 或更新版本的 NVIDIA GPU。

NunchakuLiteQuantizationConfig

config.json 文件必须包含一个 NunchakuLiteQuantizationConfig。它定义了运行时 compute_dtype 以及每种 Nunchaku Lite 量化方法的目标模块。

• compute_dtype：量化模块中浮点缓冲区的运行时数据类型，通常为 torch.bfloat16。
• svdq_w4a4：包含 precision、group_size、rank 和 targets 的 SVDQ W4A4 目标配置。
• awq_w4a16：包含 precision、group_size 和 targets 的 AWQ W4A16 目标配置。

targets 中的每个条目都必须指向一个线性层。Diffusers 会将每个 svdq_w4a4 目标替换为 SVDQ W4A4 层，并将每个 awq_w4a16 目标替换为 AWQ W4A16 层。下例展示了预期的形状，其中目标列表已作简化。

在 svdq_w4a4 或 awq_w4a16 下，列出您想要量化的每个模块。一个模块只能使用一种方法，因此请勿将同一目标同时列在两者之下。

{
  "_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"]
    }
  }
}

融合内核

原始的 Nunchaku 引擎的大部分速度优势来自于针对特定模型的融合执行路径。它将 Q、K 和 V 投影与 RMSNorm 及 RoPE 相结合，并在 MLP 中使用融合 GELU 内核。而 Nunchaku Lite 则使用标准的 Diffusers 模型和通用的量化线性层，因此不包含这些融合操作。

以下测量结果展示了在原始 Nunchaku 引擎中，这些融合操作对每步延迟的影响，测量环境为在 RTX 5090 上运行 Flux Schnell。

若未进行这些融合，GPU 会将投影、归一化、旋转嵌入和 MLP 操作作为独立的内核启动，从而增加启动开销。torch.compile 可能会减少部分此类开销，但它不一定能复现原始引擎的融合内核。请针对您的模型和工作负载对编译后的流水线进行基准测试。

torch.compile

Nunchaku Lite 内核和量化线性层与 torch.compile 兼容。 在加载流水线后，对量化后的 transformer 进行编译，以加快推理速度。

pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=True)

编译后的 Nunchaku Lite NVFP4 流水线比原始 BF16 流水线快 1.8 倍（在 RTX PRO 6000 上，从 2.271 秒降至 1.675 秒）。

资源

• diffuse-compressor
• Nunchaku 安装要求