PaperlayerView exact revision on arXiv

Source-linked AI summary

APT: Accelerating Diffusion Transformers via Attention Probability-Guided Pruning and Quantization

Sungyeob Yoo, Seeyeon Kim, Joonyong Park, Seunghee Han, Joo-Young Kim

arXiv:2608.25380v1cs.AR

TL;DR

High-resolution DiTs face an O(N^4) self-attention bottleneck that prior pruning and quantization methods do not fully address. APT combines attention probability-guided pruning, adaptive precision, FlashAttention-compatible prediction, and specialized hardware, achieving up to 8.16× speedup and 14.98× higher energy efficiency over NVIDIA A100, with corresponding gains over EXION.

  • Problem

    At 4K resolution, DiT self-attention accounts for over 90% of computation, while prior pruning and quantization techniques provide limited coverage of its dominant operations.

  • Method

    APT uses attention probabilities to jointly guide APDT pruning and dual-precision assignment, TAFA probability prediction, and a co-designed accelerator for irregular execution.

  • Results

    APT achieves up to 8.16× speedup and 14.98× higher energy efficiency over NVIDIA A100, and up to 3.01× speedup and 2.04× higher energy efficiency over EXION.

  • Takeaways & Limitations

    APT reduces high-resolution DiT computation while preserving output quality across PixArt-α, Stable Diffusion 3, and FLUX.

  • Takeaways & Limitations

    FlashAttention does not directly expose attention probabilities because tile normalization is available only after all tiles are processed.

Abstract

from arXiv · show

Recent advances in generative AI have significantly increased the demand for high-resolution image and video generation, positioning diffusion models as a core technology. Among them, Diffusion Transformers (DiTs) have emerged as the state-of-the-art (SOTA) models due to their scalability and output quality. However, self-attention in DiTs incurs significant computational overhead, leading to excessively long latency as the complexity grows with the fourth power of the output resolution. While prior works have attempted to mitigate this cost using sparsity and quantization techniques, they fall short of effectively reducing the computational cost in high-resolution DiTs. In this paper, we present APT, a software-hardware co-designed accelerator for high-resolution DiTs. APT leverages attention probabilities as a unified importance metric to jointly optimize computation through fine-grained pruning and adaptive precision scaling. At the algorithm level, we propose Attention Probability-guided Adaptive Dual Thresholding (APDT), which dynamically performs element selection and precision assignment using dual thresholds. To ensure compatibility with memory-efficient FlashAttention, we introduce Timestep-Aware FlashAttention (TAFA), which predicts attention probabilities across timesteps by exploiting temporal similarity. At the architecture level, we co-design a specialized accelerator that efficiently supports irregular sparsity and dual-precision execution, featuring dynamic mask management, address translation, dual-precision compute units, and a tile-based dataflow. Finally, we evaluate APT on SOTA DiT models, including PixArt-$α$, Stable Diffusion 3, and FLUX. APT achieves up to 8.16$\times$ speedup and 14.98$\times$ higher energy efficiency over NVIDIA A100, and up to 3.01$\times$ speedup and 2.04$\times$ higher energy efficiency over EXION, a SOTA diffusion model accelerator.

1 Introduction

High-resolution DiTs are bottlenecked by self-attention, whose cost grows as O(N^4). APT jointly uses attention probabilities to guide pruning and precision assignment, with FlashAttention-compatible prediction and specialized hardware support.

  • 1 Introduction: At 4K resolution, self-attention accounts for over 90% of total computation and scales as O(N^4).The dominant operations are QK^T and PV.
  • 1 Introduction: Attention probabilities provide a unified importance metric because Softmax concentrates mass on dominant entries and suppresses the rest.This supports decisions about which elements to prune and which to execute at reduced precision.
  • 1 Introduction: FlashAttention compatibility is challenging because it avoids materializing attention probabilities, while probability-guided decisions create irregular sparsity and precision patterns.Existing hardware cannot efficiently handle these irregular patterns.
  • 1 Introduction: APT applies APDT’s dual thresholds to classify attention elements as skipped, low-precision, or high-precision execution.TAFA predicts attention probabilities using normalization statistics reused from adjacent timesteps.
  • 1 Introduction: APT combines tile-based dataflow, dynamic mask management, address translation, and dual-precision compute units in a specialized accelerator.The design targets efficient execution of the algorithm’s irregular sparsity and dual-precision operations.
  • 1 Introduction: 8.16× speedup and 14.98× higher energy efficiency over NVIDIA A100, plus 3.01× speedup and 2.04× higher energy efficiency over EXION, are reported for APT.These are the paper’s stated headline comparisons.

