1. 项目概述当Transformer遇上存内计算Softmax为何成了“拦路虎”如果你最近关注过大语言模型LLM的硬件加速尤其是存内计算In-Memory Computing, IMC这个热门方向那你肯定听过一个说法Transformer模型里的注意力机制Attention是性能瓶颈。但更具体一点这个瓶颈的“罪魁祸首”往往指向一个看似简单的数学函数——Softmax。今天我想从一个硬件设计者的角度跟你聊聊为什么Softmax在IMC架构上这么“难搞”以及我们是如何通过一套软硬件协同设计的“组合拳”把这个瓶颈给打通的。简单来说Softmax是注意力机制里的“裁判”。它负责把一堆原始的注意力分数比如查询向量和键向量的点积结果归一化成概率分布决定模型应该“看”输入序列的哪些部分。问题就出在这个“归一化”上。传统的数字电路实现Softmax需要先对所有元素求指数exp再求和最后做除法。这个过程涉及大量的数据搬运把数据从内存搬到计算单元和复杂的非线性运算在序列长度N动辄成千上万的LLM场景下能耗和延迟都非常吓人。而IMC的初衷是把计算放到内存阵列里做减少数据搬运但Softmax这种需要全局信息的操作天然和IMC“计算局部化”的特性相悖。我们提出的方案核心思想是“算法重构硬件适配”。不是硬着头皮在IMC上跑传统的Softmax而是从算法层面改造注意力计算把计算负担从庞大的序列维度N转移到相对较小的嵌入维度D。这样一来核心的Softmax-like操作就能被巧妙地映射到IMC阵列或者专用的模拟/数字电路上。我们最终拿出了两套实现一套是面向数字IMC的比特切片Bit-Slice方案另一套是面向模拟/神经形态计算的脉冲时序编码Spiking Temporal Encoding方案。实测下来在28nm工艺下模拟方案相比数字方案在面积和能效上都有接近2倍的提升而且对模型精度的影响微乎其微。2. 核心挑战拆解为什么Softmax是IMC的“天敌”要理解我们的设计首先得搞清楚传统Softmax在硬件实现上到底卡在哪里。这不仅仅是IMC的问题也是所有AI加速器都需要面对的。2.1 Softmax的计算本质与硬件不匹配标准的Softmax函数定义如下对于序列中第i个位置的注意力分数 $a_i$ $$ \text{Softmax}(a_i) \frac{e^{a_i}}{\sum_{j1}^{N} e^{a_j}} $$这里隐藏了三个对硬件极不友好的操作指数运算Exp这是一个高度非线性的函数。在数字电路中高精度的指数运算通常通过查找表LUT或CORDIC等迭代算法实现既占面积又耗电。在模拟电路中虽然可以利用晶体管或忆阻器的自然指数特性如亚阈值区的电流-电压关系来近似但精度和动态范围的控制极其困难。全局求和Summation分母需要对所有N个元素的指数结果求和。这意味着在计算第i个结果之前必须等所有N个 $e^{a_j}$ 都算出来。这导致了严重的顺序依赖和长延迟无法充分利用硬件并行性。在IMC中计算通常在各自的内存单元或子阵列中局部完成这种需要收集全局信息的操作会引发大量的片内或片间通信开销。除法Division每个 $e^{a_i}$ 都需要除以同一个全局和。除法同样是复杂的运算尤其是在追求高能效的模拟或低精度数字电路中。在Transformer的解码生成阶段这个问题会因为KV Cache而进一步放大。模型需要为每个新生成的token计算与之前所有token的注意力序列长度N会随着生成过程不断增长使得Softmax的计算开销线性增加成为端到端延迟的主要贡献者。2.2 存内计算IMC的范式与约束IMC的核心优势在于“计算在数据所在之处进行”。典型的形式是利用非易失性存储器如RRAM, PCM或经过特殊设计的SRAM/DRAM单元在其交叉阵列Crossbar上直接完成矩阵-向量乘法VMM。这完美匹配了神经网络中占大头的乘累加MAC操作。然而IMC有其固有的约束计算模式擅长并行的、局部的乘累加。输入是电压或电流输出是累加后的电流或积分后的电压。数据表示模拟域通常用电压/电流的幅度或时间宽度来表示数据。数字域则用多位精度。非理想因素模拟计算受工艺偏差Process Variation、电压波动Voltage Noise、温度变化Temperature Drift——合称PVT——的影响巨大。器件的不一致性会导致计算误差。灵活性一旦电路固化所能执行的计算模式就相对固定。Softmax的全局归一化特性与IMC这种擅长局部、并行VMM的架构格格不入。直接把Softmax映射到IMC阵列上要么需要复杂的额外电路和全局布线来收集求和信号要么就得把中间结果搬出阵列再用数字单元处理这都违背了IMC减少数据搬运的初衷。注意这里存在一个常见的误解认为IMC能加速所有AI计算。实际上IMC是“挑食”的它最爱的是计算密集、数据复用率高、操作规则化的线性代数运算如VMM。像Softmax这种带有非线性、全局依赖的操作需要精巧的重新设计才能高效适配。3. 算法-硬件协同设计重构注意力计算既然硬件有硬约束那我们就从算法层面动刀让算法去适应硬件而不是反过来。我们的目标是将计算复杂度从序列长度N转移到嵌入维度D。因为在大模型中N例如2048通常远大于单个注意力头的维度D例如64或128。3.1 比特切片注意力Bit-Slice Attention—— 数字方案这个思路的核心是分解与近似。我们不再直接计算原始的Softmax而是利用一些数学变换和硬件友好的近似。3.1.1 算法原理从Softmax到逐元素运算首先我们回顾一下注意力分数的计算$a_i q \cdot k_i^T / \sqrt{d_k}$其中q是查询向量$k_i$是第i个键向量。Softmax之后是加权求和$output \sum_i \text{Softmax}(a_i) v_i$。一个关键的观察是如果我们能避免显式计算所有 $a_i$ 的Softmax而是直接计算其对最终输出的贡献就有可能绕过全局求和。一种途径是使用线性注意力或基于核的近似但这些方法往往在精度上损失较大。我们的数字方案采用了一种更精确的近似路径。我们受到对数域计算Log-Sum-Exp的启发但对其进行了硬件优化。核心思想是将复杂的指数和除法运算分解为一系列基于比特位的、可并行操作的步骤。具体来说我们将注意力分数 $a_i$ 的二进制表示进行切片Bit-Slicing。每一“片”对应权重的一个二进制位。通过精心设计的移位、加法和比较操作我们可以用数字逻辑电路高效地模拟出Softmax的“竞争”与“归一化”效果。这个过程避免了高精度的指数运算单元也避免了等待所有数据就绪才能进行的全局求和。每个比特片的处理可以并行进行最后再将结果组合起来。3.1.2 硬件映射与数字IMC无缝集成这种比特切片算法非常适合数字IMC的架构。数字IMC例如基于SRAM的存内计算通常在阵列边缘配备数字处理单元如加法树、激活函数单元。我们的比特切片逻辑可以作为这样一个定制化的数字处理单元紧挨着IMC阵列放置。工作流程可以简化IMC阵列计算出原始的查询-键点积结果即 $a_i$。这些结果被送入比特切片Softmax单元。该单元以流水线或高度并行的方式处理比特片输出近似归一化的权重。这些权重立刻被用于下一次IMC阵列的乘法运算与值向量V相乘生成最终的注意力输出。由于整个计算流程都在数字域且分解成了规整的逻辑操作它对PVT变化不敏感精度可控并且能够完全兼容现有的LLM模型无需重新训练。在我们的实现中该方案支持14位定点数输入在28nm工艺下每次Softmax每行每个输入token能耗仅为0.69皮焦pJ。3.2 脉冲时序注意力Spiking Temporal Attention—— 模拟/神经形态方案如果说数字方案是“精确制导”那么模拟方案就是“自然模拟”。我们尝试利用物理定律本身来优雅地实现Softmax的功能。3.2.1 算法原理用时间编码竞争想象一下赛跑。所有选手对应注意力分数 $a_i$同时起跑但起跑线根据他们的实力分数高低有前有后。实力越强的选手起跑位置越靠前。所有选手以相同的速度跑步对应一个衰减过程。那么实力最强的选手将最先冲过终点线并且他冲线时领先的距离可以间接反映出他相对于其他选手的优势程度。我们的脉冲时序编码正是基于这个类比。我们将每个注意力分数 $a_i$ 映射为一个脉冲信号的起始时间。分数越高脉冲发出越早。这个脉冲会触发一个电容的放电过程模拟指数衰减。我们同时监测所有支路电容上的电压。关键技巧在于我们并不直接测量每个电容的绝对电压值那仍然需要全局信息而是利用一个公共的放电电流源和比较器。当任意一个电容的电压下降到某个阈值时它会触发一个事件。第一个触发事件的通道就对应着最大的 $a_i$因为它的脉冲最早电容放电时间最长。更重要的是通过测量从第一个脉冲发出到第一个触发事件之间的时间间隔以及后续事件的时间关系我们可以通过模拟电路的微分方程自然地推导出归一化的权重比例。数学上这对应着将Softmax计算转化为对一组一阶线性微分方程解的时序检测。指数衰减函数 $e^{-t/\tau}$ 被自然地用RC电路的放电曲线实现。求和与除法操作则通过电路中的电流镜、电荷共享或时序竞争机制隐式地完成。3.2.2 硬件实现模拟电路的优雅与挑战这套方案可以用完全模拟的电路实现无需ADC/DAC。电路核心包括脉冲生成器将输入电压代表 $a_i$转换为延迟时间可调的脉冲。衰减单元通常是一个由脉冲控制的开关和RC放电电路。仲裁与读出电路检测哪个衰减单元最先达到阈值并可能输出与权重成比例的信号如脉冲宽度或另一个时间量。这种方案的能效极高因为大部分计算是通过电荷的流动和时间的流逝 passively被动地完成的。在我们的28nm测试芯片中模拟方案每行每token的能耗仅为0.33 pJ面积也只有37.1 µm²相比数字方案优势明显。然而模拟方案面临严峻的PVT挑战。工艺偏差会导致不同芯片、甚至同一芯片上不同单元的RC常数不一致从而影响脉冲时序和衰减速度的精确性。我们的仿真和测试表明虽然电源电压和温度的波动会对脉冲的绝对起始时间和长度产生平移效应但由于Softmax是一个比值运算这种共同的平移因子在分子和分母中会被抵消因此对最终结果影响很小。真正的威胁来自于随机偏差这可能导致本应第一个触发的单元变成了第二个严重时会影响输出。实操心得在模拟电路设计中对抗PVT变异是永恒的主题。对于这种时序编码电路采用共模反馈结构、差分对设计以及动态偏置技术可以有效抑制电源和温度的共同影响。对于随机失配则需要通过增大器件面积牺牲面积换取匹配性或引入校准电路来补偿。4. 混合架构扬长避短的实践智慧数字方案精度高、鲁棒性强但能效和面积稍逊。模拟方案能效和面积极致但对异常值Outlier敏感动态范围有限输入被限制在0.5V范围内。在实际的LLM中注意力头之间的激活值分布差异很大总有少数“活跃”的头会产生非常大的分数容易超出模拟电路的输入范围导致饱和失真。我们提出的解决方案是混合架构。这不是简单的妥协而是基于对工作负载的深刻理解做出的最优决策。4.1 混合策略按需分配精准施策核心思想是让大多数“安静”的、行为规范的注意力头使用高效的模拟Softmax电路而少数那些容易产生异常大激活值的“调皮”注意力头则分配给鲁棒的数字Softmax电路。4.1.1 如何识别“调皮”的头这通常在模型部署前的分析阶段完成。我们可以用一批校准数据运行一次模型监控所有注意力头在整个序列长度上的最大激活值分布。那些历史最大激活值 consistently 接近或超过模拟电路输入范围上限的头就会被标记为“异常值易发头”Outlier-Prone Heads。在我们的实验中即使在大模型里这样的头也只占总数很小一部分例如5%。4.1.2 硬件如何支持混合在硬件架构上这意味着一块IMC加速芯片上需要集成两种Softmax计算单元。芯片的控制器或编译器需要根据预设的“头映射表”将不同注意力头的计算任务路由到对应的模拟或数字处理单元上。这增加了一些控制逻辑的复杂度但带来的收益是巨大的。从系统层面看这种混合方案实现了“鱼与熊掌兼得”保持了高能效95%的常规计算由模拟电路完成享受其超低功耗的优势。保证了高精度5%的关键计算由数字电路处理确保了整个模型输出的准确性避免了因少数异常头导致的整体性能崩塌。无需模型重训练这种硬件层面的优化对算法透明现有的、训练好的LLM如LLaMA、GPT系列可以直接部署精度无损。我们的实验数据如表VII所示证实这种混合策略几乎可以完全恢复纯数字实现的网络精度仅下降0.5%以内同时继承了模拟方案在能效和面积上的主要优势。4.2 与神经形态计算的融合脉冲时序编码方案还有一个意想不到的“副产品”它天然契合神经形态计算Neuromorphic Computing的范式。神经形态芯片如Intel的Loihi本身就是基于脉冲Spike和时序来处理信息的。我们将这套脉冲时序注意力算法移植到了Loihi 2芯片上进行验证。尽管Loihi并非为Transformer设计但我们的算法仅使用了其基本的漏电积分-发放Leaky Integrate-and-Fire神经元和突触延迟等原语。实验表明处理一个token的注意力计算头维度D64/128KV缓存长度N2048可在1毫秒内完成整芯片功耗控制在250毫瓦以内。这个实验极具启发性。它说明通过算法重构我们可以将Transformer这样的现代AI模型“翻译”成神经形态硬件能理解的语言从而在这些超低功耗的专用平台上运行。这为未来边缘设备部署大模型打开了一扇新的大门。5. 实现细节与性能剖析纸上谈兵终觉浅是骡子是马还得拉出来溜溜。下面我结合我们的芯片实现数据具体拆解一下两种方案的性能表现和设计权衡。5.1 数字比特切片方案实现细节我们的数字实现采用28nm CMOS工艺。整个比特切片Softmax单元作为一个数字宏Macro设计目标是与数字IMC SRAM阵列紧耦合。5.1.1 电路架构该单元主要由以下几部分构成输入预处理接收来自IMC阵列的14位定点数注意力分数。可能包括一个缩放电路将输入调整到适合后续处理的动态范围。比特切片引擎这是核心。它包含多个并行的处理通道每个通道负责处理输入的一个比特片。每个通道内部有移位寄存器、加法器和比较器。累加与归一化逻辑将各个比特片处理的结果按权重累加并进行最后的归一化调整。这里我们采用了一种近似的归一化方法用移位和加法来模拟除法避免了复杂的除法器。输出缓冲将最终计算出的归一化权重通常是较低位宽如8位锁存并准备好供下一级IMC阵列用于与V矩阵相乘使用。5.1.2 性能指标与权衡面积81.2 µm²。这个面积主要消耗在并行的比特处理逻辑和中间结果缓存上。虽然比模拟方案大但相比传统需要高精度指数和除法单元的Softmax硬件面积已经大大缩减。能耗0.69 pJ/行/token。能耗主要来自数字开关活动和时钟树分布。其优势在于能耗可预测且几乎不受PVT影响。精度支持14位输入输出精度足以满足主流LLM如LLaMA 2/3的需求在标准评测集上零精度损失。延迟延迟与输入位宽和切片数有关但与序列长度N无关。对于7位有效输入其延迟与最新的数字Softmax专用硬件加速器相当并且省去了与IMC阵列间昂贵的ADC/DAC转换开销。5.2 模拟脉冲时序方案实现细节模拟方案同样基于28nm工艺其版图设计需要精心匹配IMC存储单元的尺寸Pitch以实现无缝集成。5.2.1 关键电路模块电压-时间转换器V2T这是将输入电压 $a_i$ 转换为脉冲延迟的核心。我们采用了一种基于电流积分的设计。输入电压控制一个电流源的电流大小该电流对一个电容进行充电当电容电压达到某个阈值时触发脉冲。这样输入电压越高充电电流越大达到阈值的时间越短脉冲发出越早。可配置衰减单元每个通道都有一个RC放电电路。当脉冲到来时开关打开电容开始通过一个固定电阻放电。RC时间常数τ需要精确设计它决定了衰减速度和整个计算的时间窗口。全局仲裁与读出电路所有衰减单元的输出连接到一个公共的“优胜者通吃”Winner-Take-All, WTA电路或一个带使能的比较器阵列。第一个电压衰减到阈值的单元会触发一个标志并可能通过一个脉冲宽度调制PWM电路将其放电时间转化为一个与Softmax权重成比例的脉冲宽度信号。5.2.2 非理想因素分析与应对模拟电路的性能高度依赖于对PVT变化的鲁棒性。我们进行了大量的蒙特卡洛仿真和角Corner分析工艺偏差晶体管阈值电压Vth和电阻、电容的绝对值变化会影响V2T的转换增益和RC时间常数。我们通过采用大尺寸器件降低失配和差分电路结构来抑制。电压波动电源电压VDD的波动会直接影响电流源的输出电流和比较器的阈值。仿真表明这主要导致脉冲起始时间和宽度的整体平移。如前所述由于Softmax是比值运算这种平移在最终计算中会被抵消。电路设计上我们使用带隙基准Bandgap Reference来生成稳定的偏置电流和电压。温度变化温度影响载流子迁移率和阈值电压。其效应与电压波动类似也主要产生系统性偏移对最终结果影响有限。我们可以在电路中集成简单的温度传感器对偏置进行微调。5.2.3 性能指标与权衡面积37.1 µm²仅为数字方案的45%。面积优势主要来自模拟电路的简洁性无需大量的数字逻辑门和寄存器。能耗0.33 pJ/行/token约为数字方案的一半。能量主要消耗在脉冲生成和短暂的放电过程中静态功耗极低。精度与动态范围这是模拟方案的软肋。输入被限制在0.5V的电压范围内。对于超出此范围的异常大输入电路会饱和导致计算错误。这正是我们引入混合架构的主要原因。延迟延迟由最慢的放电路径决定与序列长度N无关。在我们的设计中整个模拟Softmax包括后续驱动IMC阵列的电路的延迟约为120ns这与典型RRAM或PCM阵列完成一次VMM操作的时间相匹配可以实现流水线化操作。5.3 对比与选型指南我们将两种方案的关键指标总结如下表方便你根据应用场景做选择特性数字比特切片方案模拟脉冲时序方案说明与建议实现原理数字逻辑近似比特并行处理模拟物理定律时序竞争编码数字方案确定性高模拟方案物理直觉强。能效0.69 pJ/行/token0.33 pJ/行/token模拟方案能效优势明显适合功耗敏感场景。面积81.2 µm²37.1 µm²模拟方案面积更小集成密度更高。精度14位定点无损受限于0.5V输入范围对异常值敏感数字方案精度有保障模拟方案需处理动态范围问题。PVT鲁棒性高数字逻辑对偏差不敏感中/低需通过设计技术缓解数字方案更适合环境变化大的场合。延迟与位宽相关与N无关~120ns (匹配IMC VMM时间)与N无关两者延迟均与序列长度解耦是关键优势。与IMC集成需数字接口兼容性好直接模拟电流/电压接口更“原生”模拟方案与模拟IMC阵列连接更自然减少数据转换。适用场景高精度要求通用性强或作为混合方案中的“安全岛”超低功耗、小面积场景或与模拟IMC紧密耦合推荐混合架构模拟为主数字为辅处理异常头。6. 系统集成与未来展望将优化后的Softmax单元集成到完整的Transformer加速器中还需要考虑系统级的问题。这不仅仅是把一块电路放进去那么简单。6.1 与IMC阵列的协同工作流一个完整的注意力层加速需要协调好多个IMC阵列和Softmax单元的工作。以混合架构为例一个可能的工作流是QK^T计算查询向量Q和键矩阵K的转置通过一个或多个IMC阵列并行完成输出注意力分数向量A长度为N。头分类与路由根据预设的“头映射表”系统控制器将向量A按注意力头分组并决定每组是发送到模拟Softmax池还是数字Softmax单元。这个路由逻辑需要低延迟。并行Softmax计算模拟和数字单元同时开始计算各自分配的注意力权重。加权求和计算出的权重向量W模拟和数字结果需在数字域合并再次被送入另一组IMC阵列与值矩阵V相乘完成最终的注意力输出 $O W \cdot V$。流水线化为了隐藏延迟QK^T计算、Softmax、WV计算可以组成流线。当第t个token在进行Softmax时第t1个token的QK^T计算可以同时进行。6.2 应对长序列的挑战KV Cache与Paged Attention我们的方案将Softmax的计算复杂度从O(N)转移到了O(D)这本身就是为了应对长序列。但在实际系统中当N非常大比如数万甚至数十万时即使Softmax本身高效存储和访问巨大的KV Cache也会成为瓶颈。我们的硬件设计可以与软件层的优化技术结合Paged Attention类似操作系统管理内存将KV Cache分成一个个“页”。我们的IMC加速器可以设计相应的内存控制器高效地按需加载这些页到近存的IMC阵列中进行计算。KV Cache压缩在将键值对存入IMC阵列前先进行轻量级压缩如量化、稀疏化。我们的Softmax单元需要能处理压缩后的数据格式这可能需要对输入预处理阶段进行微调。6.3 超越Softmax更广阔的算法-硬件协同设计空间这项工作给我们最大的启示是面对硬件瓶颈最有效的往往不是更精巧的电路而是更聪明的算法。Softmax只是Transformer中的一个点。Transformer中还有其他潜在瓶颈如Layer Normalization、GELU激活函数等它们是否也能通过类似的“算法重构硬件定制”来优化更进一步当前火热的MambaSSM模型等架构试图用递归结构替代注意力其核心也是为了规避全局依赖。我们的工作表明即使不抛弃注意力通过深度的软硬件协同设计也能让它在现代硬件上高效运行。未来的AI加速芯片很可能不再是通用的矩阵乘法引擎而是为特定算法范式如注意力、状态空间模型深度定制的异构计算平台其中包含多种像我们设计的这种“算法-硬件”共生的专用加速单元。个人体会做了这么多年硬件我越来越觉得顶尖的硬件工程师必须懂算法而顶尖的算法工程师也需要理解硬件的基本约束。软硬件协同设计不是一句空话它要求双方在同一个抽象层次上对话共同定义计算范式。就像这次我们把Softmax从“全局除法”重新理解为“时序竞争”整个硬件设计的天地就豁然开朗了。这条路很难需要跨学科的知识储备但这也是它最有魅力的地方。如果你正在从事AI加速相关的工作我强烈建议你花时间深入理解你要加速的算法它的数学本质是什么有没有等价的、对硬件更友好的形式这往往是突破性能天花板的关键。
Transformer存内计算加速：软硬件协同设计攻克Softmax硬件瓶颈
1. 项目概述当Transformer遇上存内计算Softmax为何成了“拦路虎”如果你最近关注过大语言模型LLM的硬件加速尤其是存内计算In-Memory Computing, IMC这个热门方向那你肯定听过一个说法Transformer模型里的注意力机制Attention是性能瓶颈。但更具体一点这个瓶颈的“罪魁祸首”往往指向一个看似简单的数学函数——Softmax。今天我想从一个硬件设计者的角度跟你聊聊为什么Softmax在IMC架构上这么“难搞”以及我们是如何通过一套软硬件协同设计的“组合拳”把这个瓶颈给打通的。简单来说Softmax是注意力机制里的“裁判”。它负责把一堆原始的注意力分数比如查询向量和键向量的点积结果归一化成概率分布决定模型应该“看”输入序列的哪些部分。问题就出在这个“归一化”上。传统的数字电路实现Softmax需要先对所有元素求指数exp再求和最后做除法。这个过程涉及大量的数据搬运把数据从内存搬到计算单元和复杂的非线性运算在序列长度N动辄成千上万的LLM场景下能耗和延迟都非常吓人。而IMC的初衷是把计算放到内存阵列里做减少数据搬运但Softmax这种需要全局信息的操作天然和IMC“计算局部化”的特性相悖。我们提出的方案核心思想是“算法重构硬件适配”。不是硬着头皮在IMC上跑传统的Softmax而是从算法层面改造注意力计算把计算负担从庞大的序列维度N转移到相对较小的嵌入维度D。这样一来核心的Softmax-like操作就能被巧妙地映射到IMC阵列或者专用的模拟/数字电路上。我们最终拿出了两套实现一套是面向数字IMC的比特切片Bit-Slice方案另一套是面向模拟/神经形态计算的脉冲时序编码Spiking Temporal Encoding方案。实测下来在28nm工艺下模拟方案相比数字方案在面积和能效上都有接近2倍的提升而且对模型精度的影响微乎其微。2. 核心挑战拆解为什么Softmax是IMC的“天敌”要理解我们的设计首先得搞清楚传统Softmax在硬件实现上到底卡在哪里。这不仅仅是IMC的问题也是所有AI加速器都需要面对的。2.1 Softmax的计算本质与硬件不匹配标准的Softmax函数定义如下对于序列中第i个位置的注意力分数 $a_i$ $$ \text{Softmax}(a_i) \frac{e^{a_i}}{\sum_{j1}^{N} e^{a_j}} $$这里隐藏了三个对硬件极不友好的操作指数运算Exp这是一个高度非线性的函数。在数字电路中高精度的指数运算通常通过查找表LUT或CORDIC等迭代算法实现既占面积又耗电。在模拟电路中虽然可以利用晶体管或忆阻器的自然指数特性如亚阈值区的电流-电压关系来近似但精度和动态范围的控制极其困难。全局求和Summation分母需要对所有N个元素的指数结果求和。这意味着在计算第i个结果之前必须等所有N个 $e^{a_j}$ 都算出来。这导致了严重的顺序依赖和长延迟无法充分利用硬件并行性。在IMC中计算通常在各自的内存单元或子阵列中局部完成这种需要收集全局信息的操作会引发大量的片内或片间通信开销。除法Division每个 $e^{a_i}$ 都需要除以同一个全局和。除法同样是复杂的运算尤其是在追求高能效的模拟或低精度数字电路中。在Transformer的解码生成阶段这个问题会因为KV Cache而进一步放大。模型需要为每个新生成的token计算与之前所有token的注意力序列长度N会随着生成过程不断增长使得Softmax的计算开销线性增加成为端到端延迟的主要贡献者。2.2 存内计算IMC的范式与约束IMC的核心优势在于“计算在数据所在之处进行”。典型的形式是利用非易失性存储器如RRAM, PCM或经过特殊设计的SRAM/DRAM单元在其交叉阵列Crossbar上直接完成矩阵-向量乘法VMM。这完美匹配了神经网络中占大头的乘累加MAC操作。然而IMC有其固有的约束计算模式擅长并行的、局部的乘累加。输入是电压或电流输出是累加后的电流或积分后的电压。数据表示模拟域通常用电压/电流的幅度或时间宽度来表示数据。数字域则用多位精度。非理想因素模拟计算受工艺偏差Process Variation、电压波动Voltage Noise、温度变化Temperature Drift——合称PVT——的影响巨大。器件的不一致性会导致计算误差。灵活性一旦电路固化所能执行的计算模式就相对固定。Softmax的全局归一化特性与IMC这种擅长局部、并行VMM的架构格格不入。直接把Softmax映射到IMC阵列上要么需要复杂的额外电路和全局布线来收集求和信号要么就得把中间结果搬出阵列再用数字单元处理这都违背了IMC减少数据搬运的初衷。注意这里存在一个常见的误解认为IMC能加速所有AI计算。实际上IMC是“挑食”的它最爱的是计算密集、数据复用率高、操作规则化的线性代数运算如VMM。像Softmax这种带有非线性、全局依赖的操作需要精巧的重新设计才能高效适配。3. 算法-硬件协同设计重构注意力计算既然硬件有硬约束那我们就从算法层面动刀让算法去适应硬件而不是反过来。我们的目标是将计算复杂度从序列长度N转移到嵌入维度D。因为在大模型中N例如2048通常远大于单个注意力头的维度D例如64或128。3.1 比特切片注意力Bit-Slice Attention—— 数字方案这个思路的核心是分解与近似。我们不再直接计算原始的Softmax而是利用一些数学变换和硬件友好的近似。3.1.1 算法原理从Softmax到逐元素运算首先我们回顾一下注意力分数的计算$a_i q \cdot k_i^T / \sqrt{d_k}$其中q是查询向量$k_i$是第i个键向量。Softmax之后是加权求和$output \sum_i \text{Softmax}(a_i) v_i$。一个关键的观察是如果我们能避免显式计算所有 $a_i$ 的Softmax而是直接计算其对最终输出的贡献就有可能绕过全局求和。一种途径是使用线性注意力或基于核的近似但这些方法往往在精度上损失较大。我们的数字方案采用了一种更精确的近似路径。我们受到对数域计算Log-Sum-Exp的启发但对其进行了硬件优化。核心思想是将复杂的指数和除法运算分解为一系列基于比特位的、可并行操作的步骤。具体来说我们将注意力分数 $a_i$ 的二进制表示进行切片Bit-Slicing。每一“片”对应权重的一个二进制位。通过精心设计的移位、加法和比较操作我们可以用数字逻辑电路高效地模拟出Softmax的“竞争”与“归一化”效果。这个过程避免了高精度的指数运算单元也避免了等待所有数据就绪才能进行的全局求和。每个比特片的处理可以并行进行最后再将结果组合起来。3.1.2 硬件映射与数字IMC无缝集成这种比特切片算法非常适合数字IMC的架构。数字IMC例如基于SRAM的存内计算通常在阵列边缘配备数字处理单元如加法树、激活函数单元。我们的比特切片逻辑可以作为这样一个定制化的数字处理单元紧挨着IMC阵列放置。工作流程可以简化IMC阵列计算出原始的查询-键点积结果即 $a_i$。这些结果被送入比特切片Softmax单元。该单元以流水线或高度并行的方式处理比特片输出近似归一化的权重。这些权重立刻被用于下一次IMC阵列的乘法运算与值向量V相乘生成最终的注意力输出。由于整个计算流程都在数字域且分解成了规整的逻辑操作它对PVT变化不敏感精度可控并且能够完全兼容现有的LLM模型无需重新训练。在我们的实现中该方案支持14位定点数输入在28nm工艺下每次Softmax每行每个输入token能耗仅为0.69皮焦pJ。3.2 脉冲时序注意力Spiking Temporal Attention—— 模拟/神经形态方案如果说数字方案是“精确制导”那么模拟方案就是“自然模拟”。我们尝试利用物理定律本身来优雅地实现Softmax的功能。3.2.1 算法原理用时间编码竞争想象一下赛跑。所有选手对应注意力分数 $a_i$同时起跑但起跑线根据他们的实力分数高低有前有后。实力越强的选手起跑位置越靠前。所有选手以相同的速度跑步对应一个衰减过程。那么实力最强的选手将最先冲过终点线并且他冲线时领先的距离可以间接反映出他相对于其他选手的优势程度。我们的脉冲时序编码正是基于这个类比。我们将每个注意力分数 $a_i$ 映射为一个脉冲信号的起始时间。分数越高脉冲发出越早。这个脉冲会触发一个电容的放电过程模拟指数衰减。我们同时监测所有支路电容上的电压。关键技巧在于我们并不直接测量每个电容的绝对电压值那仍然需要全局信息而是利用一个公共的放电电流源和比较器。当任意一个电容的电压下降到某个阈值时它会触发一个事件。第一个触发事件的通道就对应着最大的 $a_i$因为它的脉冲最早电容放电时间最长。更重要的是通过测量从第一个脉冲发出到第一个触发事件之间的时间间隔以及后续事件的时间关系我们可以通过模拟电路的微分方程自然地推导出归一化的权重比例。数学上这对应着将Softmax计算转化为对一组一阶线性微分方程解的时序检测。指数衰减函数 $e^{-t/\tau}$ 被自然地用RC电路的放电曲线实现。求和与除法操作则通过电路中的电流镜、电荷共享或时序竞争机制隐式地完成。3.2.2 硬件实现模拟电路的优雅与挑战这套方案可以用完全模拟的电路实现无需ADC/DAC。电路核心包括脉冲生成器将输入电压代表 $a_i$转换为延迟时间可调的脉冲。衰减单元通常是一个由脉冲控制的开关和RC放电电路。仲裁与读出电路检测哪个衰减单元最先达到阈值并可能输出与权重成比例的信号如脉冲宽度或另一个时间量。这种方案的能效极高因为大部分计算是通过电荷的流动和时间的流逝 passively被动地完成的。在我们的28nm测试芯片中模拟方案每行每token的能耗仅为0.33 pJ面积也只有37.1 µm²相比数字方案优势明显。然而模拟方案面临严峻的PVT挑战。工艺偏差会导致不同芯片、甚至同一芯片上不同单元的RC常数不一致从而影响脉冲时序和衰减速度的精确性。我们的仿真和测试表明虽然电源电压和温度的波动会对脉冲的绝对起始时间和长度产生平移效应但由于Softmax是一个比值运算这种共同的平移因子在分子和分母中会被抵消因此对最终结果影响很小。真正的威胁来自于随机偏差这可能导致本应第一个触发的单元变成了第二个严重时会影响输出。实操心得在模拟电路设计中对抗PVT变异是永恒的主题。对于这种时序编码电路采用共模反馈结构、差分对设计以及动态偏置技术可以有效抑制电源和温度的共同影响。对于随机失配则需要通过增大器件面积牺牲面积换取匹配性或引入校准电路来补偿。4. 混合架构扬长避短的实践智慧数字方案精度高、鲁棒性强但能效和面积稍逊。模拟方案能效和面积极致但对异常值Outlier敏感动态范围有限输入被限制在0.5V范围内。在实际的LLM中注意力头之间的激活值分布差异很大总有少数“活跃”的头会产生非常大的分数容易超出模拟电路的输入范围导致饱和失真。我们提出的解决方案是混合架构。这不是简单的妥协而是基于对工作负载的深刻理解做出的最优决策。4.1 混合策略按需分配精准施策核心思想是让大多数“安静”的、行为规范的注意力头使用高效的模拟Softmax电路而少数那些容易产生异常大激活值的“调皮”注意力头则分配给鲁棒的数字Softmax电路。4.1.1 如何识别“调皮”的头这通常在模型部署前的分析阶段完成。我们可以用一批校准数据运行一次模型监控所有注意力头在整个序列长度上的最大激活值分布。那些历史最大激活值 consistently 接近或超过模拟电路输入范围上限的头就会被标记为“异常值易发头”Outlier-Prone Heads。在我们的实验中即使在大模型里这样的头也只占总数很小一部分例如5%。4.1.2 硬件如何支持混合在硬件架构上这意味着一块IMC加速芯片上需要集成两种Softmax计算单元。芯片的控制器或编译器需要根据预设的“头映射表”将不同注意力头的计算任务路由到对应的模拟或数字处理单元上。这增加了一些控制逻辑的复杂度但带来的收益是巨大的。从系统层面看这种混合方案实现了“鱼与熊掌兼得”保持了高能效95%的常规计算由模拟电路完成享受其超低功耗的优势。保证了高精度5%的关键计算由数字电路处理确保了整个模型输出的准确性避免了因少数异常头导致的整体性能崩塌。无需模型重训练这种硬件层面的优化对算法透明现有的、训练好的LLM如LLaMA、GPT系列可以直接部署精度无损。我们的实验数据如表VII所示证实这种混合策略几乎可以完全恢复纯数字实现的网络精度仅下降0.5%以内同时继承了模拟方案在能效和面积上的主要优势。4.2 与神经形态计算的融合脉冲时序编码方案还有一个意想不到的“副产品”它天然契合神经形态计算Neuromorphic Computing的范式。神经形态芯片如Intel的Loihi本身就是基于脉冲Spike和时序来处理信息的。我们将这套脉冲时序注意力算法移植到了Loihi 2芯片上进行验证。尽管Loihi并非为Transformer设计但我们的算法仅使用了其基本的漏电积分-发放Leaky Integrate-and-Fire神经元和突触延迟等原语。实验表明处理一个token的注意力计算头维度D64/128KV缓存长度N2048可在1毫秒内完成整芯片功耗控制在250毫瓦以内。这个实验极具启发性。它说明通过算法重构我们可以将Transformer这样的现代AI模型“翻译”成神经形态硬件能理解的语言从而在这些超低功耗的专用平台上运行。这为未来边缘设备部署大模型打开了一扇新的大门。5. 实现细节与性能剖析纸上谈兵终觉浅是骡子是马还得拉出来溜溜。下面我结合我们的芯片实现数据具体拆解一下两种方案的性能表现和设计权衡。5.1 数字比特切片方案实现细节我们的数字实现采用28nm CMOS工艺。整个比特切片Softmax单元作为一个数字宏Macro设计目标是与数字IMC SRAM阵列紧耦合。5.1.1 电路架构该单元主要由以下几部分构成输入预处理接收来自IMC阵列的14位定点数注意力分数。可能包括一个缩放电路将输入调整到适合后续处理的动态范围。比特切片引擎这是核心。它包含多个并行的处理通道每个通道负责处理输入的一个比特片。每个通道内部有移位寄存器、加法器和比较器。累加与归一化逻辑将各个比特片处理的结果按权重累加并进行最后的归一化调整。这里我们采用了一种近似的归一化方法用移位和加法来模拟除法避免了复杂的除法器。输出缓冲将最终计算出的归一化权重通常是较低位宽如8位锁存并准备好供下一级IMC阵列用于与V矩阵相乘使用。5.1.2 性能指标与权衡面积81.2 µm²。这个面积主要消耗在并行的比特处理逻辑和中间结果缓存上。虽然比模拟方案大但相比传统需要高精度指数和除法单元的Softmax硬件面积已经大大缩减。能耗0.69 pJ/行/token。能耗主要来自数字开关活动和时钟树分布。其优势在于能耗可预测且几乎不受PVT影响。精度支持14位输入输出精度足以满足主流LLM如LLaMA 2/3的需求在标准评测集上零精度损失。延迟延迟与输入位宽和切片数有关但与序列长度N无关。对于7位有效输入其延迟与最新的数字Softmax专用硬件加速器相当并且省去了与IMC阵列间昂贵的ADC/DAC转换开销。5.2 模拟脉冲时序方案实现细节模拟方案同样基于28nm工艺其版图设计需要精心匹配IMC存储单元的尺寸Pitch以实现无缝集成。5.2.1 关键电路模块电压-时间转换器V2T这是将输入电压 $a_i$ 转换为脉冲延迟的核心。我们采用了一种基于电流积分的设计。输入电压控制一个电流源的电流大小该电流对一个电容进行充电当电容电压达到某个阈值时触发脉冲。这样输入电压越高充电电流越大达到阈值的时间越短脉冲发出越早。可配置衰减单元每个通道都有一个RC放电电路。当脉冲到来时开关打开电容开始通过一个固定电阻放电。RC时间常数τ需要精确设计它决定了衰减速度和整个计算的时间窗口。全局仲裁与读出电路所有衰减单元的输出连接到一个公共的“优胜者通吃”Winner-Take-All, WTA电路或一个带使能的比较器阵列。第一个电压衰减到阈值的单元会触发一个标志并可能通过一个脉冲宽度调制PWM电路将其放电时间转化为一个与Softmax权重成比例的脉冲宽度信号。5.2.2 非理想因素分析与应对模拟电路的性能高度依赖于对PVT变化的鲁棒性。我们进行了大量的蒙特卡洛仿真和角Corner分析工艺偏差晶体管阈值电压Vth和电阻、电容的绝对值变化会影响V2T的转换增益和RC时间常数。我们通过采用大尺寸器件降低失配和差分电路结构来抑制。电压波动电源电压VDD的波动会直接影响电流源的输出电流和比较器的阈值。仿真表明这主要导致脉冲起始时间和宽度的整体平移。如前所述由于Softmax是比值运算这种平移在最终计算中会被抵消。电路设计上我们使用带隙基准Bandgap Reference来生成稳定的偏置电流和电压。温度变化温度影响载流子迁移率和阈值电压。其效应与电压波动类似也主要产生系统性偏移对最终结果影响有限。我们可以在电路中集成简单的温度传感器对偏置进行微调。5.2.3 性能指标与权衡面积37.1 µm²仅为数字方案的45%。面积优势主要来自模拟电路的简洁性无需大量的数字逻辑门和寄存器。能耗0.33 pJ/行/token约为数字方案的一半。能量主要消耗在脉冲生成和短暂的放电过程中静态功耗极低。精度与动态范围这是模拟方案的软肋。输入被限制在0.5V的电压范围内。对于超出此范围的异常大输入电路会饱和导致计算错误。这正是我们引入混合架构的主要原因。延迟延迟由最慢的放电路径决定与序列长度N无关。在我们的设计中整个模拟Softmax包括后续驱动IMC阵列的电路的延迟约为120ns这与典型RRAM或PCM阵列完成一次VMM操作的时间相匹配可以实现流水线化操作。5.3 对比与选型指南我们将两种方案的关键指标总结如下表方便你根据应用场景做选择特性数字比特切片方案模拟脉冲时序方案说明与建议实现原理数字逻辑近似比特并行处理模拟物理定律时序竞争编码数字方案确定性高模拟方案物理直觉强。能效0.69 pJ/行/token0.33 pJ/行/token模拟方案能效优势明显适合功耗敏感场景。面积81.2 µm²37.1 µm²模拟方案面积更小集成密度更高。精度14位定点无损受限于0.5V输入范围对异常值敏感数字方案精度有保障模拟方案需处理动态范围问题。PVT鲁棒性高数字逻辑对偏差不敏感中/低需通过设计技术缓解数字方案更适合环境变化大的场合。延迟与位宽相关与N无关~120ns (匹配IMC VMM时间)与N无关两者延迟均与序列长度解耦是关键优势。与IMC集成需数字接口兼容性好直接模拟电流/电压接口更“原生”模拟方案与模拟IMC阵列连接更自然减少数据转换。适用场景高精度要求通用性强或作为混合方案中的“安全岛”超低功耗、小面积场景或与模拟IMC紧密耦合推荐混合架构模拟为主数字为辅处理异常头。6. 系统集成与未来展望将优化后的Softmax单元集成到完整的Transformer加速器中还需要考虑系统级的问题。这不仅仅是把一块电路放进去那么简单。6.1 与IMC阵列的协同工作流一个完整的注意力层加速需要协调好多个IMC阵列和Softmax单元的工作。以混合架构为例一个可能的工作流是QK^T计算查询向量Q和键矩阵K的转置通过一个或多个IMC阵列并行完成输出注意力分数向量A长度为N。头分类与路由根据预设的“头映射表”系统控制器将向量A按注意力头分组并决定每组是发送到模拟Softmax池还是数字Softmax单元。这个路由逻辑需要低延迟。并行Softmax计算模拟和数字单元同时开始计算各自分配的注意力权重。加权求和计算出的权重向量W模拟和数字结果需在数字域合并再次被送入另一组IMC阵列与值矩阵V相乘完成最终的注意力输出 $O W \cdot V$。流水线化为了隐藏延迟QK^T计算、Softmax、WV计算可以组成流线。当第t个token在进行Softmax时第t1个token的QK^T计算可以同时进行。6.2 应对长序列的挑战KV Cache与Paged Attention我们的方案将Softmax的计算复杂度从O(N)转移到了O(D)这本身就是为了应对长序列。但在实际系统中当N非常大比如数万甚至数十万时即使Softmax本身高效存储和访问巨大的KV Cache也会成为瓶颈。我们的硬件设计可以与软件层的优化技术结合Paged Attention类似操作系统管理内存将KV Cache分成一个个“页”。我们的IMC加速器可以设计相应的内存控制器高效地按需加载这些页到近存的IMC阵列中进行计算。KV Cache压缩在将键值对存入IMC阵列前先进行轻量级压缩如量化、稀疏化。我们的Softmax单元需要能处理压缩后的数据格式这可能需要对输入预处理阶段进行微调。6.3 超越Softmax更广阔的算法-硬件协同设计空间这项工作给我们最大的启示是面对硬件瓶颈最有效的往往不是更精巧的电路而是更聪明的算法。Softmax只是Transformer中的一个点。Transformer中还有其他潜在瓶颈如Layer Normalization、GELU激活函数等它们是否也能通过类似的“算法重构硬件定制”来优化更进一步当前火热的MambaSSM模型等架构试图用递归结构替代注意力其核心也是为了规避全局依赖。我们的工作表明即使不抛弃注意力通过深度的软硬件协同设计也能让它在现代硬件上高效运行。未来的AI加速芯片很可能不再是通用的矩阵乘法引擎而是为特定算法范式如注意力、状态空间模型深度定制的异构计算平台其中包含多种像我们设计的这种“算法-硬件”共生的专用加速单元。个人体会做了这么多年硬件我越来越觉得顶尖的硬件工程师必须懂算法而顶尖的算法工程师也需要理解硬件的基本约束。软硬件协同设计不是一句空话它要求双方在同一个抽象层次上对话共同定义计算范式。就像这次我们把Softmax从“全局除法”重新理解为“时序竞争”整个硬件设计的天地就豁然开朗了。这条路很难需要跨学科的知识储备但这也是它最有魅力的地方。如果你正在从事AI加速相关的工作我强烈建议你花时间深入理解你要加速的算法它的数学本质是什么有没有等价的、对硬件更友好的形式这往往是突破性能天花板的关键。
