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_{j=1}^{N} e^{a_j}} $$

这里隐藏了三个对硬件极不友好的操作：

1. 指数运算（Exp）：这是一个高度非线性的函数。在数字电路中，高精度的指数运算通常通过查找表（LUT）或CORDIC等迭代算法实现，既占面积又耗电。在模拟电路中，虽然可以利用晶体管或忆阻器的自然指数特性（如亚阈值区的电流-电压关系）来近似，但精度和动态范围的控制极其困难。
2. 全局求和（Summation）：分母需要对所有N个元素的指数结果求和。这意味着在计算第i个结果之前，必须等所有N个 $e^{a_j}$ 都算出来。这导致了严重的顺序依赖和长延迟，无法充分利用硬件并行性。在IMC中，计算通常在各自的内存单元或子阵列中局部完成，这种需要收集全局信息的操作会引发大量的片内或片间通信开销。
3. 除法（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阵列放置。

工作流程可以简化��：

1. IMC阵列计算出原始的查询-键点积结果（即 $a_i$）。
2. 这些结果被送入比特切片Softmax单元。
3. 该单元以流水线或高度并行的方式处理比特片，输出近似归一化的权重。
4. 这些权重立刻被用于下一次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的注意力计算（头维度D=64/128，KV缓存长度N=2048）可在1毫秒内完成，整芯片功耗控制在250毫瓦以内。

这个实验极具启发性。它说明，通过算法重构，我们可以将Transformer这样的现代AI模型“翻译”成神经形态硬件能理解的语言，从而在这些超低功耗的专用平台上运行。这为未来边缘设备部署大模型打开了一扇新的大门。

5. 实现细节与性能剖析

纸上谈兵终觉浅，是骡子是马还得拉出来溜溜。下面我结合我们的芯片实现数据，具体拆解一下两种方案的性能表现和设计权衡。

5.1 数字比特切片方案实现细节

我们的数字实现采用28nm CMOS工艺。整个比特切片Softmax单元作为一个数字宏（Macro），设计目标是与数字IMC SRAM阵列紧耦合。

5.1.1 电路架构该单元主要由以下几部分构成：

1. 输入预处理：接收来自IMC阵列的14位定点数注意力分数。可能包括一个缩放电路，将输入调整到适合后续处理的动态范围。
2. 比特切片引擎：这是核心。它包含多个并行的处理通道，每个通道负责处理输入的一个比特片。每个通道内部有移位寄存器、加法器和比较器。
3. 累加与归一化逻辑：将各个比特片处理的结果按权重累加，并进行最后的归一化调整。这里我们采用了一种近似的归一化方法，用移位和加法来模拟除法，避免了复杂的除法器。
4. 输出缓冲：将最终计算出的归一化权重（通常是较低位宽，如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 关键电路模块

1. 电压-时间转换器（V2T）：这是将输入电压 $a_i$ 转换为脉冲延迟的核心。我们采用了一种基于电流积分的设计。输入电压控制一个电流源的电流大小，该电流对一个电容进行充电，当电容电压达到某个阈值时，触发脉冲。这样，输入电压越高，充电电流越大，达到阈值的时间越短，脉冲发出越早。
2. 可配置衰减单元：每个通道都有一个RC放电电路。当脉冲到来时，开关打开，电容开始通过一个固定电阻放电。RC时间常数（τ）需要精确设计，它决定了衰减速度和整个计算的时间窗口。
3. 全局仲裁与读出电路：所有衰减单元的输出连接到一个公共的“优胜者通吃”（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/行/token	0.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单元的工作。以混合架构为例，一个可能的工作流是：

1. QK^T计算：查询向量Q和键矩阵K的转置，通过一个或多个IMC阵列并行完成，输出注意力分数向量A（长度为N）。
2. 头分类与路由：根据预设的“头映射表”，系统控制器将向量A按注意力头分组，并决定每组是发送到模拟Softmax池还是数字Softmax单元。这个路由逻辑需要低延迟。
3. 并行Softmax计算：模拟和数字单元同时开始计算各自分配的注意力权重。
4. 加权求和：计算出的权重向量W（模拟和数字结果需在数字域合并）再次被送入另一组IMC阵列，与值矩阵V相乘，完成最终的注意力输出 $O = W \cdot V$。
5. 流水线化：为了隐藏延迟，QK^T计算、Softmax、WV计算可以组成流��线。当第t个token在进行Softmax时，第t+1个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激活函数等，它们是否也能通过类似的“算法重构+硬件定制”来优化？

更进一步，当前火热的Mamba（SSM模型）等架构，试图用递归结构替代注意力，其核心也是为了规避全局依赖。我们的工作表明，即使不抛弃注意力，通过深度的软硬件协同设计，也能让它在现代硬件上高效运行。未来的AI加速芯片，很可能不再是通用的矩阵乘法引擎，而是为特定算法范式（如注意力、状态空间模型）深度定制的异构计算平台，其中包含多种像我们设计的这种“算法-硬件”共生的专用加速单元。

个人体会：做了这么多年硬件，我越来越觉得，顶尖的硬件工程师必须懂算法，而顶尖的算法工程师也需要理解硬件的基本约束。软硬件协同设计不是一句空话，它要求双方在同一个抽象层次上对话，共同定义计算范式。就像这次，我们把Softmax从“全局除法”重新理解为“时序竞争”，整个硬件设计的天地就豁然开朗了。这条路很难，需要跨学科的知识储备，但这也是它最有魅力的地方。如果你正在从事AI加速相关的工作，我强烈建议你花时间深入理解你要加速的算法，它的数学本质是什么？有没有等价的、对硬件更友好的形式？这往往是突破性能天花板的关键。