2 Background and Motivation

DiTs progressively denoise latent representations with Transformer blocks, where self-attention computes context through QK^T, Softmax probabilities, and PV. Temporal similarity motivates probability reuse, while prior pruning, quantization, and FlashAttention leave important compatibility gaps.

  • 2.1 Overview of Diffusion Transformer: DiTs refine latent representations across timesteps using blocks containing self-attention, cross-attention, feed-forward layers, and normalization.Self-attention generates Q, K, and V, computes QK^T, applies Softmax, aggregates through PV, and projects the output.
  • 2.2 Temporal Similarity in DiT: Cosine similarity of attention probabilities exceeds 0.92 for consecutive steps in most blocks and decays with larger temporal gaps.This temporal pattern supports reusing probability information from one timestep at adjacent steps.
  • 2.3 Limitations of Prior Approaches: Softmax pruning can remove low-importance elements from PV computation, but fixed thresholds do not accommodate distributional variation across timesteps.The supplied passage identifies this as a limitation for DiTs.
  • 2.3 Limitations of Prior Approaches: Quantization alone provides limited benefit for the attention bottleneck because prior work leaves QK^T and PV in full precision.These operations dominate high-resolution attention cost.
  • 2.3 Limitations of Prior Approaches: FlashAttention saves memory by tiling computation and avoiding materialization of the full attention probability matrix.It processes key and value tiles with partial Softmax and accumulated scaling factors.
  • 2.3 Limitations of Prior Approaches: FlashAttention’s normalization becomes available only after all tiles are processed, preventing direct access to attention probabilities for Softmax pruning.A prediction mechanism is therefore required to approximate probabilities without sacrificing memory efficiency.

3 APT’s Algorithm Optimizations

APT’s algorithm jointly applies attention-probability-guided pruning and dual-precision quantization, while TAFA predicts probabilities within FlashAttention using timestep similarity. APDT adapts decisions across heads and timesteps, preserving quality while reducing memory and computation.

  • 3.2 Timestep-Aware FlashAttention: TAFA predicts attention probabilities within FlashAttention by combining previous-timestep normalization statistics with current-timestep partial sums.This avoids materializing the full attention probability matrix.
  • 3.1 Attention Probability-Guided Adaptive Dual Thresholding: APDT uses dual thresholds to skip low-probability elements, execute intermediate elements at 6-bit precision, and retain high-probability elements at 12-bit precision.The pruning and quantization thresholds are applied element-wise to attention probabilities.
  • 3.1 Attention Probability-Guided Adaptive Dual Thresholding: APDT adapts thresholds per attention head and timestep using profiled standard-deviation statistics, preserving more information early and reducing computation later.Separate parameter pairs keep the quantization threshold above the pruning threshold.
  • 3.2 Timestep-Aware FlashAttention: TAFA’s tile-ratio similarity exceeds 0.95 between consecutive timesteps, while its execution masks achieve over 97% similarity to full-Softmax masks.These similarities support reusing temporal information for probability and mask prediction.
  • 3.2 Timestep-Aware FlashAttention: TAFA reduces on-chip memory requirements by up to 103.4× versus naive APDT integration that materializes full attention probabilities at 4K resolution.The reduction comes from operating within FlashAttention’s tiled framework.

4 APT’s Hardware Architecture

