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

ESPIM架构:稀疏计算与存内计算融合,突破边缘AI推理内存墙 1. 项目概述当稀疏计算遇上存内计算在边缘设备上部署大型语言模型LLaMA、GPT等进行推理正成为一个越来越普遍的需求。无论是出于隐私保护还是为了应对有限的无线带宽本地化推理都展现出巨大吸引力。然而一个核心矛盾随之浮现这些模型的推理过程尤其是其核心的矩阵-向量乘法MV操作严重受限于内存带宽。数据在处理器和内存之间来回搬运所消耗的时间和能量常常远超实际计算本身。存内计算PIM技术被视为解决这一“内存墙”问题的利器。其核心理念简单而深刻将计算单元直接放置在内存芯片内部。这样一来数据无需经过狭窄的内存引脚Pin进行传输可以直接在内存阵列附近被处理从而充分利用DRAM内部高达外部引脚带宽数十倍甚至上百倍的内部带宽。近年来三星的FIM和SK海力士的NewtonAiM等商业化PIM产品的出现标志着这项技术正从实验室走向现实。与此同时模型稀疏化是另一个提升推理效率的经典手段。通过剪枝等技术我们可以将权重矩阵中高达80%-90%的元素置为零从而大幅减少计算量且对模型精度影响甚微。非结构化稀疏因其能保留更高的模型精度而备受青睐。那么一个自然的想法是能否将PIM的高带宽优势与稀疏化的低计算量优势结合起来实现“112”的效果这正是ESPIM面向稀疏机器学习推理的高效能存内计算架构所要回答的问题。然而结合之路并非坦途。稀疏性尤其是非结构化稀疏引入了不确定性、不规则性和负载不均衡这与传统PIM架构如Newton所依赖的全存储体Bank同步、规整数据流的设计哲学格格不入。简单地将稀疏矩阵扔给为稠密计算优化的PIM性能可能不升反降。ESPIM正是为了破解这一难题而生。它不是一个推翻重来的设计而是在继承Newton“无头”Headless即仅在DRAM中放置数据通路控制由主机负责这一高效、低面积开销架构思想的基础上针对稀疏计算的特点进行了一系列精巧的改造。接下来我们将深入拆解ESPIM是如何在严格的PIM面积与功耗约束下驯服非结构化稀疏这头“猛兽”的。2. 核心挑战与设计思路拆解要理解ESPIM的设计首先必须看清稀疏计算给传统PIM架构带来的具体挑战。我们以典型的稠密PIM架构Newton作为参照系。2.1 传统稠密PIM的工作模式与稀疏化的冲突Newton的工作流程可以概括为“广播-读取-计算”的锁步循环向量广播主机将输入向量的一个切片例如16个元素通过DRAM内部总线广播到所有存储体Bank。矩阵列读取所有存储体同步地从各自激活的DRAM行中读取一列矩阵数据同样是16个元素与向量切片宽度匹配。乘积累加MAC每个存储体内的16个MAC单元将接收到的向量切片元素与本地读取的矩阵列元素一一对应相乘并将结果累加到本地的累加器中。循环重复步骤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的设计遵循了几个核心原则继承无头架构控制复杂度坚持将复杂的控制逻辑调度、依赖判断放在主机端DRAM内部仅保留精简、规则的数据通路。这是控制芯片面积和功耗的关键。利用稀疏的静态性虽然非零元素的位置是数据依赖的但对于一个训练好的模型其权重矩阵的稀疏模式是静态且已知的。这意味着我们可以在模型部署前推理阶段之前通过一次性的、离线的分析预知所有计算的不规则性。改造数据布局与流程而非推翻硬件在硬件改动尽可能小的前提下通过重新组织数据在内存中的存放方式数据布局和计算流程来适应稀疏计算的特点。基于这些原则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硬件上的执行过程。调度器的具体工作压缩与布局根据精细粒度交织规则将原始的稀疏矩阵压缩并布局到DRAM中。生成命令流模拟器会追踪每个存储体中每个MAC单元的iFIFO和eFIFO见3.3节的状态精确判断何时进行向量切片广播。何时需要插入“广播停顿”周期因为某些存储体还未消耗完当前切片。何时需要插入“空转”的无效矩阵元素Dummy Cells以对齐时序。如何为简化后的交叉开关重排索引顺序以减少冲突见3.4节。输出调度表最终模拟器生成一个确定的、周期级的命令序列。这个序列就像一份详细的“乐谱”指明了主机在何时向ESPIM发送“加载向量”、“激活行”、“带广播计算”、“不带广播计算停顿”、“读取结果”等命令。带来的好处硬件极简DRAM内部的ESPIM硬件无需任何复杂的动态调度逻辑、依赖检测电路或分支预测。它只是一个忠实的命令执行者按照主机发来的预定序列一步步操作。这严格遵循了PIM的面积与功耗约束。确定性性能由于整个执行流程是预先确定的推理过程的延迟是可预测的这对于实时性要求高的边缘应用非常有利。注意SDDS的调度是一次性开销发生在模型部署阶段。对于需要频繁执行推理的模型这次前期投入的代价是完全可以接受的。3.3 索引与数值解耦预取隐藏向量访问延迟为了解决“数据匹配延迟”问题ESPIM引入了解耦预取机制。硬件支持 每个MAC单元配备两个小型先入先出队列FIFO索引FIFO用于预取从DRAM列读取中得到的矩阵非零元素的列索引。元素FIFO用于存储根据索引从向量广播中提取出来的对应向量元素。工作流程预取索引SDDS在安排数据布局时会有意识地将矩阵非零元素的索引提前若干周期放置在其对应的数值之前。主机可以发送专门的“加载索引”命令在不触发计算的情况下将一批索引预取到各个MAC的iFIFO中。异步匹配当向量切片广播时每个MAC单元根据其iFIFO头部的索引通过一个选择开关见3.4节从广播的16个元素中抓取自己需要的那个并将其存入eFIFO。这个过程与MAC当前正在进行的计算是并行的。计算时对齐当矩阵数值真正被读取时它不再需要等待匹配的向量元素广播因为该元素很可能已经预取并存储在eFIFO的头部。MAC直接从eFIFO中取出元素与数值相乘累加即可。效果 这相当于在向量数据流和矩阵计算流之间插入了一个小的缓冲器将原本严格的同步关系解耦。向量元素的抓取可以提前进行从而掩盖了因稀疏性导致的、等待特定向量切片广播而产生的延迟。3.4 简化交叉开关在时间维度换取空间复杂度如前所述一个16x11的完整交叉开关面积开销巨大。ESPIM利用DRAM时序特性中的一个关键参数tCCD来简化设计。简化策略tCCD是连续列读取操作之间的最小间隔通常为4个DRAM周期。ESPIM观察到向量切片广播后有4个周期的窗口期可以用来进行元素选择操作。因此它将一个庞大的16选1选择器拆解为11个并行的4选1多路复用器。操作时序在一个广播周期后的4个连续周期内每个MAC单元依次检查其所需索引是否落在[0-3],[4-7],[8-11],[12-15]这四个区间之一。每个周期所有MAC单元只处理索引落在当前区间的请求。相应的输入到开关的向量数据也按4个元素一组进行选择。通过4个周期的顺序操作最终完成所有11个MAC单元对16个向量元素的筛选。结合SDDS的优化 这种简化带来了新的挑战如果同一个MAC的iFIFO中连续两个索引都落在同一个区间由于开关每个周期只服务一个区间第二个索引就必须等待下一个广播周期导致停顿。SDDS在静态调度时会尝试对同一列读取中不同MAC的索引进行重排序尽可能让连续索引分布在不同的区间从而最大化每个广播周期的利用率减少冲突引起的停顿。4. 架构实现与灵活配置4.1 完整数据通路与工作流程结合上述所有技术ESPIM的完整数据通路和工作流程如下初始化主机将输入向量的一个“向量行”加载到PIM的全局缓冲区。命令执行主机按照SDDS生成的命令流向ESPIM发送指令序列。核心循环 a.广播与读取主机发送“COMP-BR”命令触发一次向量切片广播同时所有存储体执行一次DRAM列读取。读取的数据包含矩阵数值和预取索引。 b.索引预取与元素抓取读取的索引被压入各MAC的iFIFO尾部。同时各MAC根据iFIFO头部的索引通过简化的4x11开关从当前或上一个被锁存的广播切片中抓取向量元素压入eFIFO。 c.计算矩阵数值与eFIFO头部的向量元素进行乘积累加运算。 d.流控根据SDDS的安排可能插入“COMP-NoBR”命令广播停顿仅进行索引预取和元素抓取或进行纯计算。结果读取处理完一个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解决内存带宽瓶颈的巨大潜力。相对于稠密PIMNewton平均取得了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 给硬件工程师的实操启示静态调度是利器对于算法固定、数据模式已知的领域专用加速器DSAESPIM的SDDS策略极具借鉴意义。将复杂的运行时调度决策提前到编译时/配置时可以极大简化硬件设计降低功耗和面积。这在FPGA或ASIC设计中是常用思路。利用数据特性进行布局优化ESPIM的精细粒度交织是一个经典的“数据布局变换”优化案例。通过改变数据在内存中的存放顺序来匹配计算模式往往能以极小的硬件代价换取显著的性能提升。在设计加速器时应优先思考数据该如何“摆放”。面积与性能的权衡艺术简化交叉开关的设计是硬件资源受限下的经典权衡。它通过将空间上的并行转换为时间上的顺序操作用略微增加的计算延迟4周期 vs 1周期换取了巨大的面积节省。在芯片设计中这类“时分复用”思想无处不在。为灵活性预留空间ESPIM的混合模式设计提醒我们专用加速器也需要一定的灵活性来适应算法演进。通过硬件模块的复用和可配置的电源门控可以在不显著增加成本的前提下扩大架构的适用场景。仿真与建模的重要性ESPIM的性能高度依赖SDDS生成的优化调度表。这要求有一个周期精确的硬件模型来进行前期调度。在实际项目中搭建一个快速、准确的仿真平台对于探索设计空间、验证优化效果至关重要。ESPIM架构清晰地展示了一条路径在严格的内存工艺约束下通过算法与硬件的协同设计将稀疏计算高效地映射到存内计算平台上。它不仅是学术上的创新其设计思想对正在蓬勃发展的AI芯片、近存计算芯片等领域都有着切实的参考价值。随着稀疏化成为大型模型部署的标配类似ESPIM这样专为稀疏PIM优化的架构其重要性将日益凸显。