可重构计算挑战：在Versal AI引擎上使用Transformer进行喷注标记

# 可重构计算挑战：面向 Versal AI 引擎上喷注标记的 Transformer

来源：https://arxiv.org/html/2606.17500

###### 摘要

基于 Transformer 的模型在 CERN LHC 的喷注标记任务中表现出色，但在低延迟、资源受限的触发系统中部署它们颇具挑战性。我们提出了一种在 AMD Versal AI 引擎 (AIE) 上用于喷注标记的量化、纯整型 Transformer 的初步实现，将密集层和多头注意力 (MHA) 层映射到 AIE 瓦片上。主要贡献是一个可复用的软件框架，该框架将 Transformer 层表示为可组合的 AIE 构建模块，并能从高级 Python 模型描述自动生成相应的 Vitis 图代码。该框架为未来研究奠定了基础，并作为开源软件发布在 https://github.com/KastnerRG/particle_transformer_aie。

## I. 问题与动机

在 CERN 大型强子对撞机 (LHC) 中，质子-质子碰撞以 40 MHz 的频率发生，产生海量的粒子喷注流，必须由第一级触发 (L1T) 系统实时过滤[3 (https://arxiv.org/html/2606.17500#bib.bib1),4 (https://arxiv.org/html/2606.17500#bib.bib3)]。喷注标记是根据其起源物理过程对喷注进行分类的任务。虽然准确的喷注标记对于选择感兴趣的事件至关重要，但延迟和吞吐量必须严格控制在预算内以实现实时处理[4 (https://arxiv.org/html/2606.17500#bib.bib3)]。近期研究表明，基于 Transformer 的架构，例如粒子 Transformer (ParT)，可以直接对每组粒子的特征进行操作，从而实现高精度[8 (https://arxiv.org/html/2606.17500#bib.bib6)]。然而，此类 Transformer 模型是在离线设置的 GPU 上部署的，其中毫秒级的延迟和高功耗是可以接受的。

请参阅图注

图1: LHC 上的触发和推理延迟层级，从 40 MHz 的第一级决策到更高级别的触发和离线分析。

对于在线触发，约束条件要严格得多：L1 触发必须在几微秒内产生接受/拒绝决策，同时将输出率限制在至多 O(10⁵) 个事件每秒，并在严格的探测器上功耗和资源预算下运行[4 (https://arxiv.org/html/2606.17500#bib.bib3),2 (https://arxiv.org/html/2606.17500#bib.bib2)]。这使得在传统 CPU/GPU 平台上部署 Transformer 模型对于实时边缘推理来说不切实际。AMD Versal SoC 上的 AI 引擎 (AIE) 阵列为低延迟、高吞吐量的机器学习推理提供了一个有前景的平台。

在本工作中，我们设计并实现了一个基于 Transformer 的喷注标记加速器，目标平台为 AMD Versal VCK190。我们的主要贡献如下：

- • 我们引入了一个可复用的软件框架，该框架将 Transformer 层表示为可组合的 AIE 构建模块，并能从高级 Python 模型描述自动生成相应的 Vitis 图代码。
- • 我们实现了一个完全量化、纯整型的 Transformer，包括密集层、残差连接和一个纯整型 softmax，针对 AIE 阵列的算术和内存约束进行了定制。
- • 我们提出并评估了一种每头多头注意力 (MHA) 映射方法，将注意力头及其子投影分配给并行的 AIE 瓦片，从而提高了吞吐量。

请参阅图注

图2: 使用 AIEModel 框架从 `model.forward()` 的代码生成与验证流程。

## II. 先前及相关工作

基于 Transformer 的架构，例如粒子 Transformer (ParT)，已在离线分析中树立了喷注标记性能的最新技术水平[8 (https://arxiv.org/html/2606.17500#bib.bib6)]。近期工作已开始探索在实时触发环境中部署 Transformer，包括使用激进量化和专用数据流的亚微秒级 FPGA 实现，用于喷注标记[7 (https://arxiv.org/html/2606.17500#bib.bib7)]。与此同时，一些框架针对 Versal 类架构和 AI 引擎上的 Transformer 或 AI 工作负载，例如 CAT 定制 Transformer 加速器框架[9 (https://arxiv.org/html/2606.17500#bib.bib8)]，以及将高级模型映射到 Ryzen NPU 的 AIE 阵列上的 Vitis ONNX Execution Provider 后端[1 (https://arxiv.org/html/2606.17500#bib.bib11)]。相比之下，CAT 联合利用了 PL 和 AIE 资源，而 ONNX 后端则专注于 Ryzen 设备而非 Versal SoC。我们的工作则使用仅针对 VCK190 上 AIE 的纯整型代码生成框架，来映射高级模型。

## III. 系统概述

我们的工作重点在于构建一个轻量级代码生成框架，将用户用 Python 定义的模型映射到 Versal SoC 的 AIE 上，并在该框架内为喷注标记 Transformer 模型所需的量化计算层开发支持。

### III-A. 目标平台

我们的实现针对 AMD Versal VCK190 SoC，计算完全映射到片上的 AI 引擎 (AIE) 阵列。AIE 瓦片提供了一个空间分布的、VLIW 风格的 SIMD 计算网络，具有本地存储器和高速流互连，我们利用这些特性来实现 Transformer 层。设计使用 AMD/Xilinx Vitis 2024.1 开发。

所有神经网络计算均使用 int8 量化、纯整型算术，搭配 int32 累加器和定点缩放。在底层，我们使用 C/C++ 实现 AIE 内核，并使用 Vitis 数据流图编程模型将它们连接起来，其中内核是节点，AI 引擎瓦片之间的流通道是边。我们主要使用 AIE 硬件仿真来验证设计，这可以对我们生成的图进行周期精确的评估。

### III-B. 代码生成框架

将非平凡模型移植到 AI 引擎上是劳动密集型的。开发者必须手动编写 C/C++ 内核代码以及实例化内核、连接流并在瓦片间分配计算的 Vitis 图描述。当数据流增长到包含多个注意力头、堆叠的注意力块和残差连接时，手动管理端口、FIFO 和瓦片工作负载容易出错，并且调试预期数据流与实际数据流之间的不匹配非常困难。这些挑战促使我们建立一个代码生成框架，将高级模型规范与低级 AIE 图构建分离开来。

尽管一个 Transformer 块和完整模型看起来很复杂，但它们是由一小部分模块化、高度重复的构建模块组成的：密集投影、多头注意力、残差加法和简单的非线性函数。这种规则结构非常适合模板化的代码生成方法，其中每个逻辑层都封装了其数值行为及其硬件映射。

我们的框架暴露了一个面向用户的 Python API，以 `AIEModel` 类为中心，其中模型被定义为一系列 `AIELayer` 实例（例如，`Dense`、`ResAdd`、`MHA`），其风格类似于常见的深度学习库。在构建时，每个 `AIELayer` 消耗一组参数（张量形状、量化缩放因子和分片选择），并发出相应的 Vitis 图片段、内核实例、输入/输出端口和流连接。这些片段随后被组合成一个完整的 AIE 图，自动绑定到一个预编写的 C/C++ 内核库，该库实现了量化的密集层、残差加法和 MHA。为确保数值正确性并简化调试，相同的高级 Python 模型定义也被用来构建基于 NumPy 的 Transformer 参考实现。对于任何给定输入，框架同时运行 NumPy “黄金” 模型和 AIE 实现，并将 AIE 的数值输出与参考 NumPy 实现进行比较。

核心的面向用户入口点是 `model.forward(input)` 调用。调用时，该例程 (i) 为给定的高级模型构建相应的 Vitis AIE 图，(ii) 触发编译和 AIE 硬件仿真以获得定点加速器输出，并且 (iii) 并行构建一个等效的 NumPy 参考模型并对相同输入运行前向传播。仿真完成后，框架将 AIE 输出与 NumPy 参考进行比较，以验证 AIE 设计的数值正确性。图2 (https://arxiv.org/html/2606.17500#S1.F2) 勾勒了此流程，从高级 Python 模型描述到代码生成、编译、AIE 仿真以及基于 NumPy 的数值验证。

该框架有意设计得简单且模块化，因此各个层可以被更优化的内核、替代的分片策略或更丰富的量化方案替换，而无需更改面向用户的模型代码。不同的、中间模型定义可以轻松配置以进行测试和基准测试。

## IV. 架构与 AIE 映射

### IV-A. 喷注标记 Transformer 模型

喷注标记 Transformer 模型的特征和维度推动了我们框架的开发。模型的输入是一个形状为 (160, 8) 的 int8 张量，代表一个经过填充以兼容分片的量化粒子喷注。图3 (https://arxiv.org/html/2606.17500#S4.F3) 展示了网络的设计：一个紧凑的 Transformer 编码器，包含两个堆叠的 4 头自注意力 (MHA) 层，使用了 57 个 AIE 瓦片。每个 Transformer 块由一个 MHA 层后接一个位置级前馈子网络组成，该子网络具有 64 维的隐藏层，使用带有 ReLU 激活函数的密集层实现。在整个模型中，有七个密集层和四个残差加法连接，分别对应于注意力子层、前馈子层以及最终分类头。我们不应用任何归一化层（例如，层归一化或批归一化）。相反，稳定性通过第 IV-C 节 (https://arxiv.org/html/2606.17500#S4.SS3) 描述的量化方案来维持。归一化缩放因子可以在量化之前折叠到上游变换中[5 (https://arxiv.org/html/2606.17500#bib.bib12)]。

请参阅图注

图3: 高级喷注标记 Transformer 模型。

### IV-B. AIE 设计与映射

我们专注于展示 Transformer 到 VCK190 AIE 阵列的功能映射，并演练我们代码生成框架的完整端到端流程。所有计算密集型操作（密集层、多头注意力和残差加法）都实现为一组内核，并映射到独立的 AIE 瓦片上。我们通过头级并行性引入 MHA 层的 AIE 映射。对于每个注意力块，每头的查询、键和值投影被分配到独立的瓦片，以便不同的头子投影可以并行处理。在每个头计算出其贡献后，一小部分拼接内核收集每头的输出，并形成完整的 d_model 维度表示，然后通过一个输出投影内核。MHA 内核的映射如图4 (https://arxiv.org/html/2606.17500#S4.F4) 所示。这种高度并行的设计利用了 VCK190 AIE 的大量瓦片计数，并导致 FIFO 延迟加速，该加速与 MHA 层中定义的头数成正比，如表 I (https://arxiv.org/html/2606.17500#S5.T1) 所示。

密集层实现为单个 AIE 内核，执行 int8 矩阵-矩阵乘法，具有 int32 累加和偏置加法。累加后，重新量化的值被流式传输到下游内核。残差连接实现为独立的 AIE 内核，消耗两个 int8 输入流并执行逐元素加法。

请参阅图注

图4: 用于 MHA 层的代码生成的 Vitis AIE 图。紫色框表示单个计算瓦片，箭头表示单流 int8 数据连接。

### IV-C. 量化与纯整型计算

在我们的 AIE 上高效部署注意力模型需要整数运算。这使得线性层的量化以及诸如 softmax 等非线性运算的实现难以高效执行，因为我们必须控制量化误差，同时符合硬件对整数友好的执行模型。

AIE 上的所有神经网络计算使用对称的 int8 量化，搭配 int32 累加。对于每个权重张量，我们基于在校准过程中观察到的最大绝对值定义一个每张量缩放因子，并将实数值映射到整数范围 [-127, 127]。为了匹配硬件数据路径，有效的实数值重量化因子被实现为二进有理数 M/2ᵇ，其中 M 和 b 是预离线计算的整数。

在运行时，前向传播应用 int32 的整数矩阵乘法，随后进行重量化步骤，该步骤使用 banker's rounding (conv_even) 乘以 M 并右移 b 位。这种设计使所有中间值保持为整数格式，同时紧密逼近标准的浮点推理，从而提高 AIE 上的效率。

Softmax 在此设置中特别具有挑战性，因为指数运算通常以浮点数实现。我们采用 I-BERT [6 (https://arxiv.org/html/2606.17500#bib.bib10)] 中引入的 Int-Softmax 算法的纯整型变体，该算法将指数运算重新表述为一系列能够高效映射到我们整数硬件的定点运算。为了数值稳定性，我们首先计算最大值并将其从每个元素中减去。然后，我们将每个移位后的输入分解为：

xᵢ = -kᵢ ln 2 + rᵢ

其中 kᵢ 是整数商，rᵢ ∈ [-ln 2, 0]。这使得指数可以近似为：

eˣⁱ ≈ 2⁻ᵏⁱ · p(rᵢ),

其中 2⁻ᵏⁱ 实现为位移位，p(rᵢ) 是二阶多项式。所有量（kᵢ、rᵢ、多项式系数以及中间乘积）都以具有共享缩放因子的定点整数格式表示。所得近似引入的最大误差为 1.3×10⁻³，相对于网络的整体量化误差来说很小。

我们的量化和 softmax 设计表明，在 AIE 上实现完全面向整数的注意力堆栈是可行的。然而，我们的整数 softmax 引入了数量级的延迟增加。优化这些实现以获得更高吞吐量是未来工作的主要方向。

## V. 评估与结果

本工作中所有 AIE 实验均使用随机初始化的权重和随机生成的输入进行，旨在作为实现验证研究。量化在测试套件中以两种模式进行：动态模式，其中量化因子通过参考前向传播的校准获得，主要用于具有随机初始化权重和输入的模型；静态模式，其中固定的因子由用户提前选择，以为验证提供数值信息丰富的操作范围。在两种模式下，验证都基于 AIE 实现与相应整数参考计算之间的一致性，如第 IV-C 节所述，并由图 2 中的代码生成和验证流程所示。

我们首先评估了喷注标记 Transformer 的一个简化“骨架”版本。