基于SRAM存内计算的Transformer Softmax硬件加速方案解析
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分割如何攻克Softmax
CIMple的核心创新点在于提出了一种针对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. 超越Softmax：CIMple架构的通用化潜力
CIMple虽然针对Softmax优化，但其“基于SRAM的LUT分割计算”范式具有相当的通用性。任何可以通过查找表进行高效近似的非线性函数，理论上都可以用类似的架构加速。这为Transformer乃至更广泛的机器学习模型打开了新的优化空间。
- 激活函数：GELU、Swish、SiLU等Transformer中常用的激活函数，同样是非线性函数，且计算开销不容忽视。可以为这些函数设计专用的CIMple单元，或设计一个可配置的CIMple阵列，通过加载不同的LUT内容来支持多种函数。
- 层归一化（LayerNorm）：LayerNorm涉及计算均值和方差，其中均值的计算是规约操作，方差计算涉及平方和减法，也可以映射到存内计算阵列中进行加速。
- 多项式函数近似：许多复杂的函数（如三角函数、对数函数）在硬件上都可以通过分段多项式来近似。多项式计算可以分解为乘加运算，而乘加运算正是存内计算（尤其是模拟存内计算）擅长的事情。可以设计一个支持通用系数加载的存内乘加阵列。
- 向量-向量逐元素非线性变换：在很多模型中，都存在对特征图的逐元素非线性变换（如乘以一个可学习的缩放因子后加偏置，再经过激活函数）。这种操作数据并行度高，非常适合用CIMple风格的存内计算阵列来处理。
未来的存内计算加速器，可能不再是一个单一的、固定的模块，而是一个可重构的存内计算阵列。它像FPGA一样，可以根据加载的“计算配置文件”（即LUT内容和互联配置），在运行时动态地改变其功能，一会儿加速Softmax，一会儿加速GELU，从而实现硬件资源的最大化利用和极致的能效比。
从我实际做芯片设计的经验来看，CIMple这类工作代表了AI硬件发展的一个清晰趋势：从粗粒度的架构创新（如张量核心），转向更细粒度的、与算法和数据结构深度耦合的电路级创新。它的价值不仅在于提供了一个具体的加速方案，更在于示范了一种设计方法论——如何将算法瓶颈映射到最底层的物理器件特性上，从而获得数量级的效率提升。对于算法工程师和硬件工程师来说，这种跨层次的协同优化思维，变得越来越重要。