DxPTA: 面向光子Transformer加速器硬件/软件协同设计的基于光学数据流引导的架构设计空间探索方法

# DxPTA：基于光学数据流引导策略的光学Transformer加速器HW/SW协同设计的架构设计空间探索  
来源：https://arxiv.org/html/2606.06515  
Rachmad Vidya Wicaksana Putra, Solomon Micheal Serunjogi, Mahmoud Rasras, 和 Muhammad Shafique  
Rachmad Vidya Wicaksana Putra 隶属于 eBRAIN Lab，纽约大学阿布扎比分校工程学院，阿拉伯联合酋长国；邮箱：[email protected]。  
Solomon Micheal Serunjogi 隶属于光子研究实验室 (PRL)，纽约大学阿布扎比分校工程学院，阿拉伯联合酋长国；邮箱：[email protected]。  
Mahmoud Rasras 是光子研究实验室 (PRL) 主任，纽约大学阿布扎比分校工程学院，阿拉伯联合酋长国 (UAE)；邮箱：[email protected]。  
Muhammad Shafique 是 eBRAIN Lab 主任，纽约大学阿布扎比分校工程学院，阿拉伯联合酋长国；邮箱：[email protected]。  

###### 摘要  
基于Transformer的网络已成为具有最先进性能的杰出AI模型，潜在地为通用人工智能 (AGI) 铺平道路。然而，其庞大的规模仍然阻碍了高效实现，因此需要替代方案来实现其能效加速。最近，最先进的研究提出了光学Transformer加速器 (PTA)，与传统电子加速器相比，在速度和能效方面有显著提升。然而，它们的PTA架构是在未考虑应用约束（如面积、功耗、能量和延迟）的情况下开发的。此外，其手动设计方法也需要巨大的设计时间来为特定应用确定合适的架构，因此该方法不可扩展。为解决这些限制，我们提出DxPTA，一种新颖的设计空间探索方法，用于实现满足所有约束的合适PTA架构的高效硬件/软件协同设计。该方法通过以下方式实现：(1) 基于相干光学数据流识别PTA架构参数；(2) 分析参数的影响/重要性；(3) 利用此分析设计约束感知架构搜索算法。实验结果表明，我们的DxPTA可以为不同的基于Transformer的模型（即DeiT-T/S/B和BERT-B/L）找到合适的PTA架构。在50mm²面积、5W功耗、50mJ能量和10ms延迟的约束下，它实现了最高26mm²面积、4.8W功耗、39mJ能量和6ms延迟；搜索时间比穷举法快15.2倍。这些结果证明了DxPTA方法在实现面向多样化基于AGI的应用的高效PTA设计方面的潜力。

