让消费级显卡跑扩散模型时，INT8计算真正提速
Realizing Native INT8 Compute for Diffusion Transformers on Consumer GPUs: A Fused INT8 GEMM Kernel for Ideogram 4.0
- 用融合的INT8 GEMM内核直接调用GPU张量核心，跳过无效的反量化步骤
- 在3090上生成1024分辨率图仅需156.5秒，比FP8和NF4更快
- 适合想在消费级显卡上高效运行扩散模型的研究者与开发者
后训练的扩散变压器通常采用W8A8 INT8量化以提升速度，但在消费级Ampere GPU上反而慢于FP8和NF4。我们发现原因在于软件设计：所谓‘INT8’前向计算仅对权重和激活进行量化，随即又反量化回bf16并执行bf16矩阵乘法，未实际使用GPU的INT8张量核心，导致硬件优势完全未发挥。为此，我们在Ideogram 4.0的线性层中，用一个自适应优化的Triton融合INT8 GEMM内核替代原有路径——该内核在Ampere张量核心上实现int8xint8→int32累加，并支持按标记/按通道反量化及偏置折叠。其结果与torch._int_mm位精确一致，输出与参考值的余弦相似度为1.0且无NaN，单次GEMM速度比bf16快2.8-4.2倍。端到端性能在768px下提升约1.1倍（9-10%），1024px下生成时间仅156.5秒，优于单卡NF4（164.5秒）和FP8（172.9秒）基线，点估计质量无明显损失（PickScore/CLIPScore）。INT8由此从最慢变为最快，1024px成为单卡可行方案。主指标（超越FP8约9.5%）显著达成；与NF4差距约4.9%（单次运行n=4），在未量化波动范围内，可视为满足高目标。最后指出：该优势仅限于消费级Ampere，A100和B200因原生快速bf16/FP8路径，相同内核反而更慢。
原文摘要 · Abstract (English)
Post-training INT8 (W8A8) quantization of diffusion transformers is widely deployed as a speed optimization, yet on consumer Ampere GPUs it is frequently slower than the FP8 and NF4 alternatives it is meant to beat. We trace this to a software artifact: the production "INT8" forward quantizes weights and activations only to immediately dequantize them back to bf16 and run a bf16 matrix multiply, never engaging the GPU's INT8 tensor cores, so the hardware's compute advantage is left entirely unrealized. We close this gap with a single fused Triton INT8 GEMM (int8xint8->int32 on Ampere tensor cores, with per-token x per-channel dequantization and bias folded into the epilogue, autotuned per GEMM shape) dropped into the Ideogram 4.0 diffusion transformer's linear layers in place of the dequantize-to-bf16 path. In the kernel, the int8xint8->int32 accumulation is bit-exact against torch._int_mm and the dequantized output matches the reference at cosine similarity 1.0 with no NaNs, running 2.8-4.2x faster than bf16 per GEMM. End to end it delivers a ~1.1x (~9-10%) speedup at 768px, and at 1024px it generates an image in 156.5 s on a single RTX 3090, faster than the single-card NF4 (164.5 s) and FP8 (172.9 s) baselines, at no measurable quality cost on these point estimates (PickScore/CLIPScore). INT8 thus goes from the slowest variant to the fastest, and 1024px becomes single-GPU feasible. The primary speed criterion (beat FP8, by ~9.5%) is comfortably met; the NF4 margin (~4.9%, single-run n=4) is within run-to-run variance we did not quantify and is best read as consistent with meeting the stretch target. We close with an honest deployment map: the win is specific to consumer Ampere, and on A100 and B200 the same kernel loses to those cards' fast native bf16/FP8 paths.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。