ESPIM架构：稀疏计算与存内计算融合，突破边缘AI推理内存墙

1. 项目概述：当稀疏计算遇上存内计算

在边缘设备上部署大型语言模型（LLaMA、GPT等）进行推理，正成为一个越来越普遍的需求。无论是出于隐私保护，还是为了应对有限的无线带宽，本地化推理都展现出巨大吸引力。然而，一个核心矛盾随之浮现：这些模型的推理过程，尤其是其核心的矩阵-向量乘法（MV）操作，严重受限于内存带宽。数据在处理器和内存之间来回搬运所消耗的时间和能量，常常远超实际计算本身。

存内计算（PIM）技术被视为解决这一“内存墙”问题的利器。其核心理念简单而深刻：将计算单元直接放置在内存芯片内部。这样一来，数据无需经过狭窄的内存引脚（Pin）进行传输，可以直接在内存阵列附近被处理，从而充分利用DRAM内部高达外部引脚带宽数十倍甚至上百倍的内部带宽。近年来，三星的FIM和SK海力士的Newton（AiM）等商业化PIM产品的出现，标志着这项技术正从实验室走向现实。

与此同时，模型稀疏化是另一个提升推理效率的经典手段。通过剪枝等技术，我们可以将权重矩阵中高达80%-90%的元素置为零，从而大幅减少计算量，且对模型精度影响甚微。非结构化稀疏因其能保留更高的模型精度而备受青睐。

那么，一个自然的想法是：能否将PIM的高带宽优势与稀疏化的低计算量优势结合起来，实现“1+1>2”的效果？这正是ESPIM（面向稀疏机器学习推理的高效能存内计算架构）所要回答的问题。然而，结合之路并非坦途。稀疏性，尤其是非结构化稀疏，引入了不确定性、不规则性和负载不均衡，这与传统PIM架构（如Newton）所依赖的全存储体（Bank）同步、规整数据流的设计哲学格格不入。简单地将稀疏矩阵扔给为稠密计算优化的PIM，性能可能不升反降。

ESPIM正是为了破解这一难题而生。它不是一个推翻重来的设计，而是在继承Newton“无头”（Headless，即仅在DRAM中放置数据通路，控制由主机负责）这一高效、低面积开销架构思想的基础上，针对稀疏计算的特点进行了一系列精巧的改造。接下来，我们将深入拆解ESPIM是如何在严格的PIM面积与功耗约束下，驯服非结构化稀疏这头“猛兽”的。

2. 核心挑战与设计思路拆解

要理解ESPIM的设计，首先必须看清稀疏计算给传统PIM架构带来的具体挑战。我们以典型的稠密PIM架构Newton作为参照系。

2.1 传统稠密PIM的工作模式与稀疏化的冲突

Newton的工作流程可以概括为“广播-读取-计算”的锁步循环：

1. 向量广播：主机将输入向量的一个切片（例如16个元素）通过DRAM内部总线广播到所有存储体（Bank）。
2. 矩阵列读取：所有存储体同步地从各自激活的DRAM行中，读取一列矩阵数据（同样是16个元素，与向量切片宽度匹配）。
3. 乘积累加（MAC）：每个存储体内的16个MAC单元，将接收到的向量切片元素与本地读取的矩阵列元素一一对应相乘，并将结果累加到本地的累加器中。
4. 循环：重复步骤1-3，直到处理完当前DRAM行的所有列。然后更换DRAM行，继续处理，直到整个矩阵-向量乘法完成。

这种设计的巧妙之处在于，它通过粗粒度数据交织实现了向量数据的重用。同一个向量切片被广播给所有存储体，用于计算各自矩阵行与这个向量的内积。计算是高度同步和规整的：一次广播对应一次列读取，产生一个部分内积。

然而，当权重矩阵变得稀疏（例如90%的零元素）时，问题出现了：

