AA-DiT: Algorithm-Architecture Co-Design for Diffusion Transformer Acceleration

原文链接: IEEE TCAD 2026 Authors: Siqi Li, Flavio Ponzina, Tajana Rosing (UC San Diego, SDSU)

摘要

AA-DiT 是一种面向 Diffusion Transformer (DiT) 的算法-架构协同设计加速器，通过挖掘 DiT 执行中的多级空间和时间相似性，实现了高达 6.7 倍 的端到端加速。论文提出了静态/动态混合注意力算法，能够在运行时检测并跳过不必要的计算；同时利用跨时间步的中间结果复用，显著降低计算负担。在 7nm 工艺下，AA-DiT 仅需 27.97 mm² 面积（仅为 A100 GPU 的 3.39%），相比 A100 GPU、ViTCoD 和 Cambricon-D 等基线实现了显著的性能和能效提升。

1. 问题定义：Diffusion Transformer 的加速挑战

1.1 DiT 的崛起与计算瓶颈

Diffusion Transformer (DiT) 正在成为 AIGC 领域的主流架构：

“The backbone of diffusion models is currently shifting from U-Net-based CNNs to Diffusion Transformers (DiTs) in order to achieve better generative quality and scalability.”

DiT 的优势:

  • 更好的生成质量和可扩展性
  • 已应用于 Sora、Stable Diffusion 3 等前沿产品
  • 强大的长距离依赖建模能力

但计算挑战严峻:

  • 迭代去噪过程需要多个时间步（通常 50-1000 步）
  • 每个时间步包含大规模 Transformer 计算
  • 在 A100 GPU 上生成 512×512 图像需要 超过 10 秒

1.2 现有加速方案的局限

方案类型 代表工作 局限
NLP Transformer 加速器 ViTCoD 不考虑扩散模型的迭代特性
扩散模型加速器 Cambricon-D 仅利用时间相似性，未充分探索算法-架构协同设计空间
静态稀疏注意力 BigBird, Longformer 无法捕获动态全局信息
动态稀疏注意力 DSA 需要大量训练/微调开销

核心问题: 缺乏针对 DiT 独特特性的系统性加速方案

1.3 多级相似性的发现

论文通过系统分析 DiT 模型结构和计算过程，发现了两类关键相似性：

空间相似性 (Spatial Similarity):

  • 多头注意力中，高注意力分数集中在对角线和特定垂直带
  • 表明不同注意力头分别负责局部细节和全局结构

时间相似性 (Temporal Similarity):

  • 相邻时间步的注意力图和 MLP 输出具有显著相似性
  • 贯穿迭代计算的多个模块

2. AA-DiT 算法设计

2.1 静态/动态混合注意力算法

基于空间相似性的观察，论文提出了创新的混合注意力机制：

核心洞察:

“Important attention scores tend to cluster along a static diagonal region and a set of dynamic columns.”

算法设计:

┌─────────────────────────────────────────────────────────────┐
│           Static/Dynamic Hybrid Attention                    │
├─────────────────────────────────────────────────────────────┤
│                                                              │
│   Attention Map Pattern:                                     │
│   ┌─────────────────────────────────────┐                   │
│   │ ████░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │  Static Diagonal │
│   │ ░░███░░░░░░░░░░░░░░█░░░░░░░░░░░░░░░ │  (Always kept)   │
│   │ ░░░██░░░░░░░░░░░░░██░░░░░░░░░░░░░░░ │                  │
│   │ ░░░░█░░░░░░░░░░░░░██░░░░░░░░░░░░░░░ │  Dynamic Columns │
│   │ ░░░░░░░░░░░░░░░░░░█░░░░░░░░░░░░░░░░ │  (Runtime sel.)  │
│   │ ░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░░ │                  │
│   └─────────────────────────────────────┘                   │
│                                                              │
│   Compute Flow:                                              │
│   1. Static part: Predefined diagonal pattern                │
│   2. Dynamic part: Runtime screening of important columns    │
│   3. Skip: Low-importance attention scores                   │
│                                                              │
└─────────────────────────────────────────────────────────────┘

关键优势:

  • 无需重训练: 可直接应用于各种 DiT 模型
  • 细粒度: 在运行时动态选择重要计算
  • 硬件友好: 轻量级筛选机制

2.2 跨时间步中间结果复用

利用时间相似性，论文提出了全面的中间结果复用策略：

