1. 项目概述:当“闪电”划破存内计算的天空
最近在ISSCC 2024上读到一篇论文,编号34.3,标题里带了个挺酷的名字——“闪电”。这篇论文讲的是一个数模混合存内计算芯片,专门为了适配当下火热的Transformer和CNNs架构而设计。作为一个在芯片设计和AI加速领域摸爬滚打了十几年的老工程师,我第一眼看到这个标题就来了兴趣。存内计算喊了这么多年,真正能高效、通用地支持Transformer这种复杂模型的方案,市面上其实并不多见。大多数存内计算芯片要么偏重模拟域的高能效但精度和灵活性受限,要么偏重数字域的灵活可控但能效提升有限。“数模混合”这个路子,听起来就像是试图在鱼和熊掌之间找到那个微妙的平衡点。
这篇“闪电”芯片的论文,在我看来,其核心价值在于它不仅仅是一个电路设计,更是一套针对现代AI计算范式(尤其是Transformer中的注意力机制和CNNs中的卷积)的系统级解决方案。它试图回答一个关键问题:如何让存内计算架构不再只是实验室里的“玩具”,而是能真正适应主流、多变、计算密集的AI模型?对于芯片架构师、AI算法工程师,甚至是关注前沿硬件的投资者来说,理解这套设计思路都至关重要。它揭示了下一代AI加速器可能的一个进化方向:更贴近数据本身的计算范式,以及如何通过巧妙的电路与架构协同,来驯服Transformer这头“计算巨兽”。接下来,我就结合自己的经验,为大家拆解一下这篇论文里的门道。
2. 核心思路与架构选型:为什么是“数模混合”?
2.1 存内计算的传统困境与Transformer的新挑战
要理解“闪电”芯片的选型,我们得先看看存内计算面对的现实。传统的数字存内计算,比如基于SRAM或数字RRAM的架构,其优势是精度高、抗噪声能力强、与现有数字流程兼容性好。你可以把它想象成一个高度组织化、纪律严明的工厂,每个步骤(数据搬运、计算)都清晰可控。但它的致命伤在于“数据搬运功耗”。即便计算单元就在存储器旁边,大量的中间结果、权重数据的读取和写入,依然会产生巨大的能量开销。在Transformer模型中,这种开销被放大到了极致,因为注意力机制涉及大量的矩阵乘加运算,并且数据复用率相比CNN要低。
另一方面,纯模拟存内计算(例如利用忆阻器的电导值直接进行模拟乘加)则像一个高效的“化学反应炉”,乘加运算在物理定律下自然发生,能效极高。但它的问题在于“非理想性”太强:器件偏差、噪声、非线性、温度漂移……这些因素会让计算精度大打折扣。训练好的高精度神经网络模型(比如FP32或FP16)映射到这样的模拟阵列上,性能可能会严重衰退。这对于要求高准确率的视觉或语言任务来说,是难以接受的。
Transformer模型给硬件带来了几个特有的挑战:1)动态稀疏性:注意力矩阵通常是稀疏的,但稀疏模式是动态的、与输入相关的,无法在硬件设计时静态优化。2)混合精度需求:前向推理时,不同层、甚至同一层内的不同计算(如注意力得分计算与FFN层计算)对数值范围和精度的容忍度不同。3)复杂的数据流:涉及大量的矩阵乘加(QK^T, SV)、Softmax、LayerNorm等非线性操作,数据依赖关系复杂。
2.2 “闪电”的混合策略:分工与协同
“闪电”芯片的核心思路,正是针对上述困境和挑战,提出了一种“分工明确、协同作战”的数模混合架构。它不是简单地把模拟和数字模块拼在一起,而是根据计算任务的特性和需求,进行了精密的切分。
模拟域承担什么?论文中的模拟存内计算阵列,主要负责计算密集型、数据复用率高、且对极低功耗有苛刻要求的核心矩阵乘法。具体来说,就是Transformer中Q(查询)、K(键)、V(值)矩阵之间的乘加操作,以及CNN中的卷积核与输入特征图之间的乘加。这部分计算占据了模型推理绝大部分的运算量和能耗。利用模拟计算的高并行性和高能效特性,在这里能获得最大的收益。设计的关键在于,通过电路技术(比如论文中可能提到的电流积分、电荷共享等机制)来保证模拟计算的线性度和一致性,将其精度控制在可接受的范围内(例如8-bit以内)。
数字域承担什么?数字部分则扮演了“控制器”和“精度保障者”的角色。它主要负责:1)非线性函数计算:如Softmax、GELU、LayerNorm等。这些函数用模拟电路实现极其复杂且不精确,而数字逻辑(查找表、分段近似等)可以高效、高精度地完成。2)动态稀疏处理:识别和处理注意力机制中的动态稀疏性,跳过对零值的无效计算,这部分需要灵活的逻辑判断,数字电路更擅长。3)精度校准与纠错:对模拟计算的结果进行数字化读取,并可能通过轻量级的数字后处理电路(如校准查找表、误差补偿单元)来修正模拟非理想性带来的误差。4)全局控制与数据调度:管理整个芯片的数据流,将任务分发给模拟阵列,并收集结果。
这种混合架构的本质,是让模拟去做它最擅长的事(大量并行的乘加),让数字去处理它更擅长的任务(控制、非线性、精度提升),从而在系统层面实现能效与灵活性的帕累托最优。它避免了纯模拟方案的精度灾难,也规避了纯数字方案的数据搬运瓶颈。
3. 核心电路与关键技术点拆解
3.1 模拟存内计算阵列:如何实现高能效的矩阵乘加?
论文中模拟计算阵列的具体实现,是技术的重中之重。虽然我无法获知“闪电”芯片的全部电路细节,但基于常见的数模混合存内计算设计,我们可以推断其核心可能围绕以下几种技术:
1. 基于SRAM或新型非易失存储器的模拟计算单元:
- SRAM-based:利用SRAM单元内部分配的电荷或电流来进行乘加运算。例如,将输入激活值转换为脉冲宽度或幅度,将权重存储在SRAM的特定配置下,通过位线(BL)和字线(WL)上的操作,在阵列中同时完成多个乘积累加(MAC)操作。其优势是与CMOS工艺完全兼容,速度快。
- 非易失存储器(NVM)-based:如使用ReRAM(阻变存储器)或FeFET(铁电晶体管)。将权重以电导值(G)的形式存储在NVM器件中。输入电压(V)施加在字线上,根据欧姆定律(I = V * G),在位线上产生的电流自然就是乘法结果。同一列所有位线的电流求和(基尔霍夫电流定律),就完成了乘加运算。这种方案密度高,但器件非理想性更显著。
2. 模拟-数字转换器(ADC)的精妙设计:模拟阵列计算出的结果是模拟量(通常是电流或电压),必须被转换为数字量才能被后续数字电路处理。ADC是功耗和面积的大户,也是设计瓶颈。“闪电”芯片很可能采用了以下策略来优化ADC:
- 列共享ADC:并非每个存储单元列都配备一个ADC,而是多列共享一个高速高精度ADC,通过模拟多路复用器轮流读取,以节省面积。
- 逐次逼近型(SAR)ADC:因其在中等精度(如8-10位)下具有优异的能效比,是此类应用的常见选择。
- 利用计算本身的特点:例如,在注意力计算中,有时只需要知道最大值或相对大小(如Softmax前),可能采用更简单的比较器电路而非全精度ADC。
3. 输入/输出(IO)的数据映射与调度:如何将神经网络中的权重矩阵和输入特征图映射到物理的存储阵列上,直接影响计算效率和硬件利用率。这里涉及:
- 权重固定:通常将训练好的网络权重预先编程(烧写)到非易失存储器中,或加载到SRAM阵列的特定配置里。对于Transformer,这意味着将Q、K、V的投影矩阵权重固化。
- 输入流式处理:输入特征(如词嵌入向量或图像块)被转换为模拟信号(电压/电流),按特定顺序施加到阵列的行上。
- 脉动阵列思想:可能借鉴了数字脉动阵列的数据流,在模拟域实现数据的高效流水线处理,隐藏数据访问延迟。
实操心得:模拟计算的“校准”是关键在实际流片中,模拟阵列的性能一致性是最大的挑战之一。由于制造工艺的微小偏差,每个存储单元的电学特性(如电导值、晶体管阈值电压)都会有差异。因此,芯片必须包含一个上电或周期性的校准环节。这个环节通常由数字控制电路发起,向阵列输入一组已知的测试向量,测量输出,然后生成一个“校准系数查找表”。在正式计算时,用这个查找表对原始输出进行补偿。忽略这一步,芯片的推理准确率可能会惨不忍睹。
3.2 数字处理引擎:灵活性与精度的守护者
数字部分的设计,直接决定了芯片能否顺畅运行复杂的Transformer和CNN模型。
1. 专用指令集与可编程控制器:“闪电”芯片很可能定义了一套精简的专用指令集,用于描述神经网络层(如Linear, Conv2D, Attention, Softmax, LayerNorm)的操作。一个轻量级的RISC-V类微控制器或定制的状态机负责解析这些指令,并产生控制信号,调度模拟阵列、数字计算单元、片上缓存(SRAM Buffer)协同工作。这种设计提供了必要的灵活性,以支持不同的模型架构和层类型。
2. 非线性函数硬件加速器:
- Softmax单元:这是Transformer的核心。硬件实现Softmax通常采用“减最大值”技巧来保证数值稳定性。即先找到向量中的最大值,所有元素减去该最大值后再进行指数运算和归一化。硬件上需要并行比较树、查找表(LUT)实现指数函数、以及除法器或倒数单元。
- GELU/LayerNorm单元:GELU激活函数可以用多项式近似或查找表实现。LayerNorm涉及均值和方差的计算(归约操作),以及按元素缩放和平移,需要乘加器和一些控制逻辑。
3. 稀疏化处理单元:为了利用注意力矩阵的稀疏性,芯片需要能够快速识别并跳过零值运算。一种可能的设计是在数字部分集成一个轻量级的“稀疏性检测”模块。它可以在模拟计算之前,对输入数据(如Q和K的向量)进行预判,或者对模拟计算出的中间结果(如注意力得分)进行阈值比较,将低于阈值的值视为零,并生成一个“掩码”(Mask)。这个掩码会反馈给控制逻辑,用于跳过后续不必要的计算和数据搬运。
4. 精度提升与纠错电路:这是数模混合设计中的“粘合剂”。模拟计算的结果经过ADC量化后,会进入一个数字后处理管道。这里可能进行:
- 偏移与增益校正:根据校准阶段得到的参数,对数字输出进行线性校正。
- 非线性误差补偿:通过一个小的查找表,补偿模拟器件固有的非线性响应。
- 位宽扩展与饱和处理:将多个低精度(如4-bit)的模拟计算结果,在数字域拼接或累加成更高精度(如16-bit)的中间结果,以保持累积精度,同时防止溢出。
4. 针对Transformer和CNN的架构适配性设计
4.1 如何高效支持Transformer的注意力机制?
Transformer的注意力层是硬件设计的“硬骨头”。“闪电”芯片需要从数据流和计算组织上对其进行特殊优化。
1. 计算分解与流水线:标准的多头注意力(MHA)计算流程为:Q, K, V = input * W_q, W_k, W_v->Scores = Q * K^T->Attention = softmax(Scores) * V。 “闪电”芯片的模拟阵列可能被组织成多个“子阵列”或“块”,以并行处理多个注意力头。一个高效的数据流可能是:
- 阶段1:利用一部分模拟阵列,并行计算所有头的
Q、K、V投影(三个独立的矩阵乘法)。 - 阶段2:将计算好的
Q和K送入模拟阵列进行Q*K^T计算。这里面临一个挑战:K^T意味着需要按列访问K矩阵。这要求存储权重或中间结果的存储器有灵活的访问模式,或者通过转置缓冲器来实现。 - 阶段3:
Scores矩阵(可能很大)被送入数字部分进行Softmax计算。由于Softmax是逐行操作,可以流水线化处理,一行算完就立刻进入下一阶段。 - 阶段4:Softmax后的结果
Attention与V矩阵再次在模拟阵列中进行乘法,得到最终的注意力输出。
关键在于,这些阶段之间需要高度的流水线重叠,并充分利用片上缓存(SRAM)来存储中间矩阵(如Q,K,V,Scores),避免频繁访问片外DRAM。
2. 对因果注意力(Causal Attention)的支持:在GPT这类解码器模型中,需要因果注意力掩码(防止看到未来信息)。这可以在数字部分的Softmax单元中轻松实现,即在计算Softmax前,将未来位置的Score值设置为一个极大的负数(如-inf)。
3. 对KV Cache的硬件支持:在自回归生成(如文本生成)时,K和V会被缓存起来供后续步骤使用,以避免重复计算。“闪电”芯片的片上存储系统需要为这种缓存模式进行优化,提供高效且低功耗的KV缓存读写接口。
4.2 如何兼容经典的CNN卷积操作?
虽然论文重点在Transformer,但支持CNN能极大扩展芯片的适用场景(如视觉任务)。CNN的卷积操作与矩阵乘法有内在联系(通过im2col或Winograd变换),但数据模式不同。
1. 卷积到矩阵乘的映射:“闪电”芯片的数字预处理单元很可能集成了im2col或类似功能的硬件模块。它将输入特征图和卷积核的滑动窗口操作,展开成两个大的矩阵,从而可以直接调用模拟阵列进行矩阵乘法。这个过程需要大量的数据重排和复制,对片上存储的带宽和容量是考验。
2. 利用卷积的固有复用性:CNN中,卷积核权重在计算一个输出特征图时是固定不变的,输入数据则有大量重叠。芯片设计可以利用这一点,将卷积核权重固定在模拟阵列的某些行/列上,然后以流式方式输入不同的输入数据块,从而最大化权重数据的复用,减少配置开销。
3. 支持多种卷积参数:通过可编程控制器,芯片应能支持不同的卷积核尺寸(1x1, 3x3, 5x5等)、步长(Stride)和填充(Padding)。这主要依靠数字部分的数据地址生成器和调度逻辑来实现。
5. 芯片性能评估与潜在挑战
5.1 性能指标解读:能效比与吞吐量
这类研究芯片的论文,通常会报告几个关键指标:
- 峰值算力(TOPS):在特定精度(如INT8)下,芯片每秒能完成多少万亿次操作。这体现了计算吞吐能力。
- 能效比(TOPS/W):每瓦特功率能提供多少TOPS的算力。这是存内计算芯片的“命根子”,目标是远超传统冯·诺依曼架构的GPU/ASIC。
- 面积效率(TOPS/mm²):单位芯片面积提供的算力,衡量集成度。
- 在特定神经网络(如BERT, ViT, ResNet)上的精度与延迟:这是最终的用户体验指标。论文会展示在标准数据集上,芯片运行完整模型的准确率损失(与浮点基线相比)和推理速度。
“闪电”芯片的亮点,预计会体现在极高的能效比上,尤其是在处理Transformer的注意力层时。其数模混合架构的目标,就是在保持较高精度(例如,任务精度损失<1%)的前提下,将能效比提升一个数量级。
5.2 实际部署中的挑战与考量
尽管论文中的结果令人振奋,但将这样的芯片推向实际应用,还面临一系列工程挑战:
1. 工艺偏差与良率:模拟电路对工艺偏差极其敏感。在先进工艺节点(如7nm, 5nm)下,晶体管特性的微小波动会被放大,可能导致不同芯片之间,甚至同一芯片不同阵列之间的性能差异。这需要非常 robust 的电路设计和严格的工艺角(Corner)验证。良率问题会比纯数字芯片更突出。
2. 校准开销:如前所述,校准是必须的。但校准过程本身需要时间(影响启动延迟)和额外的电路(如高精度的参考源、测量电路)。在线校准(在运行时周期性进行)还会占用计算资源。如何设计一个快速、低开销、高精度的校准方案,是产品化的关键。
3. 编程模型与工具链缺失:这是所有新型架构硬件的共同难题。如何让AI算法工程师像使用PyTorch或TensorFlow调用CUDA那样,轻松地将模型部署到“闪电”芯片上?这需要一套完整的工具链:
- 编译器:将主流框架(PyTorch, TensorFlow)的模型,编译成芯片的专用指令流,并完成图优化、层融合、数据布局转换等。
- 量化工具:将FP32/FP16的模型,量化到芯片支持的混合精度格式(例如,权重8-bit,激活8-bit,部分累加16-bit),并可能需要进行量化感知训练(QAT)来保持精度。
- 模拟器/仿真器:在芯片实物出来之前,提供软件模型,让开发者评估性能和精度。
4. 系统集成挑战:这样一颗数模混合存内计算芯片,通常不会单独工作。它需要作为加速卡(PCIe)或移动设备SoC中的一个IP核来使用。这就涉及到与主机CPU的内存一致性、DMA数据传输、任务调度、功耗/热管理等复杂的系统级问题。
6. 未来展望与个人思考
“闪电”这类数模混合存内计算芯片,代表了一条非常务实的技术路径。它没有追求激进的纯模拟方案,而是采取了混合策略,在现有技术边界内寻求最大突破。从产业角度看,它的成功与否,不仅取决于芯片本身的PPA(性能、功耗、面积),更取决于其生态建设的速度。
我个人认为,这类芯片最先可能落地的场景,是对功耗极度敏感、但对延时有一定容忍度的边缘端设备。例如,智能手机上的实时语音助手、AR眼镜中的视觉处理、物联网网关中的视频分析。在这些场景下,高能效比直接转化为更长的续航时间,是硬性需求。而对于数据中心的大规模推理,虽然也有潜力,但需要解决多芯片互联、大规模部署下的可靠性以及与传统GPU生态竞争等问题。
另一个有趣的趋势是,随着Transformer模型本身不断演进(如出现更稀疏的模型、混合专家模型MoE),这类硬件也需要保持架构上的灵活性来适应。“闪电”芯片中强大的数字处理部分,正是这种灵活性的保障。未来的存内计算芯片,可能会更像一个“可重构的模拟计算阵列+可编程数字引擎”的紧密结合体,软件定义的程度会越来越高。
最后,对于想进入这个领域的工程师,我的建议是:跨学科的知识储备变得前所未有的重要。你需要既懂晶体管级的模拟电路设计,又懂计算机体系结构,还要理解神经网络算法。能够在这几个领域的交叉地带思考问题的人,将成为下一代芯片设计的核心力量。“闪电”论文的价值,就在于它为我们提供了一个绝佳的、具体的跨学科设计范例。它不仅仅是一篇电路设计论文,更是一份关于如何为特定算法范式进行硬件-软件协同设计的宣言书。