• 挑战一：广播带宽需求爆炸。在稠密情况下，一次列读取（16个元素）正好消耗一个向量切片（16个元素）。但在90%稀疏度下，一次列读取平均只包含1.6个非零元素。为了找到与这1.6个非零元素匹配的向量元素，理论上需要扫描大约10个向量切片（160个元素）。如果沿用Newton的锁步节奏，就需要进行10次广播才能完成一次列读取的计算，这完全抵消了稀疏化带来的计算量减少优势。
• 挑战二：计算不规则与同步失配。每个存储体中非零元素的位置是随机的、不同的。一次广播的向量切片，可能在A存储体中被一个非零元素使用，在B存储体中被三个非零元素使用，在C存储体中则完全没有被使用。这导致各个存储体消耗向量切片的速度不一致。为了确保计算正确性，必须等到所有存储体都“用完”当前广播的切片后，才能进行下一次广播。这种数据依赖的停顿在稠密PIM中不存在，而在稀疏PIM中成为常态，且动态变化，难以用简单的硬件逻辑处理。
• 挑战三：数据匹配延迟。由于向量切片是按顺序广播的，一个存储体当前读取到的矩阵非零元素，其匹配的向量元素可能位于未来某个切片中。这会导致MAC单元空转，等待所需向量数据的到来。
• 挑战四：选择逻辑复杂化。在稠密情况下，广播切片中的第i个元素固定对应列读取中的第i个矩阵元素。在稀疏情况下，这个一一对应关系被打破。每个MAC单元需要根据其非零元素的列索引，从广播的16个元素中“挑选”出正确的那一个。一个直观但粗暴的实现需要一个巨大的16x11交叉开关（假设每存储体11个MAC），这在面积和功耗上都是不可接受的。

2.2 ESPIM的核心设计哲学

面对上述挑战，ESPIM的设计遵循了几个核心原则：

1. 继承无头架构，控制复杂度：坚持将复杂的控制逻辑（调度、依赖判断）放在主机端，DRAM内部仅保留精简、规则的数据通路。这是控制芯片面积和功耗的关键。
2. 利用稀疏的静态性：虽然非零元素的位置是数据依赖的，但对于一个训练好的模型，其权重矩阵的稀疏模式是静态且已知的。这意味着我们可以在模型部署前（推理阶段之前），通过一次性的、离线的分析，预知所有计算的不规则性。
3. 改造数据布局与流程，而非推翻硬件：在硬件改动尽可能小的前提下，通过重新组织数据在内存中的存放方式（数据布局）和计算流程，来适应稀疏计算的特点。

基于这些原则，ESPIM提出了四大关键技术贡献，我们将在下一章逐一深入解析。

3. ESPIM关键技术深度解析

3.1 精细粒度交织布局：共享广播，回收带宽优势

这是ESPIM解决“广播带宽爆炸”挑战的核心创新。Newton的粗粒度交织是为了让一个向量切片在多个存储体间重用。ESPIM的精细粒度交织则是为了让一个向量切片在同一个存储体内的多个矩阵行间重用。

具体做法： 不再将单个稀疏矩阵行的所有非零元素连续存放在一个DRAM行中。相反，我们将k个连续稀疏矩阵行的第一个非零元素打包在一起，放在一个DRAM行中；然后是这k个行的第二个非零元素，以此类推。在ESPIM的配置中，k被设置为16（与向量切片宽度一致）。

工作原理与优势：

• 行间重用：在这种布局下，一次向量切片广播，可以同时用于计算k个不同矩阵行（对应k个MAC单元）的部分内积。每个MAC单元负责一个矩阵行的完整内积计算。
• 减少广播次数：由于一次广播服务了k行，平均每行所需的广播次数降至原来的1/k。在90%稀疏度下，虽然每列读取的非零元素变少，但由于一行计算被分摊到k次广播中，每次广播服务的有效非零元素数得以提升。这从根本上恢复了稀疏化带来的优势。
• 输出缓冲代价：付出的代价是，每个存储体需要维护一个k元素的输出向量（而不是Newton的一个标量），用于暂存这k行的部分和。但这部分开销是可控的。

