快速导读:KroQuant提出了一种基于Kronecker积参数化的可学习块对角激活变换,结合LoRaQ权重校正,在保持推理速度的同时显著提升了DiT模型W4A4量化的图像质量。
扩散Transformer(DiT)模型的训练后量化(PTQ)面临严峻挑战,尤其是在W4A4(4-bit权重和4-bit激活)设置下。激活值中的Outliers导致量化后图像质量严重下降。传统的激活变换方法,如SmoothQuant或全矩阵旋转,在DiT中因AdaLN层无法离线吸收(Absorption Trick失效)而面临在线计算成本高或表达能力不足的权衡。
KroQuant提出了一种新颖的解决方案:基于Kronecker积参数化的可学习块对角激活变换。该方法结合LoRaQ权重校正,在保持推理速度的同时显著提升了DiT模型W4A4量化的图像质量。本文详细解析KroQuant的核心思想、技术细节、实验结果及局限性,为计算机视觉从业者提供深入理解。
英文题目:KroQuant: Kronecker-Structured Block Transforms for Efficient Post-Training Quantization of Diffusion Transformers
论文出处:arXiv 每日论文精选 · arXiv:2607.21446
原始论文:PDF / 论文页面
对应视频标题:KroQuant:用Kronecker结构破解DiT W4A4量化难题|推荐指数:★★★★★
这篇论文解决什么问题?
在LLM量化中,常用的离线权重吸收技巧依赖于归一化层与变换矩阵的可交换性。然而,DiT中的AdaLN不满足此条件,导致变换必须在每一步在线执行。现有方法中,逐通道缩放(SmoothQuant)成本低但精度有限,而固定Hadamard变换或全矩阵学习变换(SpinQuant/FlatQuant)要么表达能力受限,要么在线计算开销过大。
具体而言,SmoothQuant仅能处理逐通道的缩放,无法捕捉激活值之间的复杂相关性。固定Hadamard变换虽然能引入一定的旋转,但其表达能力有限,难以适应不同模型的激活分布。全矩阵学习变换虽然表达能力强,但其在线计算成本过高,难以在实际部署中应用。
因此,亟需一种既能捕捉激活值复杂相关性,又具有较低在线计算成本的激活变换方法。KroQuant正是在这一背景下提出的。
核心创新
- 提出首个适用于DiT PTQ的非对角在线变换,比固定Hadamard更具表达能力,比全矩阵旋转更轻量。
- 采用Kronecker-LU参数化,将32×32块的参数量从1024压缩至15(68倍压缩),同时保持可逆性。
- 设计融合MXFP4e2量化器的Triton内核,利用Tensor Core执行32×32 GEMM,实现比SmoothQuant快14%的推理速度。
方法概览
KroQuant引入一个可学习的块对角激活变换T,每个32×32块通过5个2×2单位行列式LU因子的Kronecker积参数化(每块仅15个参数)。该变换在线应用于激活值,随后使用LoRaQ进行离线权重校正(将权重分解为4-bit残差和低秩校正分支)。
- KroQuant引入一个可学习的块对角激活变换T,每个32×32块通过5个2×2单位行列式LU因子的Kronecker积参数化。这种参数化方式将每个块的参数量从1024压缩至15(68倍压缩),同时保持可逆性。
- 该变换在线应用于激活值,随后使用LoRaQ进行离线权重校正。LoRaQ将权重分解为4-bit残差和低秩校正分支,从而在量化后恢复部分精度。
- KroQuant设计了融合MXFP4e2量化器的Triton内核,利用Tensor Core执行32×32 GEMM。这种设计使得KroQuant的推理速度比SmoothQuant快14%,同时保持了较高的图像质量。
- KroQuant的块对角结构使得每个块可以独立处理,避免了全矩阵变换中的全局耦合问题。这种局部化处理不仅降低了计算复杂度,还提高了并行效率。
逐图理解论文
直觉与失败

图1展示了W4A4量化线性层的流水线对比。SVDQuant将权重分为量化残差和SVD截断的FP16低秩分支。LoRaQ用数据无关的低秩分支替换SVD截断,该分支可量化至16位以下。KroQuant额外插入可学习的块对角激活变换T,在量化前应用,补充LoRaQ的权重侧校正。P和PSV D是4位残差矩阵,定义为低秩分量的函数。此图清晰展示了KroQuant在流水线中的位置和作用。
实验验证

在PixArt-Σ、SANA和FLUX.1-schnell上,KroQuant显著优于SVDQuant和LoRaQ。例如,在FLUX.1-schnell/SDCI上,KroQuant的LPIPS相对降低17.5%。K2内核比FWHT快1.30-2.45倍,比Scale慢最多25%。
结论与局限

