Bringing Nunchaku 4-bit Diffusion Inference to Diffusers

Massive diffusion transformers can create gorgeous photographs (and even movies, audio snippets, and now textual content), however loading a contemporary text-to-image mannequin in BF16 precision typically requires 20-30 GB of VRAM, which places these fashions out of attain of most client GPUs. Quantization is a robust resolution to this downside, and Diffusers already integrates a number of quantization backends reminiscent of bitsandbytes, GGUF, torchao, and Quanto, which we lined in Exploring Quantization Backends in Diffusers.

Most of those backends are weight-only. Because of this they retailer the weights in low precision and dequantize them again to excessive precision at compute time. This reduces reminiscence utilization considerably, nevertheless it normally doesn’t make inference sooner, and might even add a small latency overhead.

SVDQuant, the quantization methodology behind the favored Nunchaku inference engine, takes a distinct strategy. It runs the principle transformer layers with 4-bit weights and activations (W4A4), decreasing reminiscence whereas additionally dashing up the denoising loop. The small print are lined under, however till now, utilizing these checkpoints required a separate inference library.

With present Diffusers, loading a Nunchaku checkpoint is so simple as calling from_pretrained(), with no native CUDA compilation required due to the kernels package deal. As well as, the companion diffuse-compressor toolkit helps you to quantize new architectures your self and publish them as common Diffusers repositories.

Desk of Contents

Getting began with Nunchaku Lite

First, set up the necessities. You want a latest model of Diffusers and the Hugging Face kernels package deal:

pip set up -U diffusers transformers speed up kernels bitsandbytes

Then load a pre-quantized pipeline like every other Diffusers mannequin:

import torch
from diffusers import ErnieImagePipeline

pipe = ErnieImagePipeline.from_pretrained(
“lite-infer/ERNIE-Picture-Turbo-nunchaku-lite-nvfp4_r32-bnb4-text-encoder”,
torch_dtype=torch.bfloat16,
).to(“cuda”)

picture = pipe(
immediate=“A cinematic portrait of a crimson fox in a misty forest at dawn, “
“detailed fur, volumetric mild”,
top=1024,
width=1024,
num_inference_steps=8,
guidance_scale=1.0,
generator=torch.Generator(“cuda”).manual_seed(42),
).photographs[0]
picture.save(“output.png”)

No customized pipeline class or separate inference engine is required, and there may be nothing to compile domestically. The NVFP4 kernels are downloaded from the Hub via the Nunchaku Lite kernels web page the primary time they’re used. This checkpoint pairs a Nunchaku NVFP4 transformer with a bitsandbytes NF4 textual content encoder, and generates a 1024×1024 picture in about 1.7 seconds on an RTX 5090 with a peak reminiscence utilization of about 12 GB, in contrast with about 24 GB for the BF16 pipeline. You will discover extra particulars concerning the Nunchaku Lite checkpoint format within the official Diffusers documentation.

NVFP4 checkpoints require an NVIDIA Blackwell GPU (RTX 50 sequence, RTX PRO 6000, B200). For earlier generations, use the INT4 variants. See the {hardware} help desk under for particulars.

Background: SVDQuant and Nunchaku

SVDQuant is the quantization methodology behind Nunchaku, its reference CUDA inference engine. Normal 4-bit quantization is troublesome for diffusion transformers as a result of each weights and activations comprise massive outliers. SVDQuant handles this by transferring activation outliers into the weights, representing the toughest a part of every weight matrix with a small 16-bit low-rank department, and quantizing the remaining residual to 4 bits. Nunchaku makes this quick with fused kernels for the 4-bit path and the low-rank department.

Introducing Nunchaku Lite

The unique Nunchaku engine will get a lot of its pace from model-specific fused execution paths, reminiscent of fused QKV projections and fused GELU/MLP kernels. These optimizations are tied to every structure’s module format and checkpoint format, so supporting a brand new mannequin household normally requires model-specific integration work.

Nunchaku Lite is the brand new integration path in Diffusers. With it, Diffusers can load Nunchaku-style checkpoints with no customized pipeline or a separate inference engine. Below the hood, Nunchaku Lite patches the related nn.Linear modules of a inventory Diffusers mannequin with runtime SVDQ/AWQ linear layers earlier than the checkpoint is loaded. The CUDA kernels come from the Hub via the kernels package deal. Two kernel households are used:

svdq_w4a4: 4-bit weights and activations with the SVDQuant low-rank correction. This layer is used for the transformer’s consideration and MLP projections, the place practically the entire compute is spent, and is accessible in INT4 and NVFP4 variants.
awq_w4a16: 4-bit weights with 16-bit activations, used for adaptive normalization and modulation projections reminiscent of FLUX adanorm_single / adanorm_zero or Qwen-Picture modulation layers. These layers are memory-bound and precision-sensitive, making AWQ an excellent match to protect precision whereas nonetheless saving reminiscence and house.

The trade-off is that, with out architecture-specific fused kernels and modules, Nunchaku Lite can’t match the speedup of the unique Nunchaku engine. Nonetheless, the bare-bones implementation nonetheless delivers round 30% speedup whereas retaining the identical stage of VRAM discount.

Native loading in Diffusers

When you have used bitsandbytes or torchao in Diffusers, the mechanics will really feel acquainted. A Nunchaku Lite mannequin repository is an peculiar Diffusers repository. The one particular half is a quantization_config block contained in the transformer’s config.json:

“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”,
“layers.0.self_attention.to_k”,
“…”
]
},
“awq_w4a16”: {
“precision”: “int4”,
“group_size”: 64,
“targets”: [
“adaLN_modulation.1”,
“…”
]
}
}

This config tells Diffusers which modules have been quantized, which scheme they use, and which Nunchaku Lite runtime layer to instantiate (SVDQW4A4Linear or AWQW4A16Linear).

As a result of the quantized mannequin retains the precise module construction of the dense one, the whole lot downstream (schedulers, LoRA loading hooks, offloading, torch.compile) sees a traditional Diffusers mannequin.

{Hardware} help

Nunchaku Lite makes use of completely different kernel variants relying on the GPU era and checkpoint precision:

Volta and Hopper GPUs are presently not supported by the 4-bit kernels. The quantizer validates the GPU’s CUDA functionality at load time and raises a transparent error as an alternative of manufacturing incorrect outputs.

Getting extra pace and decrease reminiscence

Nunchaku Lite could be mixed with different Diffusers reminiscence and pace optimizations.

torch.compile. Compiling the transformer improves the end-to-end speedup from 1.35x to 1.8x:

pipe.transformer.compile(fullgraph=True)

pipe.transformer.compile_repeated_blocks(fullgraph=True)

Quantized textual content encoders. The transformer will not be the one part with a big reminiscence footprint. Textual content encoders reminiscent of T5 or Qwen3 can occupy a number of gigabytes on their very own. Additional quantizing the textual content encoder with bitsandbytes NF4 reduces peak VRAM by about 22% in our benchmark.

Offloading. Diffusers offloading helpers reminiscent of enable_model_cpu_offload() and enable_sequential_cpu_offload() work as ordinary if it’s essential match the pipeline onto a smaller GPU.

Benchmarks

All numbers under have been measured on an NVIDIA RTX PRO 6000 (Blackwell) at 1024×1024 utilizing rootonchair/ERNIE-Picture-Turbo-nunchaku-lite-int4-bnb4-text-encoder.

Finish-to-end latency and reminiscence

As proven above, Nunchaku reduces peak VRAM by as much as 50% whereas nonetheless bettering latency by roughly 30%. The remaining overhead comes largely from additional kernel launches, which torch.compile can mitigate, bringing the complete pipeline right down to 1.68 s, or 1.8x sooner than the BF16 baseline.

Picture high quality

Quantizing your individual mannequin

Nunchaku Lite help in Diffusers is architecture-agnostic, and the diffuse-compressor toolkit offers an end-to-end SVDQuant workflow for Diffusers fashions: calibrate, quantize, package deal, and publish.

Beneath, we stroll via quantizing FLUX.2 Klein 4B for instance. It covers the principle steps: examine the mannequin, calibrate and quantize the transformer, package deal the end result as a Diffusers pipeline, then confirm and push it to the Hub. The complete tutorial covers each flag intimately.

1. Examine what will probably be quantized

The generic scanner walks the mannequin and decides what to focus on: appropriate linears contained in the repeated transformer-block stack turn out to be SVDQ W4A4 targets, acknowledged modulation linears turn out to be AWQ W4A16 targets, and the whole lot else stays dense.

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B
–precision int4 –rank 32 –inspect-config

All the time learn this report earlier than quantizing. For FLUX.2 Klein 4B, the anticipated result’s 100 SVDQ targets, 3 AWQ targets, and 6 dense outer linears, with no lacking patterns or duplicate names.

2. Run quantization

The next command runs SVDQuant on the transformer and writes the quantized checkpoint to outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors:

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B
–precision int4
–output outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors

Substitute –precision int4 with nvfp4 to construct Blackwell-native weights.

3. Package deal a Diffusers pipeline

The converter combines the quantized transformer with the bottom pipeline’s different elements, writes the compact nunchaku_lite configuration into transformer/config.json, and might optionally convert textual content encoders to NF4:

python examples/convert_nunchaku_lite_diffusers.py
–checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors
–model-id black-forest-labs/FLUX.2-klein-4B
–bnb4-text-encoder text_encoder
–compute-dtype bfloat16
–output-dir outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder

4. Load, confirm, and push to the Hub

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
“outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder”,
device_map=“cuda”,
)
picture = pipe(
“A glass robotic in a greenhouse, cinematic lighting”,
num_inference_steps=4, guidance_scale=1.0,
generator=torch.Generator(“cuda”).manual_seed(12345),
).photographs[0]

As soon as the outputs look good, run pipe.push_to_hub(“your-name/your-model-nunchaku-lite-int4”). Different customers can then load it with the identical from_pretrained() sample proven above.

Quantizing fashions with structural rewrites

Be aware that the generic path assumes the structure could be quantized with out structural rewrites. For added speedup, the unique Nunchaku engine rewrites teams of Diffusers layers as fused modules. The generic path can’t infer these adjustments by itself, reminiscent of combining separate Q, Ok, and V projections into one module or splitting a fused projection throughout a number of modules.

FLUX.1-dev’s QKV projection is a concrete instance. Diffusers defines three separate modules:

self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)

The Nunchaku FLUX module combines these layers into one quantized to_qkv module:

to_qkv = fuse_linears([other.to_q, other.to_k, other.to_v])
self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs)

This grouped module is required as a result of Nunchaku’s fused operator consumes the QKV projection, Q/Ok normalization, and rotary embeddings collectively. By comparability, the default Diffusers path executes them individually:

question = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
worth = attn.to_v(hidden_states)

question = question.unflatten(-1, (attn.heads, –1))
key = key.unflatten(-1, (attn.heads, –1))
worth = worth.unflatten(-1, (attn.heads, –1))

question = attn.norm_q(question)
key = attn.norm_k(key)

if image_rotary_emb is not None:
question = apply_rotary_emb(question, image_rotary_emb, sequence_dim=1)
key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)

The Nunchaku path provides the grouped projection, normalization modules, and rotary embeddings to at least one fused operator:

qkv = fused_qkv_norm_rottary(
hidden_states, attn.to_qkv, attn.norm_q, attn.norm_k, image_rotary_emb
)

That is the structural rewrite that the generic path can’t infer. Diffusers has three vacation spot modules with to_q, to_k, and to_v parameter prefixes, whereas Nunchaku has one grouped module beneath to_qkv. A model-specific goal config or adapter should state that the Q, Ok, and V parameters needs to be concatenated alongside the output dimension, in that order, and loaded into to_qkv.

Structural rewrites like these are described by a model-specific goal config throughout quantization and dealt with by a small runtime adapter when the checkpoint is loaded.
The FLUX.2 Klein 4B quantization script offers a concrete target-config instance for producing a structurally rewritten checkpoint, whereas rootonchair/nunchaku-lite offers the runtime adapters wanted to load grouped QKV tensors, cut up fused projections, and different fused operations.
For the whole workflow, you’ll be able to verify the Including A New Mannequin information.

Prepared-to-use checkpoints

To get began instantly, take a look at the next repositories:

Conclusion

Nunchaku’s SVDQuant kernels are some of the efficient methods to run diffusion transformers effectively on client {hardware}, and they’re now natively supported in Diffusers. Pre-quantized checkpoints load with from_pretrained(), and the diffuse-compressor toolkit makes it attainable to quantize new architectures with out ready for engine help. By quantizing each weights and activations, the W4A4 path lowers reminiscence use whereas bettering denoising latency, protecting picture high quality near the BF16 authentic.

If you happen to quantize and publish a brand new mannequin, we might love to listen to about it. Share it on the Hub and tell us! When you have any questions on this characteristic, be happy to affix our Discord.

To study extra, take a look at the next assets:

Acknowledgements

Due to the Diffusers maintainers for evaluations and steering all through the combination, and to the MIT HAN Lab / Nunchaku staff for the unique SVDQuant work. Due to Marc Solar for offering suggestions on the weblog publish. Due to Álvaro Somoza for attempting out nunchaku-lite and for offering suggestions.

rootonchair can be grateful to SilverAI for supporting this work and offering the setting through which a lot of this improvement passed off.