APT’s hardware is co-designed for element-wise irregular sparsity and dual-precision attention execution. Its mask management, address translation, specialized MAC units, and tile-based dataflow support sparse QK^T and PV processing.

  • 4.1 Overall Architecture: APT combines an SD-MPU, VPU, quantizer, on-chip memory, controller, and DMA to support sparsity-aware dual-precision processing.The SD-MPU handles QK^T and PV through mask management, address translation, dual-precision MACs, and transpose control.
  • 4.1 Overall Architecture: APT processes each tile independently while APDT assigns individual elements to pruning, 6-bit, or 12-bit execution.Tile-level scheduling is inherited from FlashAttention, while decisions remain element-wise.
  • 4.1 Overall Architecture: The seven-stage dataflow decodes masks, translates valid addresses, performs sparse dual-precision QK^T, generates thresholds and masks, then executes sparse dual-precision PV.Masks are compressed and stored for reuse at the next timestep, with management overlapped with computation.
  • 4.1 Overall Architecture: Mask reconstruction reduces per-tile mask storage from T^2 to 2T using compressed column and row masks.The conservative reconstruction may retain otherwise pruned elements at low precision rather than skipping them.
  • 4.1 Overall Architecture: The address translator places valid sparse addresses into per-bank FIFOs, allowing one conflict-free access per bank each cycle.The resulting address translation table guides sparse memory accesses for both QK^T and PV.
  • 4.2 Processing Units: DP-MACs natively support 6-bit arithmetic and reconstruct 12-bit results from interleaved 6-bit segments, avoiding dedicated high-precision multipliers.The SD-MPU contains 16 DP-MACs performing 256 multiplications per cycle.

5 Evaluation

APT is evaluated on three DiT models at 1K, 2K, and 4K resolutions using RTL implementation and comparisons against A100, EXION, and an unoptimized APT baseline.

  • 5.1 Experimental Setup: APT is evaluated on PixArt-α, Stable Diffusion 3, and FLUX.1-dev at 1K, 2K, and 4K resolution with batch size 1.Accuracy is measured on COCO 2017 without fine-tuning or retraining.
  • 5.1 Experimental Setup: APT is implemented at RTL in SystemVerilog, synthesized with a 14 nm process, and operated at 800 MHz with HBM2E bandwidth of 2 TB/s.Per-model attention statistics are obtained through offline profiling.
  • 5.1 Experimental Setup: APT is compared with NVIDIA A100, EXION, and APT-Base under configurations scaled to match A100’s peak throughput.A100 uses FP16 FlashAttention, while APT-Base uses the same hardware without pruning or quantization.

5.2 Accuracy Evaluation

APT maintains accuracy comparable to the FP16 baseline across the reported metrics and diffusion-transformer variants. PixArt-α preserves FP16-level quality while pruning over half of elements and using limited high-precision computation.

  • 5.2 Accuracy Evaluation: APT maintains accuracy comparable to the FP16 baseline across inception score, CLIP score, and PSNR at 1K resolution.IS and CLIP variations remain within 2.5%, while PSNR indicates negligible perceptual loss.
  • 5.2 Accuracy Evaluation: 57.34% pruning and only 9.06% high-precision computation preserve FP16-level quality for PixArt-α.The evaluation uses prompts and seeds separate from the profiling data.
  • 5.2 Accuracy Evaluation: Stable Diffusion 3 and FLUX show similar accuracy trends despite differing architectural designs, supporting APDT across DiT variants.The reported results identify attention probabilities as a unified metric for element selection and precision assignment.

5.3 Performance Evaluation

APT improves self-attention and end-to-end performance across DiT models and resolutions, with larger gains at higher resolutions. Its energy efficiency and co-designed hardware support contribute to these results.

  • Self-Attention Performance: 9.71× self-attention speedup over A100 and 3.70× over EXION are reached at 4K.At 1K, APT achieves 4.35×–5.27× over A100 and 1.55×–1.81× over EXION.
  • Ablation: APT-Base already outperforms A100, while APDT and TAFA widen the advantage as resolution and sparsity increase.The tile-based dataflow provides the baseline benefit, and probability-guided pruning and precision selection add further gains.
  • End-to-End Performance: 8.16× end-to-end speedup over A100 at 4K, with up to 3.01× speedup over EXION.At 1K, APT reaches up to 3.89× over A100; the gap widens at higher resolutions.
  • Energy Efficiency: 14.98× higher energy efficiency than A100 and 2.04× higher than EXION are achieved across models and resolutions.Energy efficiency is measured as effective throughput per watt across the full denoising pipeline.
  • Hardware Cost: 160.11 mm2 total area and 128.26 W power characterize the accelerator at 0.8 V/800 MHz.DP-MACs dominate area and power, while mask management and address translation account for 10.26% of area and 7.40% of power.

