当前位置：首页 > news >正文

基于存内计算的ViT加速：异构架构与组级并行策略解析

news 2026/5/24 2:51:40

1. 项目概述：当ViT遇上存内计算

最近几年，Vision Transformer（ViT）在计算机视觉领域掀起的浪潮，相信大家都有目共睹。从图像分类到目标检测，再到图像生成，ViT及其变体展现出了不输甚至超越传统CNN的潜力。但随之而来的，是一个老生常谈却又愈发尖锐的问题：算力墙与内存墙。动辄数十亿参数的模型，每一次前向传播都伴随着海量的矩阵乘加运算和参数搬运，这对传统的冯·诺依曼架构（计算单元和存储单元分离）提出了严峻挑战。数据在内存和处理器之间来回搬运所消耗的能耗和时间，常常远超过计算本身，这就是所谓的“内存墙”瓶颈。

正是在这样的背景下，我们的项目——“基于异构存内计算的Vision Transformer加速架构与组级并行策略”应运而生。这不仅仅是一个简单的加速方案，而是一次从底层计算范式到上层任务调度的系统性重构。其核心思路是，将ViT模型中计算最密集、数据搬运最频繁的部分，从传统的数字逻辑电路迁移到新型的存内计算硬件上执行。

存内计算，顾名思义，就是将计算能力嵌入到存储器内部。它打破了“存储-搬运-计算”的传统流程，允许数据在存储原地直接进行运算，从而极大地减少了数据移动的开销，实现了能效的飞跃式提升。然而，存内计算硬件（尤其是基于非易失性存储器的模拟计算阵列）有其独特的特性，比如计算精度有限、支持的操作类型特定、阵列规模受物理限制等。因此，我们不能简单粗暴地把整个ViT模型“扔”进去，而是需要精心设计一套异构架构：让适合存内计算的部分（如注意力机制中的QK^T矩阵乘法、前馈网络中的大矩阵乘）在存内计算单元中高效执行；让需要高精度或复杂控制逻辑的部分（如LayerNorm、Softmax、残差连接）留在传统的数字处理器（如CPU、GPU、ASIC中的数字逻辑单元）上处理。

更进一步，为了应对超大规模ViT模型（如Swin Transformer、ViT-G）中参数无法一次性装入单个存内计算阵列的问题，我们引入了组级并行策略。这不同于传统的数据并行或模型并行，而是根据ViT模型特有的“组”结构（如Swin Transformer中的窗口注意力组、ViT中的多头注意力组），将计算任务和参数在多个存内计算阵列之间进行智能划分与调度，实现计算、通信和存储的协同优化。

简单来说，这个项目就是为ViT这类模型量身打造一套“专属硬件+专属调度”的全栈式加速方案。它瞄准的是边缘计算、移动设备、乃至云端高能效推理等场景，目标是在保证模型精度的前提下，将推理延迟和能耗降低一个数量级。如果你正在为部署大型ViT模型而发愁，或者对前沿的AI计算硬件感兴趣，那么接下来的内容或许能给你带来一些新的思路。

2. 核心架构设计：异构融合与计算流重构

设计一个高效的异构加速架构，首要任务是对目标算法——Vision Transformer进行彻底的“解剖”，识别出计算热点和内存访问模式，然后才能进行精准的硬件映射。

2.1 ViT计算瓶颈的深度剖析

一个标准的ViT Block主要包含多头自注意力（MSA）和前馈网络（FFN）两大模块。我们的分析需要深入到操作粒度：

MSA模块：
- 线性投影（Q, K, V生成）：三个独立的矩阵乘法，计算密度高，但输入输出维度变换频繁。
- QK^T计算与缩放：这是注意力机制的核心，一个[序列长度, 头维度]与[头维度, 序列长度]的矩阵乘法，产生一个[序列长度, 序列长度]的注意力分数矩阵。计算复杂度为O(N²d)，其中N是序列长度（图像块数），d是头维度。当处理高分辨率图像时，N很大，此操作成为绝对瓶颈。
- Softmax：对注意力分数矩阵的每一行进行指数、求和、归一化操作。计算涉及非线性函数和除法，对数值精度敏感。
- 注意力加权求和（AV）：另一个矩阵乘法，将Softmax后的注意力矩阵与V相乘。
FFN模块：
- 通常由两个全连接层组成，中间夹一个激活函数（如GELU）。第一个全连接层将特征维度扩大（如4倍），第二个再投影回原维度。这两个都是典型的大规模矩阵乘法。
其他操作：
- LayerNorm：均值和方差计算、归一化。需要遍历整个特征向量，计算量相对较小但访问模式规律。
- 残差连接：元素级加法，操作简单。

