diff --git a/docs/source/en/quantization/nunchaku.md b/docs/source/en/quantization/nunchaku.md index ceb663336ef6..825eecdae810 100644 --- a/docs/source/en/quantization/nunchaku.md +++ b/docs/source/en/quantization/nunchaku.md @@ -122,6 +122,21 @@ List each module you want to quantize under `svdq_w4a4` or `awq_w4a16`. A module } ``` +## Fused kernels + +The original [Nunchaku](https://github.com/nunchaku-ai/nunchaku) engine gets much of its speed from model-specific fused execution paths. It combines the Q, K, and V projections with RMSNorm and RoPE, and uses a fused GELU kernel for the MLP. Nunchaku Lite instead uses the standard Diffusers model with generic quantized linear layers, so it does not include these fusions. + +The following measurements show the per-step latency impact of these fusions in the original Nunchaku engine, measured with Flux Schnell on an RTX 5090. + +| Optimization | Latency with it OFF | Latency with it ON | Speedup | +| --- | --- | --- | --- | +| QKV+norm+RoPE fusion (combined) | 1.057 s | 0.834 s | **1.27×** | +| — grouping Q/K/V into one matmul, alone | 1.059 s | 1.034 s | 1.025× | +| — RMSNorm+RoPE epilogue fusion, on top of grouping | 1.034 s | 0.836 s | 1.236× | +| GELU-MLP fusion | 0.766 s | 0.722 s | **1.06×** | + +Without these fusions, the GPU launches the projection, normalization, rotary-embedding, and MLP operations as separate kernels, which adds launch overhead. `torch.compile` may reduce some of this overhead, although it does not necessarily reproduce the original engine's fused kernels. Benchmark the compiled pipeline for your model and workload. + ## torch.compile Nunchaku Lite kernels and quantized linear layers are compatible with [`torch.compile`](../optimization/fp16#torchcompile).