用混合精度量化加速扩散模型推理，速度提升超5倍且不丢质量
MixDiT: Accelerating Image Diffusion Transformer Inference with Mixed-Precision MX Quantization
- 对扩散模型激活值按重要性分层量化，关键部分保留高精度
- 在RTX 3090上实现2.10到5.32倍加速，FID无损失
- 软硬件协同设计，支持灵活精度计算与转换
扩散变换器（DiT）在图像生成任务中取得显著进展，但其推理过程计算密集，即使在数据中心级GPU上仍存在长延迟，主要源于迭代特性及编码器结构带来的大量GEMM操作。现有量化方法难以在保证高精度的同时实现显著加速。为此，本文提出MixDiT，一种算法-硬件协同的加速方案，利用混合微缩放（MX）格式对DiT激活值进行量化。MixDiT通过选择性地为基于幅度的异常值保留更高精度，实现混合精度GEMM运算。为充分发挥混合精度算术优势，我们设计了支持精度可变乘法和高效MX精度转换的MixDiT加速器。实验结果表明，MixDiT在RTX 3090上相比基线获得2.10至5.32倍加速，且无FID性能损失。
原文摘要 · Abstract (English)
Diffusion Transformer (DiT) has driven significant progress in image generation tasks. However, DiT inferencing is notoriously compute-intensive and incurs long latency even on datacenter-scale GPUs, primarily due to its iterative nature and heavy reliance on GEMM operations inherent to its encoder-based structure. To address the challenge, prior work has explored quantization, but achieving low-precision quantization for DiT inferencing with both high accuracy and substantial speedup remains an open problem. To this end, this paper proposes MixDiT, an algorithm-hardware co-designed acceleration solution that exploits mixed Microscaling (MX) formats to quantize DiT activation values. MixDiT quantizes the DiT activation tensors by selectively applying higher precision to magnitude-based outliers, which produce mixed-precision GEMM operations. To achieve tangible speedup from the mixed-precision arithmetic, we design a MixDiT accelerator that enables precision-flexible multiplications and efficient MX precision conversions. Our experimental results show that MixDiT delivers a speedup of 2.10-5.32 times over RTX 3090, with no loss in FID.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。