一个生活化的类比： 想象一个食堂有16个打饭窗口（存储体），原来每个窗口只为一个班级（矩阵行）服务。送餐车（向量广播）送来16道菜（向量切片），每个窗口只取对应自己班级的那一道。如果班级里很多人请假（稀疏），窗口可能空等。现在，我们让每个窗口同时为16个班级服务。送餐车送来16道菜，每个窗口根据自己服务的16个班级的不同需求，从这16道菜里各取所需。这样，送餐车来一趟，就能满足更多需求，效率大幅提升。

3.2 静态数据依赖调度：将动态不确定性转为静态确定性

这是ESPIM实现“无头架构”的关键，它巧妙地将运行时棘手的数据依赖问题，转化为部署前可解的静态调度问题。

SDDS的核心思想： 既然稀疏模式是静态已知的，我们可以在模型训练/剪枝之后、部署推理之前，运行一次周期精确的模拟。这个模拟器会“预演”整个稀疏矩阵-向量乘法在ESPIM硬件上的执行过程。

调度器的具体工作：

1. 压缩与布局：根据精细粒度交织规则，将原始的稀疏矩阵压缩并布局到DRAM中。
2. 生成命令流：模拟器会追踪每个存储体中每个MAC单元的iFIFO和eFIFO（见3.3节）的状态，精确判断：
   • 何时进行向量切片广播。
   • 何时需要插入“广播停顿”周期（因为某些存储体还未消耗完当前切片）。
   • 何时需要插入“空转”的无效矩阵元素（Dummy Cells）以对齐时序。
   • 如何为简化后的交叉开关重排索引顺序以减少冲突（见3.4节）。
3. 输出调度表：最终，模拟器生成一个确定的、周期级的命令序列。这个序列就像一份详细的“乐谱”，指明了主机在何时向ESPIM发送“加载向量”、“激活行”、“带广播计算”、“不带广播计算（停顿）”、“读取结果”等命令。

带来的好处：

• 硬件极简：DRAM内部的ESPIM硬件无需任何复杂的动态调度逻辑、依赖检测电路或分支预测。它只是一个忠实的命令执行者，按照主机发来的预定序列一步步操作。这严格遵循了PIM的面积与功耗约束。
• 确定性性能：由于整个执行流程是预先确定的，推理过程的延迟是可预测的，这对于实时性要求高的边缘应用非常有利。

注意：SDDS的调度是一次性开销，发生在模型部署阶段。对于需要频繁执行推理的模型，这次前期投入的代价是完全可以接受的。

3.3 索引与数值解耦预取：隐藏向量访问延迟

为了解决“数据匹配延迟”问题，ESPIM引入了解耦预取机制。

硬件支持： 每个MAC单元配备两个小型先入先出队列（FIFO）：

• 索引FIFO：用于预取从DRAM列读取中得到的矩阵非零元素的列索引。
• 元素FIFO：用于存储根据索引从向量广播中提取出来的对应向量元素。

工作流程：

1. 预取索引：SDDS在安排数据布局时，会有意识地将矩阵非零元素的索引提前若干周期放置在其对应的数值之前。主机可以发送专门的“加载索引”命令，在不触发计算的情况下，将一批索引预取到各个MAC的iFIFO中。
2. 异步匹配：当向量切片广播时，每个MAC单元根据其iFIFO头部的索引，通过一个选择开关（见3.4节）从广播的16个元素中抓取自己需要的那个，并将其存入eFIFO。这个过程与MAC当前正在进行的计算是并行的。
3. 计算时对齐：当矩阵数值真正被读取时，它不再需要等待匹配的向量元素广播，因为该元素很可能已经预取并存储在eFIFO的头部。MAC直接从eFIFO中取出元素与数值相乘累加即可。

效果： 这相当于在向量数据流和矩阵计算流之间插入了一个小的缓冲器，将原本严格的同步关系解耦。向量元素的抓取可以提前进行，从而掩盖了因稀疏性导致的、等待特定向量切片广播而产生的延迟。

3.4 简化交叉开关：在时间维度换取空间复杂度

如前所述，一个16x11的完整交叉开关面积开销巨大。ESPIM利用DRAM时序特性中的一个关键参数tCCD来简化设计。