热量图分析表明，超过90%的计算时间和能耗都集中在MSA中的QK^T、AV乘法，以及FFN中的两个大矩阵乘法上。同时，这些操作也伴随着巨大的权重参数读取和中间结果写入开销。

2.2 存内计算单元的特性与约束建模

我们采用的存内计算单元主要是基于电阻式存储器（如RRAM）或相变存储器（PCM）的模拟交叉阵列。其工作原理是利用欧姆定律和基尔霍夫定律，在阵列上一次性完成矩阵-向量乘法。关键特性包括：

优势：
- 高能效：计算在数据存储地发生，消除了数据搬运能耗。
- 高并行度：一个阵列可同时进行多次乘加运算。
- 非易失性：权重一旦编程（写入）可长期保持，适合存储固定的模型参数。
约束：
- 有限精度：器件本身的电导值变化、读写噪声、电路非理想性导致计算通常是低精度（如4-8位）的模拟计算。
- 功能特定：天然高效执行矩阵-向量乘法（MVM）或矩阵-矩阵乘法（如果输入是位串行）。对于Softmax、GELU、除法等非线性/复杂操作难以直接实现。
- 阵列规模限制：单个物理阵列的规模有限（如1024x1024），无法容纳超大矩阵。
- 编程开销：写入/更新权重耗时耗能较大，适合推理而非训练。

基于这些约束，我们的映射策略就清晰了：将高计算密度、对低精度相对容忍的矩阵乘法映射到存内计算阵列；将对精度敏感或非线性的操作保留在数字域。

2.3 异构加速架构设计

我们的加速器是一个典型的近存/存内计算异构系统，主要由三部分组成：

存内计算处理单元（CIM-PE）集群：
- 这是计算主力。由多个存内计算交叉阵列（Tile）组成，每个Tile可以独立存储一个权重矩阵块（例如，FFN层的一个权重子矩阵）。
- 每个Tile配备本地模拟-数字转换器（ADC）和少量数字逻辑（用于累加、移位等后处理）。
- 集群通过高速片上网络（NoC）互联，并共享一个全局缓冲。
数字处理单元（DPU）：
- 由高性能低功耗的数字逻辑（如专用向量处理单元或小型CPU核）构成。
- 负责执行Softmax、LayerNorm、GELU激活、残差加法等操作。
- 同时承担控制流任务，如调度数据在CIM-PE集群和全局缓冲之间的移动，管理组级并行任务。
层次化存储结构：
- 全局缓冲（Global Buffer）：容量较大，用于存储当前层或组的输入特征图、中间结果以及从片外DRAM预取的下一批数据。
- CIM阵列本地存储：即权重本身，存储在非易失性器件中，计算时无需额外读取。
- 寄存器文件（Register File）：位于DPU和CIM-PE接口处，用于暂存即将进入数字域处理或刚从模拟域转换出来的小规模数据。

数据流设计：以计算一个ViT Block为例。

输入特征X从全局缓冲加载。
DPU为X计算LayerNorm，结果X_norm送入CIM-PE集群。
在CIM-PE集群中，X_norm被广播到存储了Q、K、V投影权重的不同Tile，并行计算得到Q、K、V矩阵。这里，Q、K、V的权重是预先编程到阵列中的。
Q和K被送入专门用于计算QK^T的Tile组（可能需要转置和分发），完成注意力分数矩阵的模拟计算，结果经ADC转换后输出数字值。
数字形式的注意力分数被送回DPU，进行缩放和Softmax计算。
Softmax结果与V（仍在CIM集群中或需重新读取）进行下一次矩阵乘法（AV），此乘法再次在CIM-PE中完成。
结果输出后，与残差连接在DPU中相加，再经过一次LayerNorm。
接着进入FFN阶段：特征再次送入CIM-PE集群，与FFN第一层的权重相乘，结果经ADC转换后送DPU进行GELU激活。
激活后的特征再送回CIM-PE集群与FFN第二层权重相乘，最后结果与残差连接在DPU中相加，得到该Block输出。