图2展示了PixArt-Σ在MXFP4e2 W4A4下的视觉对比。左三列:FP16(参考)、SVDQuant、KroQuant;KroQuant使用LoRaQ后端进行离线权重校准,而SVDQuant使用其自身的SVD低秩分解。右三列:FP16(参考)、使用SmoothQuant量化器的LoRaQ,以及依赖LoRaQ后端的KroQuant。通过对比,可以看出KroQuant在图像细节和纹理上更接近FP16参考值,优于SVDQuant和LoRaQ。
实验如何设计?
- 实验在PixArt-Σ、SANA和FLUX.1-schnell三个DiT模型上进行,评估W4A4 (MXFP4 E2M1)量化性能。对比基线包括SVDQuant和LoRaQ。
- 使用MJHQ-30K和SDCI数据集,指标包括FID、LPIPS、PSNR和Image Reward。这些指标全面评估了量化后图像的视觉质量和感知相似度。
- 在AMD MI350 GPU上评估内核延迟和HBM带宽,对比Scale(逐通道缩放)、FWHT及KroQuant变体(K1/K2)。这些评估揭示了不同方法在硬件层面的效率差异。
关键结果与论文证据
- 在PixArt-Σ/MJHQ-30K上,KroQuant的LPIPS为0.347,优于LoRaQ的0.381(相对降低8.9%)。这表明KroQuant在捕捉激活值相关性方面具有优势。
- 在SANA/MJHQ-30K上,KroQuant的LPIPS为0.206,优于LoRaQ的0.229(相对降低10.0%)。这一结果进一步验证了KroQuant在不同模型上的泛化能力。
- 在FLUX.1-schnell/SDCI上,KroQuant的LPIPS为0.283,优于LoRaQ的0.343(相对降低17.5%)。这一显著改进表明KroQuant在处理复杂激活分布时的有效性。
- K2内核比FWHT快1.30-2.45倍,比Scale慢最多25%(0.75-0.96x)。这一结果表明KroQuant在保持较高图像质量的同时,实现了较高的推理效率。
阅读时需要注意
- Kronecker结构参数量随块大小增长过慢,难以捕捉GL(n)的全表达能力,因此实际限制在n=32。这意味着对于更大尺寸的激活块,KroQuant可能无法提供足够的表达能力。
- 联合多层优化不稳定,受限于单位行列式LU参数化。这可能导致在某些复杂模型中,KroQuant的优化过程难以收敛。
- KroQuant生成的图像对比度往往高于FP16参考值,理论解释尚不明确。这一现象可能影响图像的视觉自然度,需要进一步研究。
- 在FLUX.1-schnell/MJHQ-30K上,由于双流架构的激活分布异构性,KroQuant在IR和LPIPS上未全面超越基线。这表明KroQuant在处理特定架构的模型时可能存在局限性。
关联工作
- SVDQuant使用SVD分解权重,KroQuant在多数指标上优于或持平。这表明KroQuant的激活变换策略比单纯的权重分解更有效。
- LoRaQ是KroQuant采用的权重校正后端,KroQuant替换了其SmoothQuant前端。这一组合策略充分发挥了LoRaQ的权重校正能力和KroQuant的激活变换优势。
- SmoothQuant是逐通道缩放基线,KroQuant内核比其快14%。这一结果突显了KroQuant在硬件效率上的优势。
- SpinQuant/FlatQuant是全矩阵学习变换,在DiT中因在线成本过高不可行。KroQuant通过块对角结构和Kronecker参数化,成功降低了在线计算成本。
展开:论文全文中文翻译
以下译文用于快速探索和学习,技术术语按需要保留英文;正式引用和精确表述请以原论文为准。
第 1 页
KroQuant:用于高效扩散Transformer训练后量化的Kronecker结构化块变换
Yann Bouquet∗† Alireza Khodamoradi† Kristof Denolf† Mathieu Salzmann∗
摘要
将扩散Transformer(DiTs)进行W4A4精度的训练后量化(PTQ)会严重降低输出质量,因为进入每个线性层的激活值包含4-bit格式无法表示的异常值。标准的修复方法是在量化之前,对激活值应用可逆线性变换,并对权重应用其逆变换。块之间的归一化层迫使该变换在每个去噪步骤中在线运行,使其推理计算成本成为设计约束。现有方案在量化质量与推理成本之间进行权衡:逐通道幅度缩放(SmoothQuant)计算成本低,但会损害量化精度;固定Hadamard变换能产生更好的量化精度,但需要较大的块尺寸,导致较高的在线成本;学习得到的全可逆变换校准效果最佳,但每层每步都涉及昂贵的稠密 $d \times d$ 矩阵乘法(GEMM)。
我们提出了KroQuant,这是一种PTQ方法,它对激活值的每个32元素块应用学习得到的Kronecker结构化可逆变换,存储的参数少于逐通道缩放的一半。块局部结构以小型张量核心GEMM运行,在MI350 GPU上,KroQuant量化器内核比SmoothQuant内核快多达14%。离线LoRaQ权重校准[cs.LG]随后吸收剩余的逐权重量化误差。在PixArt-Σ、SANA和FLUX.1-schnell的W4A4(MXFP4e2)设置下,KroQuant在MJHQ-30K和SDCI上产生的输出比SVDQuant和LoRaQ更接近FP参考,同时保持或提高了图像质量。
1 引言
将精度量化至4位是部署十亿参数扩散Transformer(DiTs)的自然方式,当前加速器通过OCP微缩放标准[Microsoft et al., 2023]原生支持该方式。瓶颈在于激活异常值[Xiao et al., 2023]:一小部分通道的幅度比其他通道大几个数量级,因此4-bit分块量化器将其大部分动态范围浪费在这些通道上,并截断其他所有内容。
考虑线性层 $O = AW^\top$,其中 $A \in \mathbb{R}^{m \times d}$,$W \in \mathbb{R}^{o \times d}$,以及 $b$-bit量化器 $Q_b$。LLM和DiT量化器中的主要解决方案是插入一个可逆激活变换 $T \in GL(d)$,在量化之前重新分配通道幅度,作为arXiv:2607.21446v1
$A \mapsto AT, \quad W \mapsto WT^{-\top}, \quad AT \cdot (WT^{-\top})^\top = AT \cdot T^{-1}W^\top = AW^\top, \quad (1)$
这在保持线性映射不变的同时,减少了每层的量化误差
$E(T) = Q_b(AT) Q_b(WT^{-\top})^\top – AW^\top. \quad (2)$
∗瑞士洛桑EPFL。firstname.lastname@epfl.ch †美国科罗拉多州朗蒙特Advanced Micro Devices, Inc.。firstname.lastname@amd.com
预印本。
第 2 页
基线 16 位 图 1:RANK = 32 时的流水线对比 𝐿!"# 𝑅!"# W4A4 量化线性层 16 位激活 16 位输出 在 DiT 块中。 (a) SVDQuant 平滑量化器 𝑃!"# 将权重分为量化的 4 位 残差和 SVD 截断的 (a) SVDQuant。 FP16 低秩分支。 (b) LoRaQ 用数据无关的低秩分支替换 SVD 截断,该分支本身可以量化到低于 16 位。 基线 4 位 RANK 16 位 = 128 4 位 𝐿 𝑅 平滑量化器 P (c) KroQuant 额外插入 4 位 学习到的块对角激活 (b) LoRaQ。 变换 T,在量化之前, 补充 LoRaQ 的权重 我们 4 位 16 位 4 位 𝐿 RANK = 128 𝑅 方面的修正。P 和 PSV D 是 16 位激活 16 位输出 Kroquant 量化器 P 定义为低秩组 4 位 件的函数。 (c) KroQuant。
补偿变换 T −1 始终折叠到层存储的权重中;吸收技巧进一步将 T 折叠到前一层的输出中,使得到达当前层的激活已经是 T 变换过的,且运行时变换是免费的。因此,现有方法仅在 T 的结构上有所不同:对角逐通道缩放 [Xiao et al., 2023, Zhao et al., 2025, Wu et al., 2024]、密集旋转 [Ashkboos et al., 2025, Tseng et al., 2024, Liu et al., 2025] 或一般可逆矩阵 [Sun et al., 2024]。
虽然通常适用于 LLM,但吸收技巧在 DiT 中失败,特别是在 Transformer 块之间。这是因为将 T 折叠到前一层的输出需要块之间的归一化与 T 可交换: Norm(xT) T −1 = Norm(x). (3) 这并不满足位于每个 DiT 块之间的 AdaLN 变体(附录 C)。因此,DiT 量化器应用的每个变换都在每个去噪步骤的每个线性层上在线运行,且每步成本成为绑定的设计约束。虽然每步成本对于涉及对角矩阵的逐通道缩放变换来说是可承受的,但这种策略仅提供有限的表达能力,从而限制了量化精度。最近的权重分割 PTQ 流水线如 SVDQuant [Li et al., 2025] 和 LoRaQ [Bouquet et al., 2026] 采取了这种权衡,将每个权重分解为低位残差和低秩修正分支。
相比之下,现有的非对角变换产生更好的量化,但在在线约束下表现不佳。Hadamard 变换是一个固定的正交对称矩阵,具有归一化的 ±1 条目,通过加减蝴蝶操作混合通道,将异常值幅度均匀分布在块中 [Ashkboos et al., 2025]。在块大小为 32 时,它很便宜(单次 MXFP4 对齐的内核传递),但仅在每个 32 通道子向量内抑制异常值,最终表现不如逐通道缩放 [Xiao et al., 2023]。学习到的 d × d 可逆矩阵 [Liu et al., 2025, Sun et al., 2024] 表达能力强,但每层存储 O(d2) 参数,并且每一步都需要密集的 m × d × d 矩阵乘法,在推理中不可行。总之,因此需要一个学习到的且为 32 × 32 块对角的变换;即在相同块大小下比固定的 Hadamard 更具表达能力,同时在线成本低且与 MX 量化器兼容。
为了解决这一需求,我们提出了 KroQuant,这是一种学习到的 T,具有 32 × 32 块的块对角结构,并使用紧凑的参数化,每层使用的标量少于逐通道缩放,包含 Hadamard 作为初始化使用的特例,并在推理时作为每个块的单次 32 × 32 张量核心 GEMM 应用;构造细节见第 3.2 节。我们表明,KroQuant 可作为 LoRaQ 等权重分割 PTQ 流水线中激活变换阶段的即插即用替换(图 1)。在 PixArt-Σ [Chen et al., 2024a]、SANA [Xie et al., 2025] 和 FLUX.1-schnell [Black-Forest-Labs, 2024] 上,LoRaQ-with-KroQuant 在 W4A4 下匹配或优于 LoRaQ-with-SmoothQuant [Bouquet et al., 2026, Xiao et al., 2023] 和 SVDQuant 参考 [Li et al., 2025],而 KroQuant 量化器内核在 MI350 GPU 上比 SmoothQuant 内核快多达 14%。
具体而言,我们的贡献如下:
2
第 3 页
• 一种与在线约束兼容的学习型块对角激活变换。这是首个针对 DiT 训练后量化(PTQ)的非对角在线变换,其表达能力强于固定的块 Hadamard 变换,且推理成本低于全维旋转 [Liu et al., 2025] 或可逆变换 [Sun et al., 2024]。
• GL(32) 块的 Kronecker-LU 参数化。五个 2×2 单位行列式 LU 因子产生的参数规模比稠密 GL(32) 块小 68 倍,且稳定性高,Hadamard 变换可作为自然初始化。
• 融合 W4A4 内核。每个块作为单次张量核心通用矩阵乘法(Tensor Core GEMM)应用;在 MI350 上,KroQuant 量化器内核比 SmoothQuant 基线内核快达 14%。
2 相关工作
扩散 Transformer 的训练后量化。DiT 量化已从 8 位方案 [Shang et al., 2023] 进展到亚 8 位方法,这些方法将 SmoothQuant [Xiao et al., 2023] 适配到由自适应层归一化(AdaLN)诱导的时间步条件激活分布。PTQ4DiT [Wu et al., 2024] 和 ViDiT-Q [Zhao et al., 2025] 应用具有时间步依赖校准的逐通道对角缩放;Q-DiT [Chen et al., 2024b] 和 MixDQ [Zhao et al., 2024] 进一步采用混合精度和基于敏感性的分配。所有方法均保持激活变换 T 为对角形式,且在 W4A4 下性能下降,此时残差误差主要由逐通道重缩放无法消除的通道间相关性主导。
激活变换的设计空间。超越逐通道缩放的 PTQ 方法在可逆变换 $T \in GL(d)$ 的稀疏模式上有所不同。SmoothQuant 的对角 T 独立缩放通道;QuaRot [Ashkboos et al., 2025] 和 QuIP# [Tseng et al., 2024] 使用固定的稠密 Hadamard 旋转,其抑制异常值的能力随块大小增加,但其全维蝴蝶结构无法映射到张量核心;SpinQuant [Liu et al., 2025] 和 FlatQuant [Sun et al., 2024] 分别在 Stiefel 流形或 GL(d) 中学习稠密 T,代价是 $d \times d$ 参数矩阵和每层一次的稠密 GEMM。在大型语言模型(LLMs)中,所有这些方法都通过 RMSNorm 离线吸收进相邻权重;如第 1 节所述,在 DiT 中无法做到这一点,因此每步成本成为关键负担。KroQuant 介于 SmoothQuant 和 FlatQuant 之间:T 是块对角的,由 32×32 块组成,每个块在相同块大小下比固定 Hadamard 更具表达能力,并作为并行张量核心 GEMM 在线应用,每层参数少于 SmoothQuant 的对角缩放。
低秩权重分解与激活变换的作用。另一条互补的研究路线将线性层权重 W 分解为低秩分支和 4 位残差。SVDQuant [Li et al., 2025] 平滑异常值并使用数据依赖的低秩校准对平滑后的权重进行 SVD 分解;LoRaQ [Bouquet et al., 2026] 则无数据地拟合低秩分支以匹配残差量化误差,允许低秩分支本身量化至低于 16 位。这两种方法都分解了权重,但按原样量化激活,因此需要一个无归一化以融合的在线激活变换,且两者默认采用 SmoothQuant 风格的对角 T。KroQuant 是该变换的即插即用替代品:我们直接采用 LoRaQ 的权重校正,并将其 SmoothQuant 前端替换为我们学习到的块对角 T,从而得出 LoRaQ-with-KroQuant 与 LoRaQ-with-SmoothQuant 的主要对比,并以 SVDQuant 作为额外参考。
3 方法
我们在第 1 节的正式设置内工作,并按以下方式实例化其中抽象的块量化器 $Q_b$。
3.1 量化设置
我们用 $Q_b$(第 1 节)表示位宽为 b 的块级量化-反量化算子,它接受一个实矩阵并返回同形状的矩阵:输入的每行被划分为大小为 n 的不重叠块;每个块编码为一个共享缩放因子 s(可表示的缩放因子集合,例如 FP8)和一个 n 维向量 q(采用选定的 b 位格式,当 b=4 时为 FP4);反量化后的块重新组装为输出。
3
第 4 页
目标。我们旨在构建一个属于 $GL(d)$ 的变换 $T$,其为块对角形式: $T = \text{diag}(K^{(1)}, K^{(2)}, \dots, K^{(\lceil d/n \rceil)})$, (4)
其中每个对角块 $K^{(j)} \in GL(32)$ 是可学习的(不同于固定的 Hadamard 矩阵),且其参数量小于每通道缩放的一半,且比稠密的 $GL(32)$ 块小 68 倍以上。我们的目标是在保持线性映射不变($A T \cdot (W T^{-1})^\top = A W^\top$)的同时,减小 $\|E\|_F$。
3.2 基于 Kronecker 积的可逆块变换
这种 $T$ 的结构将每组 $n$ 个通道内的特征耦合起来,同时保持组间独立,这与 MXFP4e2 量化器 $Q_4$ 的块结构相匹配。
Kronecker 积构造。我们将每个 $n \times n$ 块参数化为 $\log_2 n$ 个小矩阵的 Kronecker 积: $K_n = G_1 \otimes G_2 \otimes \dots \otimes G_{\log_2 n}, \quad G_k \in \mathbb{R}^{2 \times 2}, \quad k = 1, \dots, \log_2 n$. (5)
对于 $n=32$,这给出了 $\log_2 32 = 5$ 个因子。由于可逆矩阵的 Kronecker 积仍是可逆的,且 $(G_1 \otimes \dots \otimes G_5)^{-1} = G_1^{-1} \otimes \dots \otimes G_5^{-1}$,求逆简化为五个独立的 $2 \times 2$ 求逆运算。
单位行列式 LU 参数化。为了防止退化缩放解(即变换将一个操作数推向零,同时膨胀另一个操作数),我们对每个因子强制 $|\det G_k| = 1$。我们通过具有固定单位行列式的 LU 分解来参数化每个 $2 \times 2$ 矩阵: $G_k = \begin{pmatrix} 1 & 0 \\ c_k & 1 \end{pmatrix} \begin{pmatrix} a_k & b_k \\ 0 & -1/a_k \end{pmatrix}, \quad a_k \neq 0$. (6) $\underbrace{L_k}_{\quad} \underbrace{U_k}_{\quad}$
可以验证 $\det G_k = a_k \cdot (-1/a_k) = -1$,因此对于任何 $a_k \neq 0$,都满足 $|\det G_k| = 1$。上三角因子 $U_k$ 在 $(2, 2)$ 位置使用负倒数,使得参数化将正交反射和旋转作为特例包含在内。每个 $G_k$ 由三个实数参数 $(a_k, b_k, c_k)$ 控制;每个块有五个因子,这意味着每个 $32 \times 32$ 块有 15 个参数,相比之下,无约束的可逆矩阵有 $32^2 = 1024$ 个参数,而标量通道缩放仅有 1 个参数。
Hadamard 初始化。归一化 Hadamard 矩阵 $H_n$ ($n=2^k$) 是 $k$ 个 $2 \times 2$ 因子 $\frac{1}{\sqrt{2}} \begin{pmatrix} 1 & 1 \\ 1 & -1 \end{pmatrix}$ 的 Kronecker 积。它具有行列式 $-1$,因此满足 $|\det H_n| = 1$。我们将每个 $G_k$ 初始化为 $G_k^{(0)} = \begin{pmatrix} \cos \frac{\pi}{4} & \sin \frac{\pi}{4} \\ \sin \frac{\pi}{4} & -\cos \frac{\pi}{4} \end{pmatrix}$,这等效于将 $T$ 初始化为块对角 Hadamard 旋转。经验表明,这是在基于梯度的优化之前减少激活值异常值并降低量化误差的良好初始化策略 [Ashkboos et al., 2025, Liu et al., 2025]。
3.3 逐层变换优化
目标。给定提供 $N$ 个激活样本 $\{A^{(s)}\}_{s=1}^N$ 的小型校准数据集,我们通过最小化输出均方误差来独立优化每个线性层的 $T$: $L_{\text{out}}(T) = \frac{1}{N_{\text{mo}}} \| Q_b(A T) Q_b(W T^{-1})^\top – A W \|_F^2$. (7)
正则化。仅最小化 $L_{\text{out}}$ 可能会收敛到其中一个操作数量化良好而另一个操作数较差的解——例如,完美白化激活值但使权重处于高度非均匀表示的变换。为了鼓励两个操作数都分别被良好量化,我们添加两个辅助项: $L_A(T) = \frac{1}{N_{\text{md}}} \| Q_b(A T) – A T \|_F^2, \quad L_W(T) = \frac{1}{N_{\text{od}}} \| Q_b(W T^{-\top}) – W T^{-\top} \|_F^2$. (8)
总损失为 $L(T) = L_{\text{out}}(T) + \lambda_A L_A(T) + \lambda_W L_W(T)$. (9)
4
第 5 页
λ 的自动缩放。正则化权重 λA 和 λW 被设定为使每个辅助项在初始化时与 Lout 的量级相匹配: Lout(T0) Lout(T0) λA = , λW = (10) LA(T0) LW (T0), 其中 T0 是 Hadamard 初始化。这种归一化确保在训练初期没有任何一项占主导地位,防止正则化器在优化初始阶段覆盖主要的输出误差信号,同时仍提供足够的梯度信号以避免退化解。
优化。我们使用 Adam 针对 T 的 15 · ⌈d/n⌉ 个自由参数优化 L(T)。量化器 Qb 是不可微的;我们使用直通估计器(STE)来通过量化操作传播梯度。由于不同层的变换是独立的,它们可以并行或顺序优化,相对于校准前向传播,其开销可忽略不计。使用 STE 以及可优化参数的 Kronecker 组合使得在块或模型级别上的优化不稳定,这就是我们将优化保持在层级别的原因。
3.4 最终校正:离线 LoRaQ 权重校准
在优化每层 Kronecker 变换 {T (ℓ)} 以最小化我们的目标函数后,我们将 LoRaQ [Bouquet et al., 2026] 中的离线权重校准程序应用于变换后的权重 W (ℓ) = W (ℓ)(T (ℓ))−⊤。f Bouquet et al. [2026] 将每个变换后的权重 W (ℓ) 分解为一个 b 位量化残差分支和一个加性低秩校正分支 L(ℓ)R(ℓ),其中 L(ℓ) ∈Ro×r 且 R(ℓ) ∈Rr×d,并优化这两个分支以实现无数据,将大部分 16 位激活信息携带到低秩分支,从而进一步降低线性层的整体量化误差。低秩分支本身可以被量化到低于 16 位而不破坏分解。在推理时,该层并行计算量化残差矩阵乘法和混合精度低秩校正,并求和它们的输出。
在推理时,第 ℓ 层的计算为
ˆO(ℓ) = Qb A(ℓ)T (ℓ) × Qb W (ℓ) −L(ℓ)R(ℓ) + A(ℓ)T (ℓ)L(ℓ)R(ℓ), (11) f 其中 ˆW (ℓ) 是离线预计算的。
3.5 硬件高效的变换应用
我们将块对角变换 T 实现为两个自定义 Triton 内核,将每个 32 × 32 块 K(j) 应用于每个 token 的相应 32 元素切片。两个内核都在形状为 (⌈m/BR⌉, d/32) 的 2D 网格上对激活矩阵 A ∈Rm×d 进行分块,其中 BR ∈{16, 32, 64, 128} 是自动调优的。每个线程块处理一个形状为 (BR, 32) 的图块。
3.5.1 内核变体
K1(预计算稠密)。每个 32 × 32 块 K(j) 被实例化并存储在 HBM 中。内核加载激活图块 Atile ∈RBR×32 和相应的 K(j) ∈R32×32,计算单个 tl.dot(Atile, Kblock) 并在 fp32 中累加,然后在转换回输入数据类型后存储结果。存储成本:每列块 1024 个值(2 KB);d = 4096 时总计 256 KB。 K2(即时 Kronecker 构造)。K2 不加载稠密矩阵,而是从五个 2 × 2 因子 G1, . . . , G5 的 15 个标量参数重建 K(j)(第 3.2 节)。由于 32 = 25,K(j) 的每个行和列索引都是一个 5 位整数,Kronecker 积的逐元素公式为 5 K[r, c] = Y Gk r5−k, c5−k , (12)
k=1 其中 rp 和 cp 分别表示 r 和 c 的第 p 位。内核将此公式作为五个因子的流式乘积应用,通过位索引在寄存器中逐个实例化 K 因子
5
第 6 页
标量查找,将寄存器预算留给在同一内核中融合下游 MXFP4e2 量化器。伪代码见附录 B。存储成本:每个列块 15 个标量参数,相比 K1 压缩约 68 倍。
所有内核均满足 Triton tl.dot 约束 M, N, K ≥16(对于所有 BR ≥16),因为 GEMM 形状为 (BR, 32) × (32, 32) →(BR, 32)。
K2 每个块存储 15 个标量参数(相比 K1 的 1024 个值,压缩 68 倍),吞吐量成本仅比 K1 高 5–15%(见附录 A)。对于隐藏维度较大且许多列块共享相同 Kronecker 因子的模型,K2 的参数压缩降低了变换矩阵读取路径上的 HBM 压力。在变换存储不是瓶颈的推理场景中,K1 提供更高的绝对吞吐量。
4 实验
4.1 实验设置
模型。我们在三个规模递增且架构多样的扩散 Transformer 上进行评估:
• PixArt-Σ [Chen et al., 2024a],拥有 6 亿参数,用于文本到图像生成的单流架构。 • SANA [Xie et al., 2025],拥有 16 亿参数,用于文本到图像生成的单流架构,其 FFN 模块中包含 1×1 卷积层。 • FLUX.1-schnell [Black-Forest-Labs, 2024],拥有 120 亿参数,用于文本到图像生成的双流架构。
校准数据集。遵循 SVDQuant [Li et al., 2025] 和 LoRaQ [Bouquet et al., 2026] 的方案,我们使用从 COCO 2017 验证集标题 [Chen et al., 2015] 中提取的 128 个文本提示作为校准输入。对于每个提示,对每个线性层执行一次前向传播,以收集用于变换优化和 LoRaQ 权重校准的激活统计信息。
评估数据集。我们在两个不同的数据集上报告图像质量:MJHQ-30K [Li et al., 2024] 和 SDCI [Urbanek et al., 2024],每个数据集随机采样 5,000 张图像,遵循先前工作中使用的评估分割的相同固定种子协议。
指标。遵循现有基准,我们从两个标准评估性能。为了衡量与 16 位基线的相似性,我们使用学习感知图像块相似度 (LPIPS) [Zhang et al., 2018] 和峰值信噪比 (PSNR)。为了评估整体视觉质量,我们使用 Frechet Inception Distance (FID) [Heusel et al., 2018] 和 Image Reward (IR) [Xu et al., 2023]。
数据格式。我们遵循 OCP Microscaling (MX) 标准 [Microsoft et al., 2023] 使用 MXFP4 E2M1 格式,其动机在于当前的硬件和软件栈(包括 Advanced Micro Devices [2025] 和新兴的 ML 框架)原生支持此粒度的 MX 格式,正如 Bouquet et al. [2026] 所采用。有关块量化器的描述,请参阅附录 D。
我们在附录 D 中提供了实验的完整描述以确保可复现性。
表 1:(MXFP4 E2M1, W4A4) 在 MJHQ-30K 和 SDCI(各 5k 样本)上的表现。LR 分解秩为 128,LoRaQ 主干网使用 MXFP4e2 数据格式。SVDQuant 中 LR 分解的秩为 32,使用 16 位数据格式。每个模型和数据集的最佳量化结果以粗体显示。
MJHQ-30K SDCI
模型 量化器 LR 分解 FID ↓ IR ↑ LPIPS ↓ PSNR ↑ FID ↓ IR ↑ LPIPS ↓ PSNR ↑
FP16 (参考) 16.6 0.948 — — 24.8 0.966 — — 平滑 SVDQuant 19.3 0.876 0.433 15.2 23.3 0.904 0.476 13.9 PixArt-Σ 平滑 LoRaQ 17.5 0.952 0.381 16.4 23.7 0.970 0.399 15.0 KroQuant ( ours ) LoRaQ 16.9 0.989 0.347 16.5 21.9 1.01 0.378 15.2
BF16 (参考) 16.2 1.09 — — 22.4 1.07 — — 平滑 SVDQuant 16.6 1.07 0.256 17.1 23.8 1.02 0.285 15.4 SANA 平滑 LoRaQ 17.4 1.09 0.229 17.7 24.1 1.03 0.259 15.9 KroQuant ( ours ) LoRaQ 16.6 1.09 0.206 18.4 23.6 1.04 0.241 16.4
BF16 (参考) 19.2 0.938 — — 20.8 0.932 — — 平滑 SVDQuant 20.1 0.952 0.378 15.7 22.4 0.869 0.589 13.3 FLUX.1-schnell 平滑 LoRaQ 19.9 0.921 0.331 16.7 22.0 0.961 0.343 15.8 KroQuant ( ours ) LoRaQ 19.8 0.880 0.336 16.9 21.2 0.982 0.283 16.6
6
第 7 页
图 2:PixArt-Σ 在 MXFP4e2 W4A4 下的视觉对比。左三列:FP16(参考)、SVDQuant、KroQuant;KroQuant 使用 LoRaQ 主干进行离线权重校准,而 SVDQuant 使用其基于 SVD 的低秩分解。右三列:FP16(参考)、利用 SmoothQuant 量化器的 LoRaQ,以及依赖 LoRaQ 主干的 KroQuant。
FP16 SVDQuant KroQuant (Ours) FP16 LoRaQ KroQuant (Ours)
表 2:消融实验:PixArt-Σ / MJHQ-30K(5k 样本)上的变换设计选择。所有量化行均使用 MXFP4 E2M1,b = 32,W4A4。“LR Decomposition”列中的“–”表示没有离线权重校正阶段。 量化器 LR 分解 FID ↓ LPIPS ↓ PSNR ↑
平滑 – 130.60 0.712 13.0 KroQuant (Ours) – 41.8 0.516 14.8
平滑 LoRaQ 17.5 0.381 16.4 Hadamard (32×32) LoRaQ 17.5 0.382 15.7 KroQuant (Ours) LoRaQ 16.9 0.347 16.5
4.2 主要定量结果
表 1 将 KroQuant 与 PixArt-Σ、SANA 和 FLUX.1-schnell 上的先前 PTQ 方法在两个基准测试中进行了比较。SVDQuant 和 LoRaQ 使用的平滑方法与 Li et al. [2025]、Bouquet et al. [2026] 中描述的一致。KroQuant 在 PixArt-Σ 和 SANA 的两个基准测试中均匹配或优于所有基线。在 MJHQ-30K 上,KroQuant 在 PixArt-Σ 上将 LPIPS 降低了 0.034(相对 8.9%),在 SANA 上降低了 0.023(10.0%),同时缩小了与全精度参考的 FID 差距(PixArt-Σ 上为 16.9 vs. 16.6;SANA 上匹配 16.6)。SDCI 基准测试证实了这些趋势:KroQuant 在所有三个模型上均取得了最佳的 FID、LPIPS 和 PSNR,在 FLUX.1-schnell 上 LPIPS 提升尤为显著(0.283 vs. LoRaQ 的 0.343,相对降低 17.5%)。
在 FLUX.1-schnell / MJHQ-30K 上,结果喜忧参半:KroQuant 取得了最佳的 FID 和 PSNR,但在 IR 上未能优于 SVDQuant,在 LPIPS 上未能优于 LoRaQ。我们将其归因于 FLUX.1 的双流架构,其中联合注意力块以不同的激活分布交错文本和图像 token;单层 Kronecker 变换可能无法充分拟合这种异构性。FLUX.1 上一致的 SDCI 增益表明,这一缺陷是依赖于基准测试的,而非系统性的。
可视化结果。图 2 展示了 KroQuant 与 LoRaQ 和 SVDQuant 在 PixArt-Σ / MJHQ-30K 上的定性对比。该图表明,与基线相比,KroQuant 生成的图像更清晰,伪影更少,这与改进的定量指标一致。我们在附录 F 中提供了额外的可视化对比。
7
第 8 页
4.3 消融实验
表 2 在 PixArt-Σ / MJHQ-30K 上对变换族进行了消融研究。上半部分展示了未使用 LoRaQ 的独立变换;下半部分在所有行中保持 LoRaQ 校准固定。
独立行(无 LoRaQ)隔离了变换相对于平滑缩放方法的贡献:仅 KroQuant 将 FID 从 130.6 降低至 41.8,这是一个巨大的改进,但仍远未达到全精度质量。这证实了 LoRaQ 权重校准对于缩小剩余差距至关重要,而变换则为权重校准提供了更好的起始点。
使用 LoRaQ 时,块 Hadamard-32 基线在 FID 上与逐通道平滑持平(均为 17.5),但 PSNR 略有下降(15.7 对比 16.4),表明固定的 Hadamard 旋转并不能均匀地减少量化误差,且可能对 PSNR 敏感层产生轻微损害。块 Hadamard-32 与 Kronecker(Hadamard 初始化)之间的差距,$\Delta LPIPS = 0.035$,$\Delta PSNR = 0.8$ dB,证明了每个块中 15 个可学习参数确实优于固定旋转,捕捉到了固定正交变换无法适应的逐层数值结构。
4.4 推理延迟
我们测量了 KroQuant 块对角变换(第 3.5 节)与逐通道缩放(Scale,逐层算术强度的下限)以及快速 Walsh–Hadamard 变换(FWHT [Ashkboos et al., 2025, Liu et al., 2025, Tseng et al., 2024])的墙钟时间成本。
设置。所有内核在 AMD Instinct MI350 上以 FP16 运行,在输出端融合块级 MXFP4e2 量化器,并在 100 次预热调用后对 500 次重复取中位数计时;因此报告的延迟反映了完整的逐层量化器,而不仅仅是变换本身。
我们基准测试了两个 KroQuant 变体:K1 从内存加载预计算的密集 $32 \times 32$ Kronecker 块,而 K2 从 15 个可学习参数动态重建每个块(与 K1 相比,内存压缩 68 倍)。FWHT 要求 $d$ 为 2 的幂,这迫使在生产形状上进行零填充(额外内存流量)或截断(信息丢失);而 K1 和 K2 仅要求 $d$ 能被 32 整除。
图 3 报告了内核延迟和有效 HBM 带宽随令牌数量 $m$ 的变化。K2 始终比 FWHT 快约 1.4 倍,并在每个块应用全秩 $32 \times 32$ 张量核心 GEMM 时达到 Scale 带宽的 85–88%:块对角结构将每个图块限制为一次合并读取和一次合并写入,且无块间通信,而 FWHT 的全局耦合迫使每个输出图块进行多次传递。
K1 在所有测试形状上至少与逐通道缩放一样快(1.00–1.14 倍):每个块单次 $32 \times 32$ 张量核心 GEMM 提取的有效吞吐量高于缩放操作的标量逐元素乘法,尽管对角运算在纸面上更便宜。K2 相比 Scale 最多损失 25%(0.75–0.96 倍):从 15 个参数重新生成 $32 \times 32$ Kronecker 块增加了寄存器
延迟 带宽 仅量化器 仅量化器 Scale (对角) Scale (对角) K1 (预计算) 2000 K1 (预计算) K2 (即时) (即时) (GB/s) K2 10 1 FWHT FWHT(ms) 1500 延迟 带宽 中位数 HBM 1000 500 10 2 有效
27 28 29 210 211 212 213 214 27 28 29 210 211 212 213 214 令牌数量 m 令牌数量 m
图 3:中位数内核延迟和有效 HBM 带宽与令牌数量 $m$ 的关系。$d = 4096$ 维度,在 MI350 上使用 FP16 和 MXFP4e2 量化器进行评估。K1 和 K2 紧密跟踪 Scale,且在所有 $m$ 下比 FWHT 快约 1.4 倍。
8
第 9 页
表3:在MI350、FP16精度下,维度为2的幂次时的内核延迟(µs)。所有方法均融合MXFP4e2输出量化器;Q4报告裸量化器内核作为下限。K2针对每种形状使用最佳的即时变体(重计算与数据移动)。加速比报告为“方法A / 方法B”,意为“方法B比方法A快X倍”。
Q4 Scale K1 K2 FWHT Scale/K1 Scale/K2 FWHT/K2 m d (µs) (µs) (µs) (µs) (µs) 加速比
4096 1024 9.40 10.56 9.96 14.00 18.16 1.06× 0.75× 1.30× 4096 4096 22.64 25.52 23.88 27.80 68.04 1.07× 0.92× 2.45× 4096 16384 74.44 75.44 75.32 78.52 186.04 1.00× 0.96× 2.37× 512 4096 7.20 9.24 8.12 10.28 14.40 1.14× 0.90× 1.40× 512 16384 13.68 16.36 15.36 19.32 33.96 1.06× 0.85× 1.76×
每个工作组内的压力,且在同一列块上操作的工作组之间重计算是冗余的,因此HBM流量的节省部分被每线程算术运算所抵消。尽管如此,两种变体均相比FWHT提供了一致的加速,K2在较大矩阵上最高达到2.45×。Q4列显示,在较大的d值下,融合量化器已主导内核时间,使得任何变换变体与裸量化器之间的余地很小。与FWHT不同,K1和K2原生处理任意维度,避免了非2的幂次层在生产环境中所需的填充或截断。
表3报告了代表FLUX.1和PixArt-Σ形状的2的幂次d的延迟。K1在每个形状上匹配或超过Scale(1.00–1.14×):每个块单次张量核心32×32 GEMM提取的吞吐量高于Scale的标量乘法。K2保持在Scale的25%以内(0.75–0.96×)。K1比K2快约6–10%,反映了即时Kronecker积构建的成本。即时Kronecker积重建带来的寄存器压力抵消了部分HBM流量节省。K2通过68×参数压缩(每个块15个标量 vs 1024个值)进行补偿,使其成为部署的首选变体。K1和K2均相比FWHT提供1.30–2.45×的加速,并可在任意d上运行。在较大的d值下,融合Q4量化器主导内核时间,使得任何变换变体与裸量化器之间的余地很小。
5 结论与局限性
我们提出了KroQuant,这是一种针对扩散Transformer的训练后量化方法,引入了作为在线激活预处理步骤的学习型Kronecker积结构化块变换。关键观察结果是,AdaLN阻止了LLM旋转方法使用的离线权重吸收,且全维在线变换成本过高。通过2×2因子的Kronecker积参数化的块对角学习变换,以O(d)参数成本和张量核心原生计算成本填补了这一空白。
消融实验证实,学习型Kronecker变换比固定的Hadamard初始化具有严格更强的表达能力:从Hadamard(其在FID上匹配每通道平滑)开始,对每个块的15个参数进行基于梯度的优化,在PixArt-Σ上将LPIPS降低了0.035,并将PSNR提高了0.8 dB。硬件分析表明,变换的块对角结构使其保持单遍且受内存限制,在1.4–1.9× FWHT吞吐量的情况下实现了接近带宽最优的执行。
KroQuant与权重侧校正互补:该流水线在激活变换固定后离线应用LoRaQ,且组合效果优于任一单独组件。在所有三个评估的DiT架构(PixArt-Σ、SANA、FLUX.1-schnell)上,KroQuant在MXFP4 E2M1 W4A4下实现了最佳或具有竞争力的LPIPS和PSNR。
局限性。Kronecker积构造每个n×n块仅有3 log2 n个参数,随块大小增长过慢,无法在较大块中捕获GL(n)的全部表达能力;因此KroQuant在实践中仍局限于n = 32。联合多层优化在我们的实验中不稳定,部分原因是由于单位行列式LU参数化,留下了其他参数化作为未来工作。最后,KroQuant输出的对比度往往高于FP16参考,这是一种经验效应,其理论解释超出了本文的范围。
更广泛的影响。更广泛的影响讨论见附录E。
9
第 10 页
参考文献
Advanced Micro Devices。AMD Instinct CDNA 4 指令集架构。技术报告,Advanced Micro Devices,2025年8月。
S. Ashkboos, A. Mohtashami, M. L. Croci, B. Li, P. Cameron, M. Jaggi, D. Alistarh, T. Hoefler, 和 J. Hensman。Quarot:旋转大语言模型中无异常值的4-bit推理。在《第38届神经信息处理系统国际会议论文集》,NIPS ’24,纽约州红钩,美国,2025。Curran Associates Inc. ISBN 9798331314385。
Black-Forest-Labs。Flux.1,2024。网址 https://blackforestlabs.ai/。
Y. Bouquet, A. Khodamoradi, S. Y. Shen, K. Denolf, 和 M. Salzmann。LoRaQ:针对4-bit量化的优化低秩近似,2026。网址 https://arxiv.org/abs/2604.18117。
J. Chen, C. Ge, E. Xie, Y. Wu, L. Yao, X. Ren, Z. Wang, P. Luo, H. Lu, 和 Z. Li。Pixart-Sigma:用于4K文本到图像生成的扩散Transformer的弱到强训练。在《计算机视觉 – ECCV 2024:第18届欧洲会议,意大利米兰,2024年9月29日至10月4日,论文集,第XXXII部分》,第74–91页,柏林,海德堡,2024a。Springer-Verlag。ISBN 978-3-031-73410-6。doi: 10.1007/978-3-031-73411-3_5。网址 https://doi.org/10.1007/978-3-031-73411-3_5。
L. Chen, Y. Meng, C. Tang, X. Ma, J. Jiang, X. Wang, Z. Wang, 和 W. Zhu。Q-DiT:扩散Transformer的精确训练后量化,2024b。网址 https://arxiv.org/abs/2406.17343。
X. Chen, H. Fang, T.-Y. Lin, R. Vedantam, S. Gupta, P. Dollar, 和 C. L. Zitnick。Microsoft COCO 字幕:数据收集与评估服务器,2015。网址 https://arxiv.org/abs/1504.00325。
E. Dellinger 和 A. Khodamoradi。TensorCast:PyTorch的转换/量化库。https://github.com/ROCm/tensorcast,2025。
M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, 和 S. Hochreiter。由双时间尺度更新规则训练的GAN收敛于局部纳什均衡,2018。网址 https://arxiv.org/abs/1706.08500。
D. P. Kingma 和 J. Ba。Adam:一种随机优化方法。arXiv预印本 arXiv:1412.6980,2014。
D. Li, A. Kamko, E. Akhgari, A. Sabet, L. Xu, 和 S. Doshi。Playground v2.5:增强文本到图像生成美学质量的三点见解,2024。
M. Li, Y. Lin, Z. Zhang, T. Cai, J. Guo, X. Li, E. Xie, C. Meng, J.-Y. Zhu, 和 S. Han。SVDQuant:通过低秩分量吸收异常值以实现4-bit扩散模型量化。在《第十三届国际学习表征会议》,2025。网址 https://openreview.net/forum?id=vWR3KuiQur。
Z. Liu, C. Zhao, I. Fedorov, B. Soran, D. Choudhary, R. Krishnamoorthi, V. Chandra, Y. Tian, 和 T. Blankevoort。SpinQuant:通过 learned rotations 进行大语言模型量化。在《第十三届国际学习表征会议》,2025。网址 https://openreview.net/forum?id=ogO6DGE6FZ。
Microsoft, AMD, Arm, Intel, Meta, NVIDIA 和 Qualcomm。开放计算项目 • OCP 微缩放格式 (MX) 规范,2023。网址 https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf。
Y. Shang, Z. Yuan, B. Xie, B. Wu, 和 Y. Yan。扩散模型上的训练后量化。在 CVPR,2023。
Y. Sun, R. Liu, H. Bai, H. Bao, K. Zhao, Y. Li, J. Hu, X. Yu, L. Hou, C. Yuan 等。FlatQuant:平坦度对大语言模型量化至关重要,2024。
10
第 11 页
A. Tseng, J. Chee, Q. Sun, V. Kuleshov, 和 C. De Sa。Quip#:利用 Hadamard 不相干性和格码本实现更优的大语言模型量化。在《第 41 届国际机器学习会议论文集》,ICML’24。JMLR.org,2024。
J. Urbanek, F. Bordes, P. Astolfi, M. Williamson, V. Sharma, 和 A. Romero-Soriano。一图胜千言:在密集描述上评估 CLIP 风格模型。在《IEEE/CVF 计算机视觉与模式识别会议论文集》,第 26700–26709 页,2024。
J. Wu, H. Wang, Y. Shang, M. Shah, 和 Y. Yan。Ptq4dit:扩散 Transformer 的训练后量化。在 NeurIPS,2024。
G. Xiao, J. Lin, M. Seznec, H. Wu, J. Demouth, 和 S. Han。SmoothQuant:大语言模型准确且高效的训练后量化。在《第 40 届国际机器学习会议论文集》,2023。
E. Xie, J. Chen, J. Chen, H. Cai, H. Tang, Y. Lin, Z. Zhang, M. Li, L. Zhu, Y. Lu, 和 S. Han。SANA:使用线性扩散 Transformer 进行高效高分辨率文本到图像合成。在《第十三届国际学习表征会议》,2025。URL https://openreview.net/forum?id=N8Oj1XhtYZ。
J. Xu, X. Liu, Y. Wu, Y. Tong, Q. Li, M. Ding, J. Tang, 和 Y. Dong。Imagereward:学习和评估文本到图像生成的人类偏好。在《第 37 届神经信息处理系统国际会议论文集》,第 15903–15935 页,2023。
R. Zhang, P. Isola, A. A. Efros, E. Shechtman, 和 O. Wang。深度特征作为感知度量的不合理有效性。在 CVPR,2018。
T. Zhao, X. Ning, T. Fang, E. Liu, G. Huang, Z. Lin, S. Yan, G. Dai, 和 Y. Wang。Mixdq:具有度量解耦混合精度量化的内存高效少步文本到图像扩散模型,2024。URL https://arxiv.org/abs/2405.17873。
T. Zhao, T. Fang, H. Huang, R. Wan, W. Soedarmadji, E. Liu, S. Li, Z. Lin, G. Dai, S. Yan, H. Yang, X. Ning, 和 Y. Wang。Vidit-q:用于图像和视频生成的扩散 Transformer 的高效准确量化。在《第十三届国际学习表征会议》,2025。URL https://openreview.net/forum?id=E1N1oxd63b。
11
第 12 页
附录
A 生产层形状下的内核性能
表 4 报告了源自 FLUX.1 和 PixArt-Σ 的生产线性层形状的内核延迟。 与表 3 类似,每个变体都融合了 MXFP4e2 输出量化器;Q4 是裸融合量化器 内核,作为下限,K2 报告了每个形状下最佳的对时变变体(重计算 与数据移动)。FWHT 被省略,因为这些形状大多不是 2 的幂; 包含填充至 2 的幂维度的 FWHT 的公平比较见表 3。
表 4:源自 FLUX.1 和 PixArt-Σ 的生产层形状上的内核延迟 (µs),MI350,FP16。 所有方法均融合 MXFP4e2 输出量化器。加速比报告为 Method A / Method B,意为“Method B 比 Method A 快 X 倍”。 Q4 Scale K1 K2 Scale/K1 Scale/K2 m d (µs) (µs) (µs) (µs) 加速比 加速比
512 3072 6.72 9.04 7.64 8.84 1.18× 1.02× 512 12288 11.80 14.68 13.08 17.00 1.12× 0.86× 4096 1152 10.64 14.00 10.68 14.64 1.31× 0.96× 4096 3072 18.56 21.12 19.28 23.32 1.10× 0.91× 4096 3456 20.60 24.20 21.32 25.36 1.14× 0.95× 4096 4096 22.08 25.80 23.36 27.28 1.10× 0.95× 4096 4608 25.72 28.40 26.40 30.36 1.08× 0.94× 4096 12288 59.64 59.12 59.92 62.80 0.99× 0.94×
除最大形状外,K1 在每个形状上都匹配或优于逐通道缩放(Scale/K1 范围 从 d=12288 时的 0.99× 到 d=1152 时的 1.31×),证实了在生产维度上 与表 3 相同的模式:每个块的融合 32×32 张量核心 GEMM 优于 带宽受限形状上 Scale 的标量逐元素乘法。K2 在大多数形状上保持在 Scale 的 约 7% 以内(Scale/K2 介于 0.91× 和 1.02× 之间);最坏情况是 (m=512, d=12288) 的 0.86×,其中 小批量模式暴露了 K2 的对时变 Kronecker 重建的重计算成本, 而 Scale 内核的图块被完全利用。在 (m=4096, d=12288) 处,裸 Q4 量化器本身需要 59.64 µs, 在此下限之上为任何变换变体留下的余量微乎其微。K1 和 K2 原生处理 这些任意 d 值,避免了 FWHT 在非 2 的幂层上所需的填充或截断。
B K2 流式乘积内核
K2 将第 3.5 节的 Kronecker 公式评估为五个因子 G1, . . . , G5 的累积乘积。每一步提取行索引的一位和列索引的一位,通过向量化标量广播查找当前 Gk 的 相应条目,并将运行图块与结果逐元素相乘。
i = tl.arange(0, 32)[:, None] # 行索引,形状 (32, 1) j = tl.arange(0, 32)[None, :] # 列索引,形状 (1, 32)
步骤 1: K = G1[bit_4(r), bit_4(c)]
ir = (i >> 4) & 1; jc = (j >> 4) & 1 K = tl.where(ir == 0, tl.where(jc == 0, g1_00, g1_01), tl.where(jc == 0, g1_10, g1_11))
步骤 2-5: K *= G_k[bit_{5-k}(r), bit_{5-k}(c)], k = 2, 3, 4, 5
for shift, (g00, g01, g10, g11) in zip( (3, 2, 1, 0), ((g2_00, g2_01, g2_10, g2_11), (g3_00, g3_01, g3_10, g3_11),
12
第 13 页
(g4_00, g4_01, g4_10, g4_11), (g5_00, g5_01, g5_10, g5_11)): ir = (i >> shift) & 1; jc = (j >> shift) & 1 K = K * tl.where(ir == 0, tl.where(jc == 0, g00, g01), tl.where(jc == 0, g10, g11))
每个 tl.where(ir==0, tl.where(jc==0, g_00, g_01), tl.where(jc==0, g_10, g_11)) 将 $G_k$ 的四个标量元素广播到 (32, 32) 图块中,根据行和列的单个位,在每个位置选择正确的元素。逐元素乘法 累积定义 $K[r, c]$ 的乘积 $Q_k G_k[r_{5-k}, c_{5-k}]$。 每一步只有两个 (32×32) 寄存器图块处于活动状态:运行中的乘积 $K$ 和由 tl.where 产生的临时因子 图块。临时图块被立即消耗,并在下一步之前释放。 当 BLOCK_ROWS=128 且使用 8 个 warp(256 个线程)时,每个 (32, 32) 图块每个线程占用 4 个寄存器, 因此内核为 Kronecker 积构造本身最多支付 8 个寄存器。这一预算为将下游 MXFP4e2 量化器(需要其自身的位操作 中间值)融合到同一内核中留出了空间。
C 归一化-旋转交换性
对于任何应用于归一化输入的可逆线性变换 $T \in GL(d)$,若要离线吸收进相邻权重矩阵,归一化必须满足 $Norm(xT) T^{-1} = Norm(x)$。 我们测试正交情况 $T = R \in O(d)$ 作为最简单的实例:此处的任何失败都会推广到非正交但可逆的 $T$。表 5 报告了随机正交矩阵 $R \in \mathbb{R}^{64 \times 64}$ 在随机批次 $X \in \mathbb{R}^{32 \times 64}$(种子 42)上的最大绝对偏差 $Norm(XR)R^\top – Norm(X)_\infty$。RMSNorm 匹配浮点精度,因为其分母仅依赖于 $\|x\|_2$,而正交旋转保持该值不变;它不会与非正交 $T$ 交换。 LayerNorm、BatchNorm、GroupNorm、AdaLN 和 AdaRMSNorm 均产生 $O(1)$ 偏差:归一化要么减去轴对齐的均值,要么应用学习的每通道缩放/平移,或两者兼有(AdaLN)。由于两种失败模式都是轴对齐的,它们无法反转应用于归一化输入的任何跨通道变换;只有每通道输入缩放可以折叠进仿射变换。
表 5:在随机正交旋转 $R$ 下,常见归一化方法中 $Norm(XR)R^\top$ 与 $Norm(X)$ 之间的最大绝对偏差。越小越好;只有 RMSNorm 与正交 $R$ 交换。 归一化方法 最大绝对误差 与 $R \in O(d)$ 交换? RMSNorm $1.7 \times 10^{-6}$ ✓ LayerNorm $9.2 \times 10^{-1}$ ✗ BatchNorm $1.3$ ✗ GroupNorm $2.1$ ✗ AdaLN $4.2$ ✗ AdaRMSNorm $4.4$ ✗
D 实现细节
量化。转发到低秩分支的激活值保持为 16 位(全精度)。 Pixart-Σ 使用 FP16,而 SANA 和 Flux.1 默认使用 BF16。转发到残差分支的激活值和残差权重被量化为 MXFP4e2,并应用 W4A4 矩阵乘法。SVDQuant 将低秩分支保持为全精度,秩为 32。LoRaQ 和 Kroquant 进而将低秩分支量化为 MXFP4e2,秩为 128,其内存开销与 SVDQuant 相当。所有量化算子(MXFP4e2、FP8 E8M0 块缩放以及底层 MX 块格式)均通过 TensorCast 库 [Dellinger and Khodamoradi, 2025] 实现。
13
第 14 页
D.1 KroQuant
校准。对于每个线性层,我们使用 Adam 在 128 个校准激活值上以 $10^{-2}$ 的学习率运行 100 步,通过最小化目标方程 9 来优化 Kronecker 因式分解的块变换 $T$。校准样本与 Li 等人 [2025]、Bouquet 等人 [2026] 相同。遵循与 Li 等人 [2025] 相同的种子配置,为每层随机选择 128 个激活值。这些激活值在所有实验中保持不变并用于公平比较。
D.2 平滑处理
平滑校准 考虑到我们的方法旨在替代基线所使用的校准。作为初步处理校准,Li 等人 [2025]、Bouquet 等人 [2026] 应用了逐通道平滑 [Xiao 等人, 2023] 并结合 Li 等人 [2025] 的逐层调整:对于每个线性层,向量 $\gamma \in \mathbb{R}^d$ 的分量构建如下
$$ \gamma_i = \frac{(\max_j |X_{j,i}|)^\alpha}{(\max_j |W_{i,j}|)^\beta}, \quad i \in \{0, \dots, d-1\}, $$
其中迁移强度 $\alpha, \beta$ 逐层选择,以最小化消耗量化层输出的下游模块的输出均方误差。在此搜索过程中,权重通过秩-32 的 FP16 SVD 截断进行量化,且搜索使用与 KroQuant 相同的校准数据集(第 4.1 节)。
D.3 LoRaQ
在确定 Kroquant 或平滑处理后,对每一层应用离线 LoRaQ 权重校准(第 3.4 节),在 MXFP4e2 低秩分支中使用秩 $r=128$。
优化。我们使用 Adam [Kingma and Ba, 2014] 优化 LoRaQ 的目标函数,学习率为 $10^{-3}$,每个权重 500 步。使用旋转变换对低秩分量进行量化感知微调,采用 Bouquet 等人 [2026] 中的 Cayley SGD 算法,学习率为 $5 \cdot 10^{-1}$,每个权重 500 步。
D.4 SVDQuant
校准。在平滑校准之后,Li 等人 [2025] 使用迭代算法校准其低秩分量,我们也遵循该方法以复现 MXFP4e2 中的方法:递归计算 100 个候选项。根据要校准的线性层,在线性层或 Transformer 块的输出级别计算均方误差 (MSE) 损失。按递归顺序评估每个候选项。停止算法返回目标空间中遇到的第一个局部最小值,并返回关联的低秩候选项,该候选项随后用作辅助分支的低秩分量。
计算。一个 DiT(PixArt-Σ, SANA)所有线性层的校准在单个 NVIDIA H100 上运行,每个模型的墙钟时间仅需数小时。由于其规模,Flux.1 校准需要 NVIDIA H200。实际上,我们在进行平滑处理、Kroquant 和生成图像时同时使用了 4 块 GPU。内核延迟基准测试(第 4.4 节)在 FP16 模式下的 AMD Instinct MI350 上进行,内核使用 Triton 3.6.0 针对 ROCm 7.2.2 实现。
资产许可。我们在各自的许可下使用以下外部资产:PixArt-Σ [Chen 等人, 2024a](CreativeML Open RAIL++-M);SANA [Xie 等人, 2025](Apache 2.0);FLUX.1-schnell [Black-Forest-Labs, 2024](Apache 2.0);MJHQ-30K [Li 等人, 2024];sDCI [Urbanek 等人, 2024];COCO 2017 [Chen 等人, 2015](注释采用 CC BY 4.0,底层图像受其原始 Flickr 许可约束);以及 TensorCast [Dellinger and Khodamoradi, 2025](MIT)。所有资产的使用均符合其适用的许可条款和限制。
14
第 15 页
E 更广泛的影响
KroQuant 在不改变生成能力、训练数据或安全微调的情况下,加速了三个公开发布的扩散 Transformer(PixArt-Σ、SANA、FLUX.1-schnell)的推理。我们讨论了其社会影响的两个方面,以及与相同上游检查点的基线 FP16 部署相比,变化的边界。
积极后果。一种在 W4A4 下提供接近 FP16 图像质量且具备张量核心原生在线变换的方法,降低了大型文本到图像模型每次推理的能源和内存成本。由此带来的部署成本降低,扩大了当前最先进图像生成的访问权限:在资源受限的环境中,较低配置的硬件层级对于学术研究、课堂教学和应用工作变得可行,大规模服务也降低了每个生成样本的碳足迹。通过将旋转参数量保持在 O(d) 并避免任何重新训练,KroQuant 还降低了在商用加速器上部署公共 DiTs 的工程和计算门槛。
消极后果。同样的效率提升使得上游扩散 Transformer 的双重用途风险在大规模运营时成本更低。非自愿的深度伪造、定向虚假信息以及侵犯版权的合成内容,变得对以前无法承担推理成本的行动者可用,而在单个商用加速器上运行的低于 16 位的流水线,相对于 FP16 部署,扩大了潜在滥用操作者的群体。
增量影响的边界。KroQuant 没有引入新的生成能力、新的模态、安全微调的放宽,也没有在提交时发布新的数据集或预训练模型。上游检查点(PixArt-Σ 遵循 OpenRAIL-M,SANA 和 FLUX.1-schnell 遵循 Apache 2.0)的定性威胁模型因此保持不变;只有运营这些模型的成本曲线发生了偏移。滥用缓解的责任仍由原始上游许可条款承担,这些条款已经约束了发布的权重,且并未因量化而放宽。我们将 KroQuant 的增量社会影响视为受限于其所加速的上游模型,同时承认任何推理成本的降低在部署规模下会被放大。
F 额外的视觉结果
15
第 16 页
图4:在不同量化配置下,遵循表1设置,PixArt-Σ生成的图像对比
FP16 SVDQuant LoRaQ KroQuant
16