简化策略：tCCD是连续列读取操作之间的最小间隔，通常为4个DRAM周期。ESPIM观察到，向量切片广播后，有4个周期的窗口期可以用来进行元素选择操作。因此，它将一个庞大的16选1选择器，拆解为11个并行的4选1多路复用器。

操作时序：

1. 在一个广播周期后的4个连续周期内，每个MAC单元依次检查其所需索引是否落在[0-3],[4-7],[8-11],[12-15]这四个区间之一。
2. 每个周期，所有MAC单元只处理索引落在当前区间的请求。相应的，输入到开关的向量数据也按4个元素一组进行选择。
3. 通过4个周期的顺序操作，最终完成所有11个MAC单元对16个向量元素的筛选。

结合SDDS的优化： 这种简化带来了新的挑战：如果同一个MAC的iFIFO中连续两个索引都落在同一个区间，由于开关每个周期只服务一个区间，第二个索引就必须等待下一个广播周期，导致停顿。SDDS在静态调度时，会尝试对同一列读取中不同MAC的索引进行重排序，尽可能让连续索引分布在不同的区间，从而最大化每个广播周期的利用率，减少冲突引起的停顿。

4. 架构实现与灵活配置

4.1 完整数据通路与工作流程

结合上述所有技术，ESPIM的完整数据通路和工作流程如下：

1. 初始化：主机将输入向量的一个“向量行”加载到PIM的全局缓冲区。
2. 命令执行：主机按照SDDS生成的命令流，向ESPIM发送指令序列。
3. 核心循环： a.广播与读取：主机发送“COMP-BR”命令，触发一次向量切片广播，同时所有存储体执行一次DRAM列读取。读取的数据包含矩阵数值和预取索引。 b.索引预取与元素抓取：读取的索引被压入各MAC的iFIFO尾部。同时，各MAC根据iFIFO头部的索引，通过简化的4x11开关，从当前（或上一个被锁存的）广播切片中抓取向量元素，压入eFIFO。 c.计算：矩阵数值与eFIFO头部的向量元素进行乘积累加运算。 d.流控：根据SDDS的安排，可能插入“COMP-NoBR”命令（广播停顿），仅进行索引预取和元素抓取，或进行纯计算。
4. 结果读取：处理完一个DRAM行后，主机读取所有存储体的k元结果向量。

4.2 对稠密与稀疏模型的灵活支持

尽管专注于稀疏，但ESPIM也考虑了现实需求：并非所有模型都经过剪枝。因此，ESPIM设计了一种灵活配置模式。

硬件复用：

• 每个存储体物理上集成了16个MAC单元。
• 当运行稀疏模型时，启用其中的11个MAC，并配套其iFIFO、eFIFO和选择开关。
• 当运行稠密模型时，启用全部16个MAC，并关闭为稀疏计算服务的FIFO和开关电路（通过门控电源），以节省能耗。稠密模型的数据布局是连续的16个元素，无需索引。

数据布局兼容： 为了简化数据通路的多路复用器设计，稀疏矩阵的存储格式被设计为：连续存放11个数值，紧接着连续存放这11个数值对应的索引。这样，无论是稠密还是稀疏数据，前11个数值的位置是固定的，硬件只需在数据路径末尾添加一组2选1多路复用器，为MAC选择输入是来自广播总线（稠密模式）还是来自eFIFO（稀疏模式）。

4.3 负载均衡与其他优化

• 贪婪负载均衡：ESPIM借鉴了SparTen的思想，在将矩阵行分配给各个存储体时，不是简单地按顺序分配，而是先将矩阵行按密度（非零元素数量）排序，然后以轮询方式将最密的行和最疏的行配对分配到同一存储体。这样可以在宏观上平衡不同存储体间的计算负载，避免某些存储体因处理非常稀疏的行而过早空闲。
• 激活函数卸载：像ReLU、Softmax这类激活函数，依然由主机CPU负责执行。对于简单的逐元素操作（如ReLU），其计算可以隐藏在结果读取的延迟中。对于Softmax等需要全局扫描的操作，则在主机端进行向量化计算。
• 错误校验：与Newton类似，ESPIM假设ECC校验由主机内存控制器负责，并通过定期重载权重矩阵来应对可能的DRAM软错误。