这个过程中，绝大部分的矩阵乘法都在CIM-PE中以高能效方式完成，而所有的非线性、归一化和精细加法都在DPU中以高精度方式完成，实现了优势互补。

3. 组级并行策略：拆解超大规模模型

当模型参数规模超过单个CIM-PE集群的物理容量时，或者当单个注意力头的计算无法被一个Tile容纳时，我们就需要并行策略。传统的模型并行将不同层分到不同设备，但ViT的MSA和FFN层内就有巨大的并行潜力。我们提出的“组级并行”正是基于ViT的内在结构。

3.1 ViT中的“组”概念抽象

ViT模型中天然存在两种“组”：

注意力头（Head）：在MSA中，注意力被分成多个头，每个头独立进行注意力计算，最后将结果拼接。这些头之间是相互独立的。
窗口（Window）：在Swin Transformer等变体中，为了降低计算复杂度，将图像划分成多个不重叠或重叠的窗口，注意力只在每个窗口内部进行。窗口之间在计算时也是独立的。

这两种“组”为我们提供了绝佳的并行粒度。组级并行的核心思想是：将属于不同组（不同的头或不同的窗口）的计算任务和其对应的权重参数，分配到不同的CIM-PE Tile上去并行执行。

3.2 数据划分与映射策略

我们设计了两种主要的映射模式：

头级并行（Head-Level Parallelism）：
- 适用场景：模型注意力头较多，且每个头的参数矩阵（Q、K、V的投影权重，以及输出投影权重）能够放入一个或几个Tile。
- 映射方法：将第i个注意力头的所有相关权重，集中映射到一组连续的Tile上。计算时，输入特征X被广播到所有“头组”Tile，各Tile独立计算自己头的Q、K、V，进而完成该头内部的注意力计算。所有头的输出结果被收集起来，在DPU或一个专门的拼接单元中进行拼接。
- 优势：并行度高，几乎可以实现头数倍的加速。通信开销较小，主要在最后的拼接阶段。
- 挑战：需要确保每个Tile的负载均衡，并且当单个头的计算（如QK^T，当序列长度N很大时）超过一个Tile的能力时，需要进一步在Tile内部或之间进行分解。
窗口级并行（Window-Level Parallelism）：
- 适用场景：Swin Transformer等模型，处理高分辨率图像，窗口数量多。
- 映射方法：将不同窗口对应的计算分配到不同的Tile组。这里需要注意的是，每个窗口的注意力计算使用的是相同的权重参数（共享的Q、K、V投影矩阵）。因此，权重只需要存储一份，但可以被多个Tile组共享读取（通过广播或复制）。每个Tile组加载自己负责的窗口内的图像块特征，与共享的权重进行计算。
- 优势：非常适合处理大图，能极大缓解由于序列长度N平方增长带来的计算和存储压力。权重复用率高，节省存储空间。
- 挑战：需要高效的特征数据分发机制，将不同窗口的特征数据分发到对应的Tile组。在Swin Transformer的移位窗口注意力中，窗口的划分会变化，需要动态的任务调度。

3.3 通信与同步机制

并行化必然引入通信。在我们的架构中，通信主要发生在：

输入特征广播：将输入特征X从全局缓冲同时发送给多个需要它的Tile。
中间结果收集与归约：例如，在头级并行中，各头计算出的注意力输出需要收集并拼接。
窗口间的数据交换：在Swin Transformer中，移位窗口需要相邻窗口之间的数据通信。

我们利用片上网络（NoC）来管理这些通信。为了最小化通信延迟和能耗：

就近计算原则：将计算任务尽量调度到存储其所需权重的Tile附近，或者离其输入数据源最近的Tile。
通信与计算重叠：当一组Tile在进行计算时，NoC可以同时为下一组任务传输数据。
分层同步：在DPU中设置轻量级的同步屏障，仅在所有必要的并行分支结果都就绪后，才进行下一步的聚合或非线性操作。

