소비자용 GPU에서 Diffusion Transformer를 위한 네이티브 INT8 연산 구현: Ideogram 4.0을 위한 Fused

Diffusion Transformer의 사후 학습 INT8 (W8A8) 양자화 (Quantization)는 속도 최적화를 위해 널리 배포되고 있지만, 소비자용 Ampere GPU에서는 이를 능가해야 할 FP8 및 NF4 대안보다 오히려 느린 경우가 빈번합니다. 우리는 이 원인을 소프트웨어적 결함에서 찾았습니다. 실제 운영 중인 "INT8" 방식은 가중치 (Weights)와 활성화 값 (Activations)을 양자화한 직후 즉시 bf16으로 역양자화 (Dequantize)하여 bf16 행렬 곱셈 (Matrix Multiply)을 수행하며, GPU의 INT8 텐서 코어 (Tensor Cores)를 전혀 사용하지 않기 때문에 하드웨어의 연산 이점을 전혀 활용하지 못합니다.

우리는 역양자화 후 bf16으로 변환하는 경로 대신, Ideogram 4.0 Diffusion Transformer의 선형 레이어 (Linear Layers)에 단일 Fused Triton INT8 GEMM (Ampere 텐서 코어 상의 int8xint8->int32 연산, 토큰당 x 채널당 역양자화 및 에필로그 (Epilogue)에 바이어스 (Bias) 폴딩 포함, GEMM 형상별 오토튜닝 적용)을 삽입하여 이 격차를 해소했습니다. 해당 커널 내에서 int8xint8->int32 누적 (Accumulation)은 torch._int_mm과 비트 단위로 일치하며, 역양자화된 출력은 NaN 없이 코사인 유사도 (Cosine Similarity) 1.0으로 참조값과 일치하며, GEMM당 bf16보다 2.8~4.2배 더 빠르게 실행됩니다.

엔드 투 엔드 (End-to-end) 측면에서 768px 해상도에서 약 1.1배 (~9-10%)의 속도 향상을 제공하며, 1024px 해상도에서는 단일 RTX 3090에서 156.5초 만에 이미지를 생성합니다. 이는 단일 카드 NF4 (164.5초) 및 FP8 (172.9초) 베이스라인보다 빠르며, 이러한 점 추정치 (PickScore/CLIPScore) 상에서 측정 가능한 품질 저하는 없습니다. 이로써 INT8은 가장 느린 변형에서 가장 빠른 변형으로 바뀌었으며, 1024px 생성은 단일 GPU로 실행 가능한 수준이 되었습니다. 주요 속도 기준 (FP8을 약 9.5% 차이로 능가)은 여유롭게 충족되었습니다. NF4와의 격차 (~4.9%, 단일 실행 n=4)는 우리가 정량화하지 않은 실행 간 편차 범위 내에 있으며, 도전적인 목표를 달성한 것으로 해석하는 것이 가장 적절합니다. 마지막으로 정직한 배포 지도를 제시합니다: 이 성과는 소비자용 Ampere에 특화된 것이며, A100 및 B200에서는 동일한 커널이 해당 카드들의 빠른 네이티브 bf16/FP8 경로에 밀리게 됩니다.