复用范围:

  • 注意力模块: 复用前一时间步的注意力图
  • MLP 模块: 复用前一时间步的 MLP 输出

复用机制:

Timestep N:   Compute full Attention(Q, K, V) → Store
                    ↓
Timestep N-1: Reuse stored attention where similar
              Compute only changed parts
                    ↓
Timestep N-2: Continue reuse with update

质量保证:

  • 相似性阈值控制，确保生成质量不下降
  • 自适应复用策略，根据内容动态调整

3. AA-DiT 架构设计

3.1 跨层级协同设计

算法创新带来了三个架构挑战，论文提出了系统性的解决方案：

挑战 1: 运行时动态检测的硬件支持

  • 解决方案: 轻量级硬件筛选单元
  • 开销: 仅占总面积 0.008%，总功耗 0.1%

挑战 2: 不规则计算和内存访问模式

  • 解决方案: 块级聚集数据流 (Block-wise Gathering Dataflow)
  • 效果: 缓解混合注意力引起的不规则性

挑战 3: 复用算法的高效调度

  • 解决方案: 复用感知调度机制
  • 效果: 重叠数据移动和计算，最小化延迟

3.2 硬件架构概览

┌─────────────────────────────────────────────────────────────┐
│                    AA-DiT Accelerator                        │
├─────────────────────────────────────────────────────────────┤
│                                                              │
│  ┌─────────────┐    ┌─────────────┐    ┌─────────────┐     │
│  │   PE Array  │◄──►│  Screening  │◄──►│    SFU      │     │
│  │  (1.52mm²)  │    │  (0.0022mm²)│    │  (0.064mm²) │     │
│  │  7759.87mW  │    │   13.68mW   │    │  337.28mW   │     │
│  └──────┬──────┘    └─────────────┘    └─────────────┘     │
│         │                                                    │
│         ▼                                                    │
│  ┌─────────────────────────────────────────────────────┐    │
│  │              Memory Hierarchy                        │    │
│  │  ┌───────────────┐      ┌───────────────┐          │    │
│  │  │ Activation    │      │ Weight SRAM   │          │    │
│  │  │ SRAM (15.83mm²)│      │ (10.55mm²)    │          │    │
│  │  │ 2856.69mW     │      │ 1904.46mW     │          │    │
│  │  └───────────────┘      └───────────────┘          │    │
│  └─────────────────────────────────────────────────────┘    │
│                                                              │
│  Total: 27.97 mm², 12871.98 mW (7nm technology)            │
│                                                              │
└─────────────────────────────────────────────────────────────┘

3.3 关键硬件组件

1. 轻量级筛选单元 (Screening Unit):

  • 面积: 0.0022 mm² (几乎可以忽略)
  • 功耗: 13.68 mW (总功耗的 0.1%)
  • 功能: 运行时检测注意力分数重要性

2. 块级聚集数据流:

  • 处理混合注意力引起的不规则访问
  • 优化内存带宽利用率
  • 支持动态稀疏模式的高效执行

3. 复用感知调度器:

  • 协调跨时间步的数据复用
  • 重叠计算和数据移动
  • 最小化端到端延迟

4. 实验结果

4.1 端到端性能

AA-DiT 在多个 DiT 模型上进行了评估：

基线 平均加速比 备注
A100 GPU 6.73× 最先进的 GPU
Cambricon-D 2.55× 专用扩散模型加速器
ViTCoD 1.85× 专用 Transformer 加速器

分解分析:

模块 vs Cambricon-D vs ViTCoD
Attention 2.3× 1.46×
MLP 2.12× 1.92×

4.2 能效分析

指标 vs Cambricon-D vs ViTCoD
能效提升 2.18× 1.69×

4.3 面积与功耗

参数 AA-DiT A100 GPU 比例
面积 27.97 mm² 826 mm² 3.39%
工艺 7nm 7nm -

面积分解:

模块 面积 (mm²) 占比
PE Array 1.52 5.4%
Screening 0.0022 0.008%
SFU 0.064 0.23%
Activation SRAM 15.83 56.6%
Weight SRAM 10.55 37.7%

4.4 生成质量

AA-DiT 在显著加速的同时，保持了生成质量：

  • FID 分数与原始模型相当
  • 视觉质量无明显下降
  • 算法-架构协同设计确保质量-效率平衡

5. 为什么对 AI 硬件重要