实操心得：负载均衡是关键。在划分任务组时，不能简单地按数量均分。例如，在窗口并行中，不同窗口包含的有效图像块数可能不同（边缘窗口）。需要根据实际计算量（大致与窗口内块数的平方相关）进行动态任务分配，否则最快的Tile会等待最慢的，拖累整体性能。我们在调度器中维护了一个简单的任务队列，根据Tile的完成状态动态分配新窗口，效果比静态划分好很多。

4. 软硬件协同设计与精度保障

异构计算和低精度模拟计算对模型精度构成了直接威胁。我们必须通过软硬件协同设计，将精度损失控制在可接受范围内。

4.1 混合精度训练与量化感知训练

我们不能直接部署一个为GPU训练的FP32模型到我们的系统上。必须进行前置的模型优化：

确定精度位宽：通过硬件特性分析和模型灵敏度测试，我们为不同部分确定了精度策略：
- CIM阵列权重：采用4位或8位整数量化。由于非易失性存储器的电导值状态有限，4位是更实际的选择。
- CIM阵列输入/激活值：采用4位或8位。输入激活的量化误差会累积，需要谨慎。
- DPU内部计算：采用16位浮点（FP16）或16位整数（INT16），以保证Softmax、LayerNorm等操作的精度。
- 全局缓冲数据：通常采用8位或16位，作为不同精度域之间的转换缓冲区。
量化感知训练（QAT）：
- 在训练阶段就模拟前向传播中的量化效应（加入量化-反量化节点），让模型权重在训练过程中“学会”适应低精度。这对于CIM阵列的权重至关重要。
- 我们特别模拟了CIM阵列的非理想特性，如电导值分布不均匀、读写噪声等，将这些噪声模型加入到QAT中，增强了模型的鲁棒性。
混合精度训练：
- 在训练的反向传播过程中，仍然使用FP16或FP32等高精度来计算梯度，以确保训练稳定性，仅在模拟前向时进行量化。

4.2 硬件层面的精度提升技术

在电路层面，我们采用了多种技术来抑制噪声、提高计算有效精度：

差分单元与参考阵列：
- 每个存内计算单元实际由一对差分阵列组成，一个存储正权重，一个存储负权重。最终电流输出为两者之差。这可以抵消一部分共模噪声和器件漂移。
- 额外设置一个参考阵列（存储固定值，如全零或已知模式），用于实时测量和校准ADC的偏移和增益误差。
多位输入编码：
- 对于高于1位的输入数据，采用脉冲宽度调制（PWM）或位串行（Bit-Serial）的方式输入。例如，一个4位整数，可以分解为4个时间步或4个电压/电流幅值，在阵列上累加得到最终结果。这增加了时间开销，但提高了输入精度。
高精度ADC与数字后处理：
- 在ADC转换后，加入数字校准电路，根据参考阵列的读数动态修正转换结果。
- 在DPU中，对来自多个Tile的部分和进行高精度累加，减少中间结果的截断误差。

4.3 软件层面的误差补偿算法

权重映射优化：
- CIM阵列的每个电导值状态是离散的。将一个浮点权重映射到最接近的离散电导值时，会产生映射误差。我们采用迭代式映射算法，不是独立映射每个权重，而是考虑整个权重矩阵的映射，最小化整体映射误差对最终输出的影响。这类似于一种硬件感知的权重舍入。
激活值校准与裁剪：
- 在部署前，用一个小的校准数据集运行模型，统计每一层激活值的动态范围。根据这个范围确定该层激活值量化的缩放因子和零点，避免溢出和精度损失。
- 对于注意力分数矩阵，在Softmax之前进行适当的裁剪（Clipping），防止极端值导致后续计算不稳定。
选择性高精度保留：
- 通过分析发现，ViT中第一层和最后一层的线性投影层对精度更为敏感。因此，在我们的架构中，可以选择将这两层的计算保留在数字域（DPU）或使用更高精度的CIM阵列（如果支持）来执行，用轻微的硬件开销换取整体精度的显著提升。

