Search papers, labs, and topics across Lattice.
This paper introduces KroQuant, a novel post-training quantization method for diffusion transformers that employs a learned Kronecker-structured invertible transform to optimize the quantization of activations while maintaining output quality. By leveraging a block-local structure, KroQuant significantly reduces the computational burden associated with traditional methods, achieving up to 14% faster inference on MI350 GPUs compared to existing approaches like SmoothQuant. The results demonstrate that KroQuant not only preserves image quality but also outperforms other quantization methods, producing outputs closer to full precision references on various datasets.
KroQuant achieves superior output quality in post-training quantization of diffusion transformers while being up to 14% faster than conventional methods.
Post-training quantization (PTQ) of diffusion transformers (DiTs) to W4A4 severely degrades output quality, because activations entering each linear layer contain outliers that 4-bit formats cannot represent. The standard fix applies an invertible linear transform to the activations and its inverse to the weights before quantizing both. Normalization layers between blocks force this transform to run online at every denoising step, making its inference computation cost the binding design constraint. Existing options trade quantization quality for inference cost: per-channel scaling (SmoothQuant) is computationally cheap but impacts the magnitude of the channels, which can harm quantization accuracy; fixed Hadamard transforms yield better quantization accuracy but require large block sizes that incur a high online cost; learned full-$d$ invertible transforms calibrate best but entail a prohibitive dense $d \times d$ matrix multiplication (GEMM) per layer per step. We propose KroQuant, a PTQ method that applies a learned Kronecker-structured invertible transform to each 32-element block of the activation, storing less than half the parameters of per-channel scaling. The block-local structure runs as small tensor-core GEMMs, and on an MI350 GPU the KroQuant quantizer kernel is up to $14\%$ faster than the SmoothQuant kernel. Offline LoRaQ weight calibration then absorbs the residual per-weight quantization error. On PixArt-$\Sigma$, SANA, and FLUX.1-schnell at W4A4 (MXFP4e2), KroQuant produces outputs closer to the FP reference than SVDQuant and LoRaQ on MJHQ-30K and SDCI, while preserving or improving image quality.