5.1 算法-架构协同设计的典范

AA-DiT 展示了算法-架构协同设计的强大威力：

传统方法:

  • 算法优化 → 硬件适配
  • 硬件设计 → 算法适配
  • 两者分离，次优解

协同设计:

  • 算法洞察指导硬件设计
  • 硬件约束反馈算法优化
  • 全局最优解

关键启示:

“The dedicated accelerator achieves significantly more compact integration compared with an advanced GPU.”

仅用 A100 3.39% 的面积实现 6.7 倍加速，证明了协同设计的价值。

5.2 对生成式 AI 芯片的启示

1. 多级相似性挖掘:

  • 空间相似性: 注意力模式的结构化稀疏
  • 时间相似性: 迭代计算的跨步复用
  • 未来可探索: 层间相似性、模型间相似性

2. 混合稀疏模式:

  • 静态稀疏: 预定义模式，硬件友好
  • 动态稀疏: 运行时适配，灵活性高
  • 两者结合: 兼顾效率和适应性

3. 轻量级硬件支持:

  • 筛选单元仅占 0.008% 面积
  • 证明复杂算法可以通过精巧设计高效硬件化
  • 为其他动态算法硬件化提供参考

5.3 对神经符号 AI 的关联

虽然 AA-DiT 针对扩散模型，但其设计思想对神经符号 AI 有重要启示：

结构化稀疏性:

  • 神经符号 AI 通常具有明确的结构
  • 可以设计针对性的稀疏计算模式
  • 类似混合注意力的静态/动态结合

迭代推理加速:

  • 符号推理通常是迭代的
  • 可以借鉴跨时间步复用的思想
  • 中间结果缓存和复用

硬件-算法协同:

  • 神经符号 AI 的离散特性需要专门硬件支持
  • 协同设计可以显著提升效率

6. 局限与未来方向

6.1 当前局限

  • 任务范围: 主要针对图像生成任务，未扩展到视频生成
  • 模型范围: 在特定 DiT 变体上验证，通用性待进一步验证
  • 精度约束: 当前设计针对 FP16/BF16，低精度扩展有待探索
  • 多模态: 未涉及文本-图像联合生成的特殊优化

6.2 未来研究方向

技术演进:

  • 扩展到视频扩散模型（如 Sora）
  • 支持更低精度（INT8/INT4）的量化感知设计
  • 多模态融合计算的专用架构
  • 与光计算、存算一体等新兴技术结合

应用扩展:

  • 实时交互式生成应用
  • 边缘设备上的高效推理
  • 与 LLM 结合的统一生成架构

7. 总结

AA-DiT 代表了扩散模型加速器设计的重要进展：

  1. 算法创新: 静态/动态混合注意力 + 跨时间步复用
  2. 架构创新: 轻量级筛选单元 + 块级聚集数据流 + 复用感知调度
  3. 性能突破: 6.7 倍加速，仅 3.39% 面积，2.18 倍能效提升
  4. 设计范式: 算法-架构协同设计的成功范例

对于 AI 芯片设计，AA-DiT 提供了以下关键启示：

  • 多级相似性挖掘: 空间、时间、层间相似性的系统利用
  • 混合稀疏模式: 静态规则性与动态适应性的平衡
  • 轻量级硬件支持: 复杂算法的高效硬件化路径
  • 协同设计方法论: 算法洞察与硬件约束的深度融合

随着生成式 AI 的快速发展，类似 AA-DiT 的专用加速器将在 AIGC 应用中发挥越来越重要的作用，为实时、高效、低成本的图像和视频生成提供硬件基础。

参考文献

  1. Li, S., Ponzina, F., & Rosing, T. (2026). AA-DiT: An Algorithm-Architecture Co-Design for Diffusion Transformer Acceleration. IEEE Transactions on Computer-Aided Design of Integrated Circuits and Systems. DOI: 10.1109/TCAD.2026.3674447
  2. Peebles, W., & Xie, S. (2023). Scalable diffusion models with transformers. ICCV.
  3. Rombach, R., et al. (2022). High-resolution image synthesis with latent diffusion models. CVPR.
  4. Chen, J., et al. (2023). ViTCoD: Vision Transformer Acceleration via Dedicated Algorithm and Accelerator Co-Design. HPCA.
  5. Lu, Y., et al. (2023). Cambricon-D: A fully integrated diffusion model accelerator with conditional computing. ISCA.