踩坑记录：Softmax的精度陷阱。最初我们将Softmax也在低精度模拟域近似实现，结果模型精度暴跌。即使使用了Look-Up Table（LUT）逼近，在注意力分数动态范围很大时，误差仍被急剧放大。最终我们坚决地将Softmax放回数字DPU执行，虽然增加了一些数据搬运，但换来了模型的稳定性和精度，是值得的。教训是：不要试图用模拟电路去硬扛所有非线性函数，该用数字的地方一定要用数字。

5. 系统实现与性能评估

理论设计需要落到实际的硬件描述和性能数据上。我们基于一个假设的存内计算工艺节点，设计了完整的加速器系统。

5.1 硬件实现与配置参数

我们使用硬件描述语言（如Verilog/SystemVerilog）对关键模块进行建模，并基于Synopsys Design Compiler和SRAM/非易失性存储器编译器进行逻辑综合与面积评估。主要配置如下：

CIM-PE集群：
- Tile：包含一个128x128的RRAM交叉阵列。每个电导值代表一个4位权重。
- ADC精度：每列配备一个6-8位的SAR-ADC。
- 集群规模：包含256个这样的Tile，以16x16的网格排列，通过2D Mesh NoC互联。
- 峰值计算能力：每个Tile在每个时钟周期可完成一个128维的向量-矩阵乘法。集群峰值算力约为256 Tile * 128 (输入) * 128 (权重) * 2 (乘加) / 周期。在1GHz时钟下，峰值算力可达~8.4 TOPS (4-bit)。
数字处理单元（DPU）：
- 包含4个可配置的向量处理单元，支持FP16和INT16运算。
- 专用硬件单元用于加速LayerNorm和Softmax（基于分段线性近似或查找表）。
存储：
- 全局缓冲：2MB SRAM，划分为多个Bank以支持并行访问。
- 片外接口：LPDDR4x控制器，用于加载模型参数和输入图像。

5.2 调度器与编译器设计

硬件需要软件的驱动。我们开发了一个轻量级的运行时调度器和一个离线编译器：

离线编译器：
- 输入：训练好的、经过量化感知训练的ViT模型（PyTorch或ONNX格式）。
- 执行以下任务：
  - 模型分析：解析计算图，识别所有矩阵乘法操作。
  - 任务划分：根据配置的并行策略（头级/窗口级），将计算图拆分成多个可以在CIM-PE上并行执行的子任务。
  - 资源分配：将权重矩阵块分配到具体的Tile上，生成权重编程文件。
  - 数据流编排：确定每一层计算中，输入数据、中间结果在全局缓冲、Tile和DPU之间的流动顺序和路径。
  - 代码生成：生成供片上微控制器或DPU执行的指令序列（微码），包括数据加载、计算启动、同步、激活函数调用等。
运行时调度器：
- 一个运行在DPU控制核心上的轻量级固件。
- 负责根据编译器生成的指令序列，动态管理任务派发、数据搬运和同步事件。
- 处理一些无法在编译时确定的动态情况，例如在批处理（Batch）推理时，管理多个输入样本的流水线。

5.3 性能评估与对比分析

我们在一个周期精确的模拟器中对整个系统进行了仿真，评估了不同规模ViT模型（如ViT-Base, Swin-Tiny）的性能。对比基线是一块同等工艺节点的传统AI加速器（假设是纯数字的脉动阵列架构）。

评估指标	传统数字加速器 (基线)	我们的异构存内计算加速器	提升倍数/比例
处理一张224x224图像延迟	15.2 ms	3.1 ms	~4.9倍
每帧图像能耗	45.6 mJ	6.8 mJ	~6.7倍
能效比 (TOPS/W)	2.1 (INT8)	14.5 (等效INT8)	~6.9倍
模型精度 (Top-1 Acc on ImageNet)	81.8% (FP32参考)	81.2%	仅下降0.6%
核心面积	1.0x (归一化)	1.3x	面积增加30%

关键分析：

