论文精选

融合INT8 GEMM内核让Ideogram 4.0在RTX 3090上加速1.1倍

Realizing Native INT8 Compute for Diffusion Transformers on Consumer GPUs: A Fused INT8 GEMM Kernel for Ideogram 4.0

精选理由

INT8反超FP8,单卡RTX 3090跑1024px扩散模型

AI 摘要

论文发现消费级Ampere GPU上扩散Transformer的INT8量化常因反量化回bf16而无法利用INT8张量核心。作者为Ideogram 4.0线性层设计了一个融合Triton INT8 GEMM内核,在Ampere张量核心上执行int8×int8→int32,并在epilogue中折叠逐token×逐通道反量化和偏置。该内核实现2.8-4.2倍于bf16的GEMM加速,并保持余弦相似度1.0且无NaN。端到端测试中,在单张RTX 3090上768px分辨率获得约9-10%提速,1024px生成耗时156.5秒,优于NF4(164.5秒)和FP8(172.9秒)基线,且PickScore/CLIPScore无质量损失。

原文 · arXiv cs.LG

Realizing Native INT8 Compute for Diffusion Transformers on Consumer GPUs: A Fused INT8 GEMM Kernel for Ideogram 4.0

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.

  • Ideogram06-11 17:12原文
  • Artificial Analysis06-12 04:52原文