1. 项目概述当ViT遇上存内计算最近几年Vision TransformerViT在计算机视觉领域掀起的浪潮相信大家都有目共睹。从图像分类到目标检测再到图像生成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 TableLUT逼近在注意力分数动态范围很大时误差仍被急剧放大。最终我们坚决地将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 ms3.1 ms~4.9倍每帧图像能耗45.6 mJ6.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芯片设计者而言存内计算提供了一个突破能效瓶颈的崭新路径对于算法工程师而言理解底层硬件的约束将成为设计下一代高效模型的重要素养。