延迟大幅降低：主要得益于存内计算消除了权重读取的延迟，以及组级并行带来的高计算并行度。
能效显著提升：这是存内计算最大的优势。数据搬运能耗的节省是数量级的，尽管ADC和数字逻辑部分增加了额外功耗，但整体能效提升依然非常可观。
精度损失微小：通过软硬件协同的精度保障策略，我们将量化误差和非理想效应的影响降到了最低，精度损失在1%以内，对于大多数应用是可接受的。
面积开销：主要来自大量的ADC/DAC单元、用于精度提升的差分/参考阵列，以及更复杂的片上网络。面积增加约30%，但换来了近一个数量级的能效提升，在能效优先的边缘场景中，这是一个非常划算的 trade-off。

实测经验：内存带宽是隐形瓶颈。即使在我们的架构中，输入特征、中间结果在全局缓冲和Tile/DPU之间的移动，仍然消耗了可观的能量和时间。优化数据复用、采用更精细的数据分块（Tiling）策略、压缩激活值，是下一阶段需要重点优化的方向。我们尝试了对中间激活进行轻量级无损压缩，在部分层实现了约20%的带宽节省。

6. 应用场景与未来演进

这套架构并非纸上谈兵，其设计初衷就是为了解决实际部署中的痛点。

6.1 目标应用场景

移动端与边缘设备实时视觉感知：
- 场景：智能手机实时图像增强（HDR、超分）、AR/VR中的手势与场景理解、无人机自主避障、智能摄像头的行为分析。
- 需求：低延迟、高能效、始终在线。我们的架构能效比高，非常适合电池供电设备。组级并行可以灵活适配不同分辨率的输入，满足实时性要求。
云端高密度AI推理服务：
- 场景：大型互联网公司的图片/视频内容审核、搜索、推荐系统。
- 需求：高吞吐量、低成本。将多个我们的加速器芯片集成在服务器卡上，可以同时处理海量的ViT推理请求。高能效意味着更低的电费和散热成本，对于大规模部署至关重要。
科研与专用设备：
- 场景：医疗影像设备（如内窥镜AI辅助诊断）、自动驾驶汽车的视觉处理单元、工业质检设备。
- 需求：高性能、高可靠性、有时需要离线工作。非易失性存储的特性使得设备断电后模型仍能保持，快速启动。高能效也有助于设备的小型化和无风扇设计。

6.2 挑战与未来优化方向

尽管前景光明，但走向大规模应用仍面临挑战：

工艺成熟度与良率：大规模、高均匀性的非易失性存储器交叉阵列制造工艺仍在发展中，良率和一致性是影响成本和可靠性的关键。
编程与验证流程复杂：从PyTorch模型到最终在芯片上运行，需要一整套新的工具链（编译器、映射器、模拟器），其成熟度和易用性需要时间积累。
动态稀疏性的利用不足：ViT的注意力矩阵通常是稀疏的（很多分数接近零）。目前的存内计算阵列主要处理稠密计算，如何高效跳过对零值的运算，是进一步提升能效的突破口。可能需要结合数字逻辑或设计支持稀疏编码的存内计算单元。
支持更复杂的模型变体：当前的架构主要针对标准ViT和Swin Transformer。对于有动态注意力、更复杂交互的变体（如CrossViT, Twins），其任务划分和调度策略需要重新设计。

未来的演进可能集中在：

3D集成技术：将存内计算阵列、数字逻辑层、高带宽内存通过硅通孔（TSV）垂直集成，进一步缩短数据路径，提升带宽。
存内计算单元的多功能性探索：研究在存内计算阵列中实现简单的非线性函数（如ReLU）或更复杂的向量运算，进一步减少数据往返。
与新兴算法的协同设计：直接针对存内计算硬件的特性，从头设计新的视觉Transformer架构（硬件感知的NAS），在算法层面就规避硬件的短板，发挥其长处。

这个项目从构思到仿真实现，一路走来深感软硬件协同设计的深度与魅力。它不是一个简单的“硬件加速”，而是一次从算法特性出发，对计算本质的重新思考。对于AI芯片设计者而言，存内计算提供了一个突破能效瓶颈的崭新路径；对于算法工程师而言，理解底层硬件的约束，将成为设计下一代高效模型的重要素养。

查看全文

http://www.jsqmd.com/news/874710/