5.4 Comparison with GPU-Based Optimization

GPU-based optimization struggles to exploit DiT self-attention sparsity and low-bit computation effectively. A100 can remain slower with sparse execution, while GPU quantization leaves dominant attention operations in full precision.

  • GPU Sparsity: A100 remains slower than dense FlashAttention up to 50% sparsity even with structured sparsity.Kernel launch overhead, synchronization, and scattered memory accesses limit sparse GPU execution.
  • GPU Quantization: APT directly targets QK^T and PV, whereas SVDQuant leaves both operations in full precision.These operations dominate high-resolution self-attention computation.
  • Irregular Sparsity: APT’s runtime element-wise sparsity is expected to degrade GPU performance further because it creates irregular access patterns.The compared structured sparsity pattern is more GPU-friendly than APDT’s element-wise pattern.

5.5 Design Space Analysis

The design-space analysis favors 6/12-bit dual precision and shows that pruning and precision reduction increasingly reduce computation at higher resolutions.

  • Dual-Precision Configuration: 6/12-bit achieves perceptual quality nearly identical to 8/16-bit with fewer bit-operations.Configurations involving 4-bit computation introduce visible artifacts and a clear PSNR drop.
  • Dual-Precision Configuration: 6/12-bit yields lower BOPs than 4/16-bit despite the latter’s lower base precision.The configuration provides the best quality–computation trade-off among those evaluated.
  • Computation Reduction: 53.7% average BOPs reduction at 4K comes from pruning alone, increasing to up to 76.4% with pruning and dual precision.Both techniques become more effective as self-attention occupies a larger share of computation at higher resolutions.

6 Related Works

Prior diffusion and attention optimizations face limits in precision, timestep coverage, FlashAttention compatibility, or sparsity granularity. APT combines dynamic probability-guided decisions with hardware support for runtime irregular sparsity.

  • Diffusion Accelerators: Ditto and Diff-DiT support low-bit computation but are limited to at most 8-bit precision, causing noticeable quality degradation on modern DiTs.Ditto also requires recovering full activations from off-chip memory, while Diff-DiT applies approximate attention only to a subset of timesteps.
  • Diffusion Accelerators: EXION combines top-k attention pruning with feed-forward approximation but achieves limited sparsity on modern DiTs and handles only output sparsity.APT instead guides element-wise pruning and dual-precision decisions using predicted attention probabilities, supporting input and output sparsity.
  • Attention Accelerators: Sanger, FACT, and SpAtten require explicit attention-score construction, which is incompatible with FlashAttention’s tiled execution.APT uses TAFA to predict attention probabilities from previous-timestep normalization statistics within tiled execution.
  • Sparse Quantization: Static weight-statistics methods miss the evolving importance of activation–activation attention operations across diffusion timesteps.APT uses attention probabilities as a dynamic importance signal and operates at finer granularity than block- or token-group-based sparse attention.
  • Software Techniques: ToCa is incompatible with FlashAttention because it constructs the full attention map, while DuCa achieves only limited speedup.APT integrates APDT decisions with mask management and address translation for runtime dynamic sparsity.

7 Conclusion

APT addresses the self-attention bottleneck in high-resolution DiTs through co-designed probability-guided pruning, quantization, and hardware support. It achieves substantial speedup and energy-efficiency gains over A100 and EXION, with benefits increasing at higher resolutions.

  • APT jointly uses attention probabilities to guide pruning and quantization in high-resolution DiTs.APDT dynamically prunes and quantizes attention elements while preserving output quality across three SOTA DiT models.
  • TAFA enables probability-guided optimization within FlashAttention by reusing normalization statistics across timesteps.This avoids full attention-probability materialization while supporting APDT.
  • The SD-MPU architecture handles dynamic element-wise sparsity through mask management, address translation, and dual-precision MAC units.These components are co-designed with APDT and incur modest area overhead.
  • 8.16× and 3.01× speedup, and 14.98× and 2.04× higher energy efficiency are achieved over A100 and EXION, respectively.The advantage grows at higher resolutions.
Loading 2608.25380v1…