## I. 引言  
基于Transformer的网络模型 [31 (https://arxiv.org/html/2606.06515#bib.bib58)]，如视觉Transformer (ViTs) 和大语言模型 (LLMs)，已成为解决各种机器学习任务（如视觉和自然语言处理 (NLP)）中具有最先进性能的杰出AI模型 [6 (https://arxiv.org/html/2606.06515#bib.bib57),30 (https://arxiv.org/html/2606.06515#bib.bib33),11 (https://arxiv.org/html/2606.06515#bib.bib55),9 (https://arxiv.org/html/2606.06515#bib.bib56)]，因此潜在地为通用人工智能 (AGI) 铺平道路 [18 (https://arxiv.org/html/2606.06515#bib.bib4)] [33 (https://arxiv.org/html/2606.06515#bib.bib5)]。然而，这种最先进的性能伴随着高计算和内存需求，如图1 (https://arxiv.org/html/2606.06515#S1.F1) (a) 所示，从而导致巨大的功耗/能量消耗 [9 (https://arxiv.org/html/2606.06515#bib.bib56)]。这种情况使得在跨不同应用的广泛部署中处理Transformer模型时难以获得高性能效率。

一个潜在的解决方案是使用专用加速器来加速Transformer推理，从而最小化功耗并提高能效 [17 (https://arxiv.org/html/2606.06515#bib.bib17),22 (https://arxiv.org/html/2606.06515#bib.bib14),32 (https://arxiv.org/html/2606.06515#bib.bib29),26 (https://arxiv.org/html/2606.06515#bib.bib30),36 (https://arxiv.org/html/2606.06515#bib.bib27),10 (https://arxiv.org/html/2606.06515#bib.bib3),35 (https://arxiv.org/html/2606.06515#bib.bib28)]；见图1 (https://arxiv.org/html/2606.06515#S1.F1) (b)。然而，传统的电子加速器面临其性能效率下降（即单位面积功耗增加和性能提升变慢）的挑战，因为晶体管电路达到了Dennard缩放的极限 [28 (https://arxiv.org/html/2606.06515#bib.bib13)]。

最近，基于电子-光子集成电路 (EPIC) 的解决方案，即光子加速器 [25 (https://arxiv.org/html/2606.06515#bib.bib19),23 (https://arxiv.org/html/2606.06515#bib.bib24),8 (https://arxiv.org/html/2606.06515#bib.bib26),34 (https://arxiv.org/html/2606.06515#bib.bib8)]，已被研究作为电子加速器的替代方案，以实现显著的加速和效率提升，因为它具有超高速、高带宽和低能耗的特点 [8 (https://arxiv.org/html/2606.06515#bib.bib26)]。因此，使用光子加速器加速神经网络 (NN) 工作负载的研究正在积极进行。例如，光子张量核心 (PTC) 的开发利用了诸如 Mach-Zehnder 干涉仪 (MZI) [24 (https://arxiv.org/html/2606.06515#bib.bib23)]、微环谐振器 (MRR) 组 [29 (https://arxiv.org/html/2606.06515#bib.bib21)] [27 (https://arxiv.org/html/2606.06515#bib.bib22)] 和相变材料 (PCM) 交叉阵列 [7 (https://arxiv.org/html/2606.06515#bib.bib20)] 等光学组件。然而，这些工作仍然针对加速传统的卷积神经网络 (CNN) 工作负载 [37 (https://arxiv.org/html/2606.06515#bib.bib32)]，因此表明需要进一步研究以实现Transformer模型的高性能和节能推理。

因此，本文针对的研究问题是：如何有效地利用基于光子的加速器实现Transformer模型的高性能和节能推理？解决此问题可能有助于基于Transformer的模型在基于光子的计算系统上的高效部署。

参见图注  
图 1: (a) 基于Transformer的模型通常以更大的内存大小（即更多的参数数量）为代价来提升性能；数据来自 [9 (https://arxiv.org/html/2606.06515#bib.bib56)]。(b) 使用不同计算平台运行 Data-efficient Image Transformer Base (DeiT-B) [30 (https://arxiv.org/html/2606.06515#bib.bib33)] 的实验结果：CPU、GPU、基于CMOS的加速器（即 AutoViT-4bit [16 (https://arxiv.org/html/2606.06515#bib.bib25)] 和 HeatViT-8bit [5 (https://arxiv.org/html/2606.06515#bib.bib31)]）和基于光子的 Lightening-Transformer (LT) 加速器（即 LT-Base-4bit 和 LT-Large-4bit）；数据来自 [37 (https://arxiv.org/html/2606.06515#bib.bib32)]。

### I-A. 光学Transformer加速器 (PTA) 的最新进展及其局限性  
最近的研究提出了采用静态操作PTC的PTA设计，使用MRR组 [14 (https://arxiv.org/html/2606.06515#bib.bib7),13 (https://arxiv.org/html/2606.06515#bib.bib6),3 (https://arxiv.org/html/2606.06515#bib.bib15),1 (https://arxiv.org/html/2606.06515#bib.bib10),2 (https://arxiv.org/html/2606.06515#bib.bib12)] 和PCM交叉阵列 [15 (https://arxiv.org/html/2606.06515#bib.bib9)]。另一项研究提出了基于动态操作PTC的 Lightening-Transformer (LT) [37 (https://arxiv.org/html/2606.06515#bib.bib32)]。它启发了在数模转换器 (DAC) 设计 [12 (https://arxiv.org/html/2606.06515#bib.bib11),4 (https://arxiv.org/html/2606.06515#bib.bib18)] 和可重构性 [38 (https://arxiv.org/html/2606.06515#bib.bib16)] 方面的进一步研究。LT通过实现全范围输入操作数的动态操作，提升了Transformer推理的性能和效率，使其成为最先进的PTA设计。尽管有这些优势，所有这些工作仍存在以下局限性。

- • 它们的架构是在未考虑应用约束（例如面积、功耗、能量和延迟）的情况下开发的，因此其设计不能直接应用于目标应用，导致性能和效率提升欠佳。
- • 其手动设计方法需要巨大的设计时间和功耗/能量消耗来为特定应用开发合适的架构，因此该方法不可扩展。

为了说明最新技术的局限性和相关研究挑战，我们进行了一个实验案例研究，将在第I-B节 (https://arxiv.org/html/2606.06515#S1.SS2) 中讨论。

### I-B. 案例研究与研究挑战  
参见图注  
图 2: 考虑4位LT加速器中不同架构参数配置（即Nt和Nc）的实验结果，针对DeiT-Base [30 (https://arxiv.org/html/2606.06515#bib.bib33)]：(a) 功耗和面积；(b) 能量消耗和延迟。  
我们探索了最先进的4位LT加速器 [37 (https://arxiv.org/html/2606.06515#bib.bib32)] 中不同架构参数的影响。这里，我们改变瓦片数量 (Nt) 和每个瓦片的核心数量 (Nc)。工作负载方面，我们考虑DeiT-Base模型 [30 (https://arxiv.org/html/2606.06515#bib.bib33)]。LT硬件架构和实验设置的详细信息分别在第II-B节 (https://arxiv.org/html/2606.06515#S2.SS2) 和第IV节 (https://arxiv.org/html/2606.06515#S4) 中提供。实验结果如图2 (https://arxiv.org/html/2606.06515#S1.F2) 所示，我们从中得出以下关键观察结果。

- • 不同的配置导致不同的面积、功耗、能量和延迟特性，从而突显了开发光子加速器的广泛设计选择。
- • 增加Nt或Nc会导致更高的功耗和更大的面积，因为电路更复杂（见A），但由于并行性增加，可能会减少延迟和能量消耗（见B）。这显示了加速器设计中进行权衡分析的必要性。
- • 最先进的LT设计（Nt=4，Nc=2）可能无法满足约束。例如，在最大5W的低功耗应用中，该设计会消耗显著更多的功率（约15W），如C所示，表明需要定制架构。

这些观察暴露了为目标研究问题设计解决方案时的几个研究挑战，概述如下。

- • 解决方案应利用光子器件和光学数据流的特性来找到满足所有约束的PTA架构，确保其适用于多样化的应用。
- • 解决方案应最小化PTA架构的搜索时间，从而加快设计时间，并为多样化的应用提供可扩展的设计方法。

### I-C. 我们的新颖贡献  
为了解决目标研究问题和相关挑战，我们提出**DxPTA**，一种新颖的架构设计空间探索方法，利用相干光学数据流引导策略，实现满足多个约束（即面积、功耗、能量和延迟）的光学Transformer加速器的高效硬件/软件 (HW/SW) 协同设计。它采用以下关键步骤（参见图3 (https://arxiv.org/html/2606.06515#S1.F3) 的概述和图4 (https://arxiv.org/html/2606.06515#S1.F4) 的细节）。

- • **识别PTA的架构参数** (第III-A节 (https://arxiv.org/html/2606.06515#S3.SS1))：旨在分析PTA架构（包括其层次结构和光子器件）及其相干光学数据流的特性，以识别主要的架构参数。
- • **分析架构参数的影响** (第III-B节 (https://arxiv.org/html/2606.06515#S3.SS2))：通过观察架构参数（例如Nt和Nc）对面积、功耗、能量和延迟的影响，确定其重要性。这些信息将用于架构搜索。
- • **设计约束感知搜索算法** (第III-C节 (https://arxiv.org/html/2606.06515#S3.SS3))：通过利用相干光学数据流和参数重要性，探索架构候选，评估它们的能量-延迟乘积 (EDP)，然后选择具有最低EDP且满足所有约束的架构。

**关键结果：** 我们使用Python实现评估了我们的DxPTA方法，然后在Nvidia RTX 6000 Ada GPU机器上运行，同时考虑多样化的Transformer工作负载（即DeiT-T/S/B¹和BERT-B/L²）。实验结果表明，DxPTA成功找到了满足所有约束的加速器架构。对于所有研究的模型，在50mm²面积、5W功耗、50mJ能量和10ms延迟的约束下，它实现了最高26mm²面积、4.8W功耗、39mJ能量和6ms延迟，搜索时间比穷举法快15.2倍。

参见图注  
图 3: 我们在这项工作中的新颖贡献。  
参见图注  
图 4: 我们的新颖DxPTA方法及其关键步骤：(1) 识别PTA的架构参数；(2) 分析架构参数的影响；(3) 设计约束感知搜索算法。

## II. 预备知识  
### II-A. 基于Transformer的网络  
基于Transformer的网络通常由多个相同的块组成，称为编码器和解码器块。每个块包括一个多头自注意力 (MHA) 模块、一个前馈网络 (FFN)、快捷连接以及层归一化 (LN) [37 (https://arxiv.org/html/2606.06515#bib.bib32)]。此外，解码器块还具有交叉注意力和掩码自注意力模块。基本的编码器块可以用公式 (1) 和 (2) 表示。这里，X_l 是第l层的输入序列。  
公式 (1): $\hat{\mathbf{X}}_{l+1} = \textit{MHA}(\textit{LN}(\mathbf{X}_l)) + \mathbf{X}_l$  
公式 (2): $\mathbf{X}_{l+1} = \textit{FFN}(\textit{LN}(\hat{\mathbf{X}}_{l+1})) + \hat{\mathbf{X}}_{l+1}$  
多头自注意力 (MHA) 模块有H个自注意力头，每个头将输入向量转换为独立的向量：查询 (Q)、键 (K) 和值 (V) 向量。这些输入向量之间的注意力函数可以使用公式 (3) 计算。这里，d_k 是Q和K的维度。  
公式 (3): $Attention(\mathbf{Q},\mathbf{K},\mathbf{V}) = softmax\left(\frac{\mathbf{Q}\mathbf{K}^\intercal}{\sqrt{d_k}}\right)\mathbf{V}$

### II-B. 光学Transformer加速器 (PTA)  
参见图注  
图 5: LT加速器的架构；基于 [37 (https://arxiv.org/html/2606.06515#bib.bib32)]。  
在这项工作中，我们聚焦于LT加速器。