1. 项目背景当Transformer的Softmax成为算力瓶颈最近在优化一个部署在边缘设备上的Transformer模型时我遇到了一个棘手的问题推理速度始终上不去。经过Profiling分析发现一个反直觉的现象——在注意力机制的计算中耗时最长的并非我们通常认为的矩阵乘法MatMul而是那个看起来“人畜无害”的Softmax归一化操作。尤其是在处理长序列比如长文本或高分辨率图像分块时计算QK^T后得到的庞大得分矩阵其Softmax运算带来的数据搬运和指数计算开销竟然成了整个推理流程的“阿喀琉斯之踵”。这促使我开始深入探究Softmax的硬件加速方案。传统的做法无非两种一是算法层面近似用ReLU、线性注意力等替换二是在通用处理器CPU/GPU上做算子优化。但前者往往牺牲模型精度后者则受限于冯·诺依曼架构固有的“内存墙”问题——数据需要在计算单元和存储单元之间来回搬运能效比极低。直到我看到了“CIMple”这篇论文的思路它提出了一种基于标准单元SRAM的存内计算架构专门用于加速Transformer中的Softmax尤其是通过LUT查找表分割技术来应对高动态范围的输入。这个设计让我眼前一亮因为它直接从最底层的存储器件入手试图在数据存储的地方完成计算从而从根本上规避数据搬运的开销。今天我就结合自己的理解来拆解一下CIMple架构的核心思想、实现难点以及它给边缘AI计算带来的可能性。2. 理解存内计算为什么是SRAM为什么是现在在深入CIMple之前我们必须先搞清楚存内计算到底是什么以及为什么SRAM是当前最可行的载体之一。存内计算的核心思想是“计算发生在数据存储的地方”这打破了冯·诺依曼体系中计算和存储分离的范式。其最大的优势就是能显著减少数据在存储器和处理器之间的移动从而大幅降低功耗并提升能效这对于电池供电的物联网设备至关重要。那么为什么选择SRAM而不是其他存储器呢这需要对比几种主流存储介质DRAM密度高、成本低但读写速度相对慢且需要定时刷新存内计算电路设计复杂干扰大。Flash非易失性密度极高但写入速度慢、寿命有限且模拟计算精度受制于器件耐久性。RRAM/MRAM等新型存储器具有非易失、高密度潜力是研究热点但工艺尚未完全成熟器件一致性、良率仍是产业化的挑战。相比之下SRAM作为CPU缓存的核心拥有无与伦比的速度和可靠性。它采用标准的CMOS工艺制造与逻辑电路兼容性极好易于集成在SoC中。更重要的是SRAM的存储单元6T或8T晶体管结构规整便于在其阵列周边或内部嵌入模拟计算电路。虽然SRAM的密度不如DRAM和Flash且是易失性的但对于需要频繁、高速访问的片上缓存或专用加速器缓冲区来说它的性能优势是决定性的。CIMple选择基于标准单元SRAM来设计意味着它可以直接利用现有成熟、高可靠性的半导体工艺流片降低了技术落地门槛。存内计算的具体实现方式主要分数字和模拟两种。数字存内计算通常在SRAM阵列中嵌入额外的逻辑门在读出数据的同时进行简单的位运算如与、或、加法。而模拟存内计算则利用存储器件的电学特性如电流、电压、电荷来直接进行模拟量的乘加运算能效比潜力更高。CIMple针对Softmax的加速更倾向于一种混合或数字化的方案因为它需要高精度的非线性函数计算这对纯模拟计算是一个挑战。3. Softmax的硬件之痛指数运算与高动态范围要理解CIMple的价值必须深刻理解Softmax在硬件实现上的难点。Softmax函数的公式很简单$Softmax(x_i) \frac{e^{x_i}}{\sum_{j}e^{x_j}}$。但在硬件上它至少带来三大挑战指数运算e^x的非线性与复杂性指数函数没有简单的数字逻辑电路可以直接实现。在通用处理器上通常采用查找表LUT结合多项式近似的方法来计算。但这需要多次内存访问和算术运算消耗大量时钟周期和能量。高动态范围的输入值在Transformer注意力中$x_i$ 是QK^T的结果其值域可能非常广。当$x_i$很大时$e^{x_i}$会溢出超出浮点数表示范围当$x_i$很小时$e^{x_i}$会下溢为零。为了防止溢出常规做法是减去最大值$x_i - max(x)$。但这第一步“求最大值”就是一个需要遍历整个向量的规约操作增加了数据依赖和延迟。除法与求和分母是所有元素的指数和这引入了全局数据依赖。必须先计算完所有$e^{x_j}$并求和才能进行除法。这限制了计算的并行性并且除法运算本身在硬件上也是比较昂贵的操作。在传统的处理器架构中这些操作意味着反复将张量从缓存SRAM加载到寄存器堆再送入ALU运算结果写回。数据像钟摆一样在存储和计算单元间摆动产生了巨大的“搬运功耗”。CIMple的思路就是能否在数据存放的SRAM阵列里直接完成这一系列操作的核心部分4. CIMple架构深度拆解LUT分割如何攻克SoftmaxCIMple的核心创新点在于提出了一种针对Softmax的、基于SRAM的存内计算架构其关键是一种称为“LUT分割”的技术。下面我们来一步步拆解它的工作原理。4.1 整体计算流程映射首先CIMple并不是要在存内一步到位算出完整的Softmax。更务实的思路是将Softmax计算中瓶颈最突出、最适合存内计算的部分剥离出来。通常这个部分就是指数运算 $e^{x}$ 的近似计算。CIMple将经过最大值规约后的输入向量 $x_i$ 值域被限定在非正数例如[ -B, 0 ]直接送入SRAM存内计算阵列。这个SRAM阵列被设计成不仅能存储数据还能根据地址即输入值$x_i$查表输出一个对应的近似值 $f(x_i) \approx e^{x_i}$。也就是说SRAM阵列在这里扮演了一个巨大的、分布式的查找表角色。4.2 LUT分割的精髓精度、面积与功耗的平衡如果用一个单一的、高精度的查找表来覆盖整个输入值域那么这个表会非常庞大。例如如果输入是16位定点数一个完整的LUT需要 $2^{16}$ 个表项每个表项存储一个输出值比如16位这将消耗巨大的芯片面积SRAM位单元数量和静态功耗完全不现实。CIMple的“LUT分割”技术巧妙地解决了这个问题。其核心思想是将输入值域分割成若干个子区间每个子区间对应一个较小精度的LUT再通过一个共享的、高精度的“基准表”和简单的插值逻辑来重建高精度输出。具体来说分割输入值域将输入 $x$ 的二进制表示分成两部分高位部分MSBs和低位部分LSBs。高位部分用于选择“子LUT”或“区间”低位部分用于在该区间内进行精细查找或插值。基准表与差值表基准表存储每个区间起点由高位部分决定对应的高精度函数值 $f(x_{base})$。这个表较小因为区间数量远少于总输入状态数。差值表/子LUT存储每个区间内基于低位部分的函数值增量或一个低精度的完整映射。由于每个区间覆盖的范围小函数变化相对平缓因此可以用很低的精度较少的比特位来表示这个增量或映射从而极大地压缩了表格大小。重建输出最终输出通过组合基准表输出和差值表输出来获得。例如$f(x) \approx f(x_{base}) \Delta f(x_{LSB})$。这里的加法操作可以利用存内计算阵列中嵌入的简单加法逻辑来实现。通过这种分割CIMple用“一小块高精度内存多块低精度内存”的组合实现了接近全精度LUT的效果同时在面积和功耗上取得了极佳的平衡。这正是其能实用化的关键。4.3 SRAM阵列的微架构实现那么这种LUT如何在物理的SRAM阵列中实现呢CIMple likely采用了一种数字存内计算的思路。我们可以想象对传统的SRAM阵列进行改造字线WL作为输入输入值 $x$或其高位部分被解码为SRAM阵列的地址激活某一条或一组字线。位线BL输出计算结果在每条位线上通过精心设计的感测放大器Sense Amplifier或额外的在位线末端/内部的模拟-数字混合电路来生成对应的LUT输出值。对于分割LUT可能需要同时激活多条字线对应基准值和差值并在位线层面进行电流或电压的叠加模拟加法操作。嵌入轻量级逻辑在SRAM阵列的外围或子阵列之间加入必要的控制逻辑和简单的算术单元如加法器用于完成插值或后续的累加、归一化步骤。这种设计使得一次SRAM读操作不仅能读出数据还能“读出”一个经过计算的结果实现了真正的“存算一体”。5. 从单元到系统CIMple如何集成并加速Transformer注意力单个存内计算单元解决了指数近似问题但要加速完整的注意力模块还需要系统级的架构设计。5.1 注意力计算的数据流重构标准注意力计算流程为$Attention(Q, K, V) Softmax(\frac{QK^T}{\sqrt{d_k}})V$。CIMple的介入点通常在计算 $S QK^T$ 之后。假设我们有一个存内计算单元阵列CIM Array其设计流程可能如下行最大值规约计算得分矩阵 $S$ 每一行的最大值。这个规约操作本身也可以通过存内计算的方式加速例如利用SRAM阵列内的比较逻辑树。减最大值与值域限定将 $S$ 的每个元素减去其所在行的最大值得到 $S‘$值域被限定在-∞ 0]。这一步通常需要在数字逻辑中完成因为涉及减法。存内指数近似将 $S‘$ 的各个元素广播或分发到CIMple阵列的不同单元。每个CIMple单元接收一个标量输入 $s‘{ij}$并通过其内部的LUT分割机制输出对应的近似值 $exp_approx(s‘{ij})$。这一步是性能提升的关键它避免了将庞大的 $S‘$ 矩阵搬出内存进行指数计算。行求和规约对每一行近似指数结果进行求和得到分母 $sum_i$。这同样可以利用存内计算阵列内的加法树结构来实现结果存回存储体或寄存器。除法归一化最后将每一行的每个 $exp_approx(s‘_{ij})$ 除以该行的 $sum_i$得到注意力权重。除法运算可能仍需在传统的数字逻辑单元如专用除法器或DSP中完成但此时数据量已经过指数压缩和规约搬运和计算开销大大减小。5.2 与现有计算单元的协同CIMple不是一个孤立的加速器它需要与现有的计算体系如向量处理器、张量核心紧密协同。一种可能的集成方式是作为现有AI加速器如NPU的“智能内存”或“近内存计算”模块。例如在NPU的片上缓存通常是SRAM中划出一部分区域采用CIMple结构进行设计。当NPU的矩阵乘法单元计算出 $QK^T$ 后直接将结果写入这片特殊的SRAM区域。随后由CIMple逻辑控制完成Softmax的核心计算计算结果再被NPU读回用于与V矩阵相乘。这种集成方式最小化了数据移动实现了“计算向数据靠拢”。它本质上是一种异构计算用定制化的存内计算单元处理特定、高开销的算子Softmax而通用的矩阵乘等操作则由高度优化的张量核心处理各司其职。5.3 精度与能效评估任何近似计算硬件都必须面对精度问题。CIMple的LUT分割方法通过分段线性插值或高阶近似可以将计算误差控制在很低的水平。论文中应该会通过定点数量化、误差分析和实验证明在8位或16位精度下其输出的注意力权重与浮点软件计算结果之间的差异对于模型整体精度如BLEU分数、图像分类准确率的影响可以忽略不计。能效提升则是更显著的亮点。由于避免了将中间张量在存储层级间反复搬运尤其是避免了高功耗的指数函数在通用ALU上的迭代计算CIMple有望将Softmax操作的能效提升一个数量级以上。这对于手机、AR眼镜等设备的长续航至关重要。6. 实战思考设计自己的存内计算加速单元阅读论文固然重要但将其思想转化为实际的设计考量更有价值。如果你要为一个特定的神经网络算子设计存内计算加速单元应该如何思考以Softmax为例算子剖析与热点定位首先使用性能分析工具如NSight VTune或周期精确模拟器精确量化目标模型中该算子的执行时间、能耗占比和数据移动开销。确认Softmax确实是瓶颈。计算模式抽象将Softmax分解为基本操作序列寻最大值、减法、指数、求和、除法。分析哪些操作是数据并行的如每个元素的指数运算哪些是规约的求最大、求和哪些是全局依赖的除法。数据并行部分最适合存内计算。数值范围分析与定点化确定输入$QK^T$的典型值域。通过统计大量推理数据确定最大值规约后的输入范围如[-10, 0]。这个范围决定了LUT需要覆盖的区间直接影响硬件复杂度。接着需要做定点量化确定整数和小数部分的位宽平衡精度和硬件成本。LUT分割方案选型等间隔分割 vs. 非等间隔分割等间隔设计简单但可能在某些函数变化剧烈的区间精度不够。非等间隔如基于函数二阶导数可以用更少的区间达到相同精度但寻址逻辑更复杂。线性插值 vs. 多项式插值线性插值只需存储区间起点值和终点值或斜率硬件实现最简单一个乘法器和一个加法器。二次或三次插值精度更高但需要存储更多系数计算逻辑也更复杂。基准表与差值表的位宽分配需要通过误差仿真权衡基准表位宽和差值表位宽。通常基准表需要较高精度如12-16位而差值表由于区间内变化小4-8位可能就足够了。SRAM电路级设计这是最挑战的部分。需要考虑如何修改标准6T SRAM单元或引入8T、10T单元来支持计算功能。是在读放电路Sense Amplifier上做文章使其能输出模拟电流和并进行叠加还是在位线之间加入模拟乘法器亦或是采用完全数字化的方案将SRAM阵列作为真值表配合多路选择器和加法器这需要深厚的电路设计知识并与工艺厂紧密合作。系统集成与验证设计完成后需要在SoC层面集成。考虑总线接口AMBA AXI、与主处理器或NPU的通信协议、数据搬运DMA的设计等。最后必须进行完整的验证从RTL功能仿真到FPGA原型验证再到ASIC流片后的硅后测试并使用真实的Transformer模型如BERT ViT进行端到端的精度和性能评估。注意存内计算电路设计尤其是模拟存内计算对工艺波动PVT非常敏感。晶体管参数的微小差异可能导致计算结果的偏差。因此片上校准电路、冗余设计、误差校正编码等技术至关重要这是在设计初期就必须考虑的。7. 超越SoftmaxCIMple架构的通用化潜力CIMple虽然针对Softmax优化但其“基于SRAM的LUT分割计算”范式具有相当的通用性。任何可以通过查找表进行高效近似的非线性函数理论上都可以用类似的架构加速。这为Transformer乃至更广泛的机器学习模型打开了新的优化空间。激活函数GELU、Swish、SiLU等Transformer中常用的激活函数同样是非线性函数且计算开销不容忽视。可以为这些函数设计专用的CIMple单元或设计一个可配置的CIMple阵列通过加载不同的LUT内容来支持多种函数。层归一化LayerNormLayerNorm涉及计算均值和方差其中均值的计算是规约操作方差计算涉及平方和减法也可以映射到存内计算阵列中进行加速。多项式函数近似许多复杂的函数如三角函数、对数函数在硬件上都可以通过分段多项式来近似。多项式计算可以分解为乘加运算而乘加运算正是存内计算尤其是模拟存内计算擅长的事情。可以设计一个支持通用系数加载的存内乘加阵列。向量-向量逐元素非线性变换在很多模型中都存在对特征图的逐元素非线性变换如乘以一个可学习的缩放因子后加偏置再经过激活函数。这种操作数据并行度高非常适合用CIMple风格的存内计算阵列来处理。未来的存内计算加速器可能不再是一个单一的、固定的模块而是一个可重构的存内计算阵列。它像FPGA一样可以根据加载的“计算配置文件”即LUT内容和互联配置在运行时动态地改变其功能一会儿加速Softmax一会儿加速GELU从而实现硬件资源的最大化利用和极致的能效比。从我实际做芯片设计的经验来看CIMple这类工作代表了AI硬件发展的一个清晰趋势从粗粒度的架构创新如张量核心转向更细粒度的、与算法和数据结构深度耦合的电路级创新。它的价值不仅在于提供了一个具体的加速方案更在于示范了一种设计方法论——如何将算法瓶颈映射到最底层的物理器件特性上从而获得数量级的效率提升。对于算法工程师和硬件工程师来说这种跨层次的协同优化思维变得越来越重要。