QFlash: Bridging Quantization and Memory Efficiency in Vision Transformer Attention

2026-04-28

AI Summary

Integer quantization of Vision Transformer attention mechanisms faces significant challenges, including numerical instability, inefficiency of exponential operations, and constraints imposed by coarse quantization granularity, which hinder a fully integer implementation of FlashAttention. This work proposes QFlash, the first end-to-end integer-domain FlashAttention framework, which overcomes three key obstacles—scaling explosion in block-wise accumulation, inefficient GPU shift-based exponentiation, and restrictive uniform quantization scales—through a custom Triton kernel that efficiently computes softmax entirely in the integer domain. By integrating integer arithmetic, block-wise computation strategies, and tensor-level quantization, QFlash achieves up to 8.69× speedup on ViT, DeiT, and Swin models, reduces energy consumption by 18.8% compared to FP16 FlashAttention, and maintains lossless Top-1 accuracy on ViT and DeiT while preserving competitive performance on Swin.

Abstract

FlashAttention improves efficiency through tiling, but its online softmax still relies on floating-point arithmetic for numerical stability, making full quantization difficult. We identify three main obstacles to integer-only FlashAttention: (1) scale explosion during tile-wise accumulation, (2) inefficient shift-based exponential operations on GPUs, and (3) quantization granularity constraints requiring uniform scales for integer comparison. To address these challenges, we propose \textit{QFlash}, an end-to-end integer FlashAttention design that performs softmax entirely in the integer domain and runs as a single Triton kernel. On seven attention workloads from ViT, DeiT, and Swin models, QFlash achieves up to 6.73$\times$ speedup over I-ViT and up to 8.69$\times$ speedup on Swin, while reducing energy consumption by 18.8\% compared to FP16 FlashAttention, without sacrificing Top-1 accuracy on ViT/DeiT and remaining competitive on Swin under per-tensor quantization. Our code is publicly available at https://github.com/EfficientCompLab/qflash.

Problem

• quantization
• Vision Transformer
• FlashAttention
• integer-only computation
• memory efficiency

Innovation

• QFlash
• integer-only attention
• quantized softmax
• Vision Transformer
• memory-efficient attention