1. AI芯片软硬件协同设计的核心逻辑
1.1 为什么软硬件必须一起设计
做AI芯片这行的人都有一个共识:硬件堆算力不难,难的是让软件能把硬件的算力真正吃满。我见过太多团队,流片回来的芯片理论算力标称几百TOPS,实际跑模型连三分之一都跑不到。问题出在哪?出在软硬件脱节。
传统芯片设计流程是硬件团队先定义指令集、定架构,然后软件团队再往上适配编译器。这种串行模式在通用CPU时代还能凑合,因为CPU的灵活性足够高,编译器有足够的优化空间。但到了AI芯片领域,尤其是做深度学习加速器,这个逻辑就完全行不通了。AI负载的特征太鲜明了——大量的矩阵乘法、卷积运算、激活函数,数据流模式高度规律。如果硬件设计的时候不考虑编译器怎么写、算子怎么映射,最后出来的芯片就是一堆死算力。
软硬件协同设计(HW/SW Co-Design)的核心思想是:在架构定义阶段就让编译器团队、算法团队深度参与,把数据流、存储层次、计算单元的组织方式跟上层框架的算子实现一起考虑。举个例子,Google的TPU从第一代开始就是软硬件一起设计的典型代表。TPU的脉动阵列尺寸、片上存储大小、指令集格式,都是根据TensorFlow里最常见的算子模式反推出来的。反过来,TensorFlow的XLA编译器也会针对TPU的硬件特性做专门的算子融合和内存调度。
这个逻辑放到国内做AI芯片的团队也一样适用。你不能先拍脑袋定一个256x256的脉动阵列,然后指望编译器能把所有模型都高效映射上去。实际做的时候,你得先看目标场景——是跑推荐模型还是跑视觉模型,是训练还是推理,batch size大概多大,精度要求是什么。这些信息决定了你的阵列规模、数据位宽、片上缓存策略。
1.2 从算法到硅片:一条完整的设计链路
一个AI芯片项目从启动到流片,大致要经过这么几个阶段:
算法分析与负载建模。这个阶段要回答的问题是:我们的目标模型有哪些?它们的计算特征是什么?比如Transformer类模型,核心是QKV矩阵乘和FFN层的大矩阵乘,计算密度高,对带宽需求相对可控。而卷积神经网络,尤其是depthwise卷积,计算密度低,对内存带宽极其敏感。你得把这些特征量化出来,形成负载模型。
架构探索与性能建模。基于负载模型,开始设计计算阵列、存储层次、互联结构。这个阶段通常用C++或Python写一个周期精确或者近似周期精确的模拟器,快速迭代不同的架构参数。比如脉动阵列的大小、PE的MAC数量、SRAM的bank划分方式。每次改参数,跑一遍负载模型,看吞吐、延迟、能耗的变化。
指令集与编译器设计。架构基本定型后,开始定义指令集。AI芯片的指令集通常分两层:一层是粗粒度的算子级指令(比如“执行一个卷积”),另一层是细粒度的微指令(控制数据搬运、PE阵列的配置)。编译器负责把上层框架的图切分成算子,再把算子映射成指令序列。
RTL实现与验证。指令集和微架构确定后,硬件团队开始写RTL。这个阶段最怕的是发现某个算子映射效率极低,回头改架构成本巨大。所以前期软件团队必须把主流算子都在模拟器上跑通,确认没有性能悬崖。
流片与回片调试。芯片回来之后,软件团队要快速把编译器后端对接上,跑真实模型。这时候经常发现模拟器和实际芯片有偏差,比如某个数据通路的带宽比预期低,或者某个指令的延迟比建模时多。这些问题需要软硬件团队一起定位。
这条链路里,任何一个环节脱节都会导致最终芯片的实际性能远低于预期。我个人的经验是,架构探索阶段至少要留出整个项目周期的30%时间,而且软件团队必须从第一天就介入。
1.3 当前主流AI芯片架构的取舍
市面上做AI芯片的公司,架构路线大致分几类:
脉动阵列派。以Google TPU为代表,核心是一个二维的PE阵列,数据从左边和上边流入,在阵列中流动的过程中完成乘累加。这种结构的优势是数据复用率高,权重和激活值可以在阵列中多次使用,减少对内存的访问。缺点是灵活性差,阵列一旦固定,很难高效处理稀疏或者不规则的计算。
SIMD/SIMT派。以GPU为代表,大量的小核心并行执行,通过warp调度隐藏延迟。优势是通用性强,能处理各种算子。缺点是能效比相对低,因为指令调度和寄存器堆的开销大。
数据流架构派。以一些初创公司的方案为代表,根据算子的数据依赖关系动态配置计算单元和存储单元。灵活性介于前两者之间,但编译器复杂度极高。
存内计算派。把计算单元嵌入到SRAM或者DRAM里面,减少数据搬运。这个方向学术界很热,但工程落地还有距离,主要问题是工艺一致性和良率。
选择哪种架构,取决于目标场景。如果是做云端推理,追求极致能效比,脉动阵列或者数据流架构更合适。如果是做边缘端,模型变化快,SIMD路线可能更稳妥。没有绝对的好坏,只有适不适合。
2. 脉动阵列的硬核原理与设计细节
2.1 脉动阵列到底是怎么“脉动”的
脉动阵列(Systolic Array)这个概念最早是H.T. Kung在1982年提出的,当时是为了解决VLSI时代计算单元和内存之间带宽不匹配的问题。它的核心思想是:让数据像血液在血管里一样,有节奏地流过计算阵列,每个PE(Processing Element)在数据流过的瞬间完成一次乘累加,然后把结果传给下一个PE。
我拿一个最简单的4x4脉动阵列做矩阵乘法来举例。假设我们要算C = A × B,其中A是4x4,B是4x4。在脉动阵列里,A的元素从左边流入,B的元素从上边流入。每个PE内部有一个乘法器和一个累加器。在每一个时钟周期,A的元素向右移动一格,B的元素向下移动一格。PE(i,j)在第t个周期接收到的A元素和B元素相乘,累加到自己的寄存器里。
关键点在于:A的每一行元素是错开时钟周期进入阵列的,B的每一列元素也是错开进入的。这样设计的结果是,当A(i,k)和B(k,j)在PE(i,j)相遇时,正好是它们应该相乘的时刻。整个计算过程不需要任何全局的地址广播,数据在阵列内部自然流动,极大地减少了控制逻辑和内存访问。
用生活化的类比:想象一个工厂流水线,每个工位(PE)只负责一道工序(乘累加),原料(数据)从流水线的一端进入,经过每个工位时被加工一次,最终从另一端出来就是成品(计算结果)。工位之间不需要互相喊话协调,因为流水线的节奏是固定的。
脉动阵列的尺寸选择是个权衡。阵列越大,数据复用率越高,但利用率可能越低。比如一个256x256的阵列,跑一个batch size为1的矩阵乘,很多PE会闲置。实际设计中,通常会根据目标模型的最大矩阵维度来定阵列大小,同时支持把大矩阵切分成小块,分时复用阵列。
2.2 数据复用与带宽瓶颈的博弈
脉动阵列最大的优势是数据复用。在一个NxN的阵列里,每个权重值可以被复用N次(沿着列方向流动),每个激活值也可以被复用N次(沿着行方向流动)。这意味着从片外内存读取的数据量可以减少到原来的1/N。对于计算密集型的矩阵乘,这个复用率直接决定了芯片的能效比。
但复用率高不代表没有瓶颈。实际做设计的时候,你会发现真正的瓶颈往往在片上存储的带宽和容量上。举个例子:一个256x256的脉动阵列,每个周期需要从左边流入256个激活值,从上边流入256个权重值。如果PE是8位乘法器,那每个周期需要512字节的输入带宽。这个带宽如果全部从SRAM读,SRAM的bank数量和端口数就得精心设计,否则根本喂不饱阵列。
更麻烦的是,矩阵乘只是模型的一部分。卷积、激活、归一化这些算子也需要存储和带宽。所以实际芯片的片上存储通常分多级:寄存器堆、PE本地缓存、全局SRAM、最后才是片外DRAM。每一级的带宽和延迟都不一样,编译器需要根据算子的数据依赖关系,把数据在不同层级之间调度。
我踩过的一个坑是:早期设计的时候只关注了阵列的峰值算力,忽略了SRAM的带宽匹配。结果流片回来发现,跑大矩阵乘的时候阵列利用率只有60%,因为SRAM的读取速度跟不上。后来在架构里加了一级ping-pong buffer,把数据预取和计算重叠起来,利用率才提到85%以上。
2.3 脉动阵列的变体与优化方向
基础的脉动阵列是权重固定(Weight Stationary)的,权重在PE里不动,激活值流动。但实际芯片里,根据算子不同,还有输出固定(Output Stationary)和行固定(Row Stationary)等变体。
权重固定适合权重复用率高的场景,比如全连接层。权重预加载到PE阵列里,然后一批激活值流过,每个激活值跟所有PE里的权重相乘。这种模式在推理场景很常见,因为推理时权重是固定的。
输出固定适合卷积层。每个PE负责计算一个输出像素的部分和,输入像素和权重在PE之间流动。这种模式可以减少输出部分和的搬运次数。
行固定是Eyeriss架构提出的,结合了权重固定和输出固定的优点,在行方向上复用权重,在列方向上复用激活值。适合卷积神经网络的各种层。
实际芯片设计里,很少只用一种数据流。通常是可配置的,编译器根据算子类型选择最优的数据流模式。比如矩阵乘用权重固定,卷积用行固定,depthwise卷积用输出固定。这种灵活性会增加控制逻辑的复杂度,但能显著提升整体利用率。
还有一个优化方向是稀疏化支持。真实模型的权重和激活值都有大量零值,如果PE能跳过零值计算,等效算力可以提升好几倍。但稀疏化对硬件的要求很高,需要额外的索引存储和调度逻辑。目前学术界有很多稀疏脉动阵列的方案,工程落地的还不多。
3. 数值格式:从FP32到FP8的演进逻辑
3.1 为什么AI芯片需要低精度格式
训练和推理对精度的需求完全不同。训练的时候,梯度更新需要高精度累加,否则模型收敛会出问题。所以训练芯片通常支持FP32或者BF16,累加器用FP32。推理的时候,模型权重已经固定,对精度的容忍度高很多,用FP16甚至INT8都能保持不错的准确率。
低精度格式带来的好处是直接的:数据位宽减半,内存带宽需求减半,乘法器的面积和功耗也大幅降低。一个FP32乘法器的面积大概是一个FP16乘法器的4倍,功耗是3倍左右。对于动辄几百TOPS的AI芯片,这个差距直接决定了芯片的成本和能效。
但低精度不是没有代价的。FP16的动态范围有限,遇到特别大或者特别小的值会溢出或者下溢。INT8的量化误差更明显,尤其是对激活值做量化的时候,不同层的数值分布差异很大,统一的量化参数往往效果不好。所以实际部署的时候,需要做量化感知训练(QAT)或者训练后量化(PTQ),把量化误差控制在可接受范围内。
3.2 FP8的两种格式与选择依据
FP8是最近两年最热的话题之一。它有两种主流格式:E4M3和E5M2。E4M3表示4位指数、3位尾数,E5M2表示5位指数、2位尾数。
E4M3的精度更高,但动态范围小。它的最大正规数是448,最小正规数是2^-6。E5M2的动态范围大,最大正规数是57344,最小正规数是2^-14,但精度低,只有2位尾数。
实际使用的时候,权重通常用E4M3,因为权重的数值分布相对集中,精度更重要。激活值用E5M2,因为激活值的动态范围大,需要更大的指数位来避免溢出。梯度也用E5M2,因为梯度的数值范围变化剧烈。
NVIDIA的H100和国内一些AI芯片都支持FP8。实测下来,FP8训练相比FP16训练,在大多数模型上准确率损失很小,但吞吐可以翻倍。推理场景下,FP8的收益更明显,因为推理对精度的要求更低。
但FP8不是万能的。有些模型对精度极其敏感,比如涉及大量小数值累加的模型,FP8的尾数位太少,累加误差会累积。这时候还是得用FP16或者BF16。所以好的AI芯片应该支持多种精度格式,让编译器根据模型特点自动选择。
3.3 量化格式与硬件设计的联动
数值格式的选择直接影响硬件设计。FP8乘法器比FP16乘法器小一半,但FP8的累加器通常还是用FP16或者FP32,因为累加过程需要更高的精度。这就意味着PE内部的数据通路是混合精度的:输入是FP8,乘法是FP8×FP8,累加是FP16或FP32。
这种混合精度设计对PE的布局布线有影响。乘法器和累加器之间的位宽转换需要额外的逻辑,如果处理不好会成为时序瓶颈。我见过一些设计,为了省面积把累加器也做成FP8,结果模型准确率掉得厉害,根本没法用。
另一个联动点是量化参数的存储和计算。INT8量化需要存储scale和zero_point,FP8量化需要存储scale。这些参数通常放在片上SRAM里,跟权重一起加载。编译器在生成指令的时候,要把反量化操作融合到计算流程里,避免额外的内存访问。
还有一个容易被忽略的点是溢出处理。FP8的指数位少,遇到极端值容易溢出。硬件需要支持饱和运算(saturate)或者无穷大表示,否则一个溢出值会污染整个累加结果。实际测试的时候,要专门构造一些边界case,验证溢出处理逻辑是否正确。
4. 软硬件协同的实操流程与关键环节
4.1 从模型到指令:编译器的核心工作
编译器在AI芯片里的角色,相当于一个翻译官加调度员。它要把PyTorch或者TensorFlow的模型图,翻译成芯片能执行的指令序列,同时还要做各种优化,让指令序列跑得尽可能快。
第一步是图优化。编译器会先对计算图做算子融合,把连续的Conv+BN+ReLU融合成一个算子,减少中间结果的存储和搬运。然后是常量折叠,把可以在编译期算出来的值提前算好。还有死代码消除、公共子表达式提取等等。
第二步是算子映射。每个融合后的算子,编译器要决定用芯片的哪些计算资源来执行。比如一个矩阵乘,是放到脉动阵列上跑,还是用向量单元跑。映射的时候要考虑数据布局、分块大小、循环顺序。这些决策直接影响性能。
第三步是内存分配。编译器要决定每个张量放在片上SRAM还是片外DRAM,什么时候预取,什么时候写回。这个阶段最复杂,因为片上SRAM容量有限,要尽可能把频繁访问的数据留在片上,同时避免bank冲突。
第四步是指令生成。把前面的决策翻译成具体的指令序列,包括数据搬运指令、计算指令、同步指令。指令的调度要尽量填满流水线,避免气泡。
我个人的经验是,编译器的优化空间比硬件大得多。同一颗芯片,好的编译器能让性能提升2-3倍。所以软件团队的投入不能省,尤其是做推理芯片,编译器的质量直接决定客户体验。
4.2 性能建模与瓶颈定位
性能建模是软硬件协同设计里最容易被低估的环节。很多团队觉得写个模拟器太费时间,不如直接上FPGA原型。但FPGA原型的迭代速度慢,改一个参数要重新综合几小时,根本没法做架构探索。
一个好的性能模型应该包含几个层次:计算单元的吞吐模型、存储层次的带宽和延迟模型、互联网络的冲突模型。输入是算子的参数(矩阵大小、数据位宽、分块策略),输出是周期数、能耗、利用率。
建模的精度不需要做到周期精确,但趋势要准。比如你把脉动阵列从128x128改成256x256,模型应该能预测出吞吐提升多少、利用率下降多少。这样架构师才能快速做权衡。
瓶颈定位是性能模型的另一个用途。跑完一个模型,模型会告诉你哪个算子最耗时、哪个存储层次最拥堵。常见的瓶颈有:计算阵列利用率低(通常是分块策略不好)、SRAM带宽不够(bank冲突或者端口数不足)、DRAM访问太频繁(数据复用没做好)。
我常用的一个方法是:把模型的每个算子的理论最优周期数和实际周期数对比,差距大的算子重点分析。通常80%的性能损失来自20%的算子,把这20%优化好,整体性能就能大幅提升。
4.3 回片调试的实战经验
芯片回来之后的调试阶段,是最考验软硬件协同能力的。模拟器再准,跟真实芯片也有偏差。常见的偏差来源有:时序违例导致某些路径降频、SRAM的读写冲突比建模时严重、指令发射的逻辑有bug。
调试的第一步是跑通基本功能。用一个最简单的矩阵乘,验证数据通路是通的。然后逐步增加复杂度,跑卷积、跑完整的模型。每一步都要对比模拟器的结果和实际芯片的结果,定位偏差。
第二步是性能调优。先跑一遍profiling,看每个算子的实际耗时。然后跟模拟器对比,找出偏差大的算子。常见的性能问题包括:数据预取没生效、指令流水线有气泡、SRAM bank冲突严重。
第三步是精度验证。低精度格式的芯片,必须验证量化后的模型准确率。通常用几个标准数据集(ImageNet、COCO)跑一遍,看准确率损失是否在可接受范围内。如果损失太大,可能需要调整量化策略,或者在某些层用更高的精度。
我踩过的一个坑是:回片后发现某个算子的性能只有模拟器的50%。查了很久才发现,是SRAM的某个bank在特定访问模式下有冲突,导致实际带宽减半。后来在编译器里加了一个地址重映射的pass,把冲突避开了。这个问题在模拟器里完全没暴露,因为模拟器的SRAM模型太理想化了。
5. 常见问题与排查技巧实录
5.1 脉动阵列利用率低的排查思路
脉动阵列利用率低是最常见的问题。表现是理论算力很高,但实际跑模型的时候阵列大量闲置。排查的时候按这个顺序来:
先看矩阵维度。如果矩阵的M、N、K维度跟阵列尺寸不匹配,比如阵列是256x256,但矩阵是100x100,那大部分PE会闲置。解决办法是把小矩阵拼成大的,或者用更小的阵列配置。
再看数据流。权重固定模式下,如果权重加载时间太长,阵列会等数据。检查权重预加载的指令是否跟计算重叠了。理想情况下,权重加载应该在上一轮计算还没结束的时候就完成。
然后看分块策略。大矩阵切成小块的时候,块的大小要跟阵列尺寸匹配。如果块太小,阵列利用率低;如果块太大,片上存储放不下,会频繁换入换出。
最后看同步开销。阵列计算完一个块之后,需要同步才能开始下一个块。如果同步逻辑太重,气泡会很多。可以考虑用双缓冲或者异步流水线来隐藏同步开销。
5.2 FP8精度损失的定位与补偿
FP8精度损失通常表现为模型准确率下降。定位的时候,先做逐层分析:把每一层的输入输出跟FP32的结果对比,看哪一层的误差最大。
常见的误差来源有几个:一是权重或激活值的动态范围超出了FP8的表示范围,导致溢出。这时候需要调整scale,或者对这一层用更高的精度。二是累加器的精度不够,大量小数值累加的时候误差累积。解决办法是把累加器升级到FP16或FP32。三是量化参数的粒度太粗,比如整个张量用一个scale,但张量内不同区域的数值分布差异很大。可以改用per-channel或者per-group的量化。
补偿的方法包括:量化感知训练(QAT),在训练的时候模拟量化误差,让模型适应低精度。混合精度,对敏感层用FP16,其他层用FP8。还有残差补偿,把量化误差作为残差传到下一层,在下一层补偿回来。
5.3 软硬件接口的常见坑
软硬件接口是很多问题的根源。我整理了一个速查表:
| 问题现象 | 可能原因 | 排查方法 | 解决方案 |
|---|---|---|---|
| 指令执行结果错误 | 指令编码和解码不一致 | 对比RTL仿真和模拟器的指令trace | 统一指令集定义,用同一份头文件 |
| 数据搬运丢失 | 地址对齐问题 | 检查DMA的地址和长度寄存器 | 强制地址对齐,或者支持非对齐访问 |
| 性能远低于预期 | 编译器调度不合理 | profiling看流水线气泡 | 优化指令调度,增加并行度 |
| 精度不达标 | 量化参数配置错误 | 逐层对比量化前后输出 | 调整scale和zero_point |
| 芯片发热严重 | 时钟门控没生效 | 检查空闲模块的时钟使能 | 增加细粒度时钟门控 |
| 多核同步失败 | 同步指令的语义不一致 | 用示波器抓同步信号 | 明确同步原语的内存序 |
这些坑我基本都踩过。最麻烦的是指令编码不一致,因为模拟器和RTL是两拨人写的,很容易出现理解偏差。后来我们强制要求指令集定义用同一份YAML文件,两边都从这份文件生成代码,问题就少了很多。
5.4 给新入行同学的建议
如果你刚入行做AI芯片,我的建议是:先把一个完整的链路跑通,哪怕是很小的一个设计。从算法分析开始,写一个简单的性能模型,定义一个最小指令集,写RTL,跑仿真,最后在FPGA上验证。这个过程中你会遇到各种问题,但每解决一个,你对软硬件协同的理解就深一层。
不要一上来就追求大而全。我见过太多项目,架构设计得极其复杂,结果流片回来一堆bug,根本跑不起来。反而是那些架构简单、但软硬件打磨得很好的芯片,实际表现更出色。
还有一点:多跟做算法的同学交流。AI芯片的最终目标是跑模型,如果不懂模型的特点,硬件设计就是空中楼阁。我每周都会花时间看最新的模型论文,了解算子层面的变化趋势。这个习惯让我在设计架构的时候能提前预判需求,而不是等流片回来才发现不支持某个新算子。
最后,工具链的投入不能省。一个好的模拟器和编译器,能让整个团队的效率提升好几倍。我见过一些团队,为了省人力,用Excel做性能建模,结果架构决策全靠拍脑袋,流片回来性能不达标,损失的钱够养十个工具链团队。这个账要算清楚。