5. 性能评估与工程启示

根据论文中的仿真实验，在LLaMA-7B模型、90%稀疏度的设定下，ESPIM展现出了显著的优势：

5.1 性能表现

• 相对于GPU：平均取得了127倍的加速比。这充分体现了PIM解决内存带宽瓶颈的巨大潜力。
• 相对于稠密PIM（Newton）：平均取得了2倍的加速比（最高达4.2倍）。这证明了ESPIM针对稀疏优化的有效性，成功将稀疏性的理论优势转化为了实际性能提升。
• 各技术贡献的收益：通过逐项启用技术的对比实验表明，精细粒度交织是性能提升的基础，解耦预取带来了稳定增益，而针对简化开关的索引重排序和贪婪负载均衡在高稀疏度下贡献了额外的性能提升。

5.2 能效与面积

• 能耗：平均比Newton降低34%（最高达63%）。能耗降低主要来源于两点：一是计算量的真实减少（稀疏），二是访问数据量的减少（压缩格式）。
• 面积开销：
  • 仅支持稀疏模式的ESPIM，其硬件逻辑面积相比传统DRAM增加了约31%，但相比Newton仅增加了不到5%。
  • 支持稀疏/稠密混合模式的灵活配置，面积开销比传统DRAM增加约40%，比Newton增加不到12%。
  • 值得注意的是，论文中使用触发器（Flip-Flop）而非更高效的SRAM来实现FIFO，因此实际的面积开销可能比报告的数字更优。

5.3 关键参数敏感性分析

• FIFO大小：FIFO深度从4增加到16，性能持续提升，尤其是在高稀疏度（90%）下提升更明显。这说明更大的FIFO能更好地吸收稀疏性带来的不规则性。但深度达到8后，收益增长曲线放缓，需要在面积开销和性能收益间取得平衡。
• 存储体数量：性能随着存储体数量增加而近乎线性提升，证明了架构的良好可扩展性。高稀疏度下的不规则性会轻微削弱扩展效率，但影响有限。

5.4 给硬件工程师的实操启示

1. 静态调度是利器：对于算法固定、数据模式已知的领域专用加速器（DSA），ESPIM的SDDS策略极具借鉴意义。将复杂的运行时调度决策提前到编译时/配置时，可以极大简化硬件设计，降低功耗和面积。这在FPGA或ASIC设计中是常用思路。
2. 利用数据特性进行布局优化：ESPIM的精细粒度交织是一个经典的“数据布局变换”优化案例。通过改变数据在内存中的存放顺序来匹配计算模式，往往能以极小的硬件代价换取显著的性能提升。在设计加速器时，应优先思考数据该如何“摆放”。
3. 面积与性能的权衡艺术：简化交叉开关的设计是硬件资源受限下的经典权衡。它通过将空间上的并行转换为时间上的顺序操作，用略微增加的计算延迟（4周期 vs 1周期）换取了巨大的面积节省。在芯片设计中，这类“时分复用”思想无处不在。
4. 为灵活性预留空间：ESPIM的混合模式设计提醒我们，专用加速器也需要一定的灵活性来适应算法演进。通过硬件模块的复用和可配置的电源门控，可以在不显著增加成本的前提下，扩大架构的适用场景。
5. 仿真与建模的重要性：ESPIM的性能高度依赖SDDS生成的优化调度表。这要求有一个周期精确的硬件模型来进行前期调度。在实际项目中，搭建一个快速、准确的仿真平台，对于探索设计空间、验证优化效果至关重要。

ESPIM架构清晰地展示了一条路径：在严格的内存工艺约束下，通过算法与硬件的协同设计，将稀疏计算高效地映射到存内计算平台上。它不仅是学术上的创新，其设计思想对正在蓬勃发展的AI芯片、近存计算芯片等领域，都有着切实的参考价值。随着稀疏化成为大型模型部署的标配，类似ESPIM这样专为稀疏PIM优化的架构，其重要性将日益凸显。