1. AI芯片项目从哪里开始:从算力需求反推架构
这个系列写到第6篇,前面对AI芯片的基本原理、编程模型、工具链都有过铺垫,这篇打算硬核一点,直接聊聊我在搭建AI芯片软硬件系统时最核心的一些判断逻辑和实操取舍。很多人拿到一个AI芯片项目,第一反应是"用哪种加速器架构""上不上存内计算""矩阵单元做多大",实际上我在实际项目里踩过很多次坑之后发现,最该先做的事是把工作负载特征彻底摸清楚,让需求去反推架构,而不是先定了架构再去适配模型。
为什么这么说?AI芯片和通用CPU最大的区别在于,CPU面对的是不确定的工作负载,设计目标是把指令流执行得足够快;而AI芯片面对的是一个相对收敛的算子集合,卷积、矩阵乘、激活、归一化、注意力机制,就这些玩意的组合。这意味着我们可以把这些算子的数据流特征算得很清楚,访存量算得很精确,甚至能推算出一次推理总共需要搬运多少字节的数据、产生多少MAC运算。有了这些数字,芯片的面积、功耗、内存层次、主频、带宽约束就都能落到一个具体的工程区间里。
我通常把一个AI芯片项目的需求拆成四块:峰值算力、有效算力、访存带宽、能效比。峰值算力是一个宣传数字,等于MAC阵列数量乘主频乘2;有效算力才是用户真正感知到的性能,它取决于数据能不能喂得进去、计算单元利用率高不高;访存带宽决定了数据搬运的上限;能效比直接决定芯片能不能在目标散热条件下长期运行。一个成熟的软硬件团队,在这四个数字上一定是在一个合理的平衡区间里反复校准,而不是单点堆参数。
1.1 先把目标模型的算子统计做出来
做架构设计前,我习惯先做一份算子统计表。把目标跑通的一批主流模型,比如ResNet系列的卷积模型、BERT/Transformer系列的注意力模型、还有轻量级检测模型,用Profiling工具跑一遍,统计每个算子的耗时占比、计算量(MACs)、访存量(Bytes)、参数规模。这份表格就是整个架构设计的输入。
表格里最关键的三个指标是:计算密度(MACs/Bytes)、算子分布(卷积占比多少、矩阵乘占比多少)、张量形状分布。这几个指标直接决定了你有多少比例的计算可以复用数据、需要多大的片上缓存、矩阵单元的位宽和形状怎么定。比如某个模型卷积占比超过70%,那数据复用路径的设计优先级就远高于稀疏计算支持;如果Transformer的矩阵乘占了主头,那计算核心更偏向大矩阵乘法和高吞吐的访存调度。
这些数字做出来之后,再去看架构选型就有据可依了。举个例子,假设目标场景要求跑一个7B参数的Transformer模型,单次推理总计算量大约是14 TFLOPs(乘加),权重复用次数等于序列长度乘batch size。如果batch size取1、序列长度4096,那每层权重都会从HBM里被反复读取,访存瓶颈就非常明显,这时候你需要在芯片上堆large SRAM、做权重驻留,甚至考虑模型切分到多核并行,而不是一味追求MAC阵列的面积。
1.2 算力与带宽的预算公式
算力预算有一组我从项目里验证过很多次的公式。假设工艺目标频率是1GHz,MAC阵列是128x128,那峰值算力大概是128x128x1Gx2,等于32.768 TOPS。这里的2是乘加算两次浮点操作。但如果你的片上SRAM带宽设计不够,数据搬运跟不上,每拍指令发起一次矩阵乘,实际上阵列大部分时间在等数据,能效利用率可能只有40%到60%。所以带宽预算是紧接着算力预算要做的第二件事。
具体公式是:一次矩阵乘需要的带宽约等于输入特征图字节数加权重字节数加输出字节数,再除以单次计算耗时。以128x128 MAC阵列跑FP16为例,单拍能完成128x128x2字节约32KB的数据吞吐需求。如果主频1GHz,那理论带宽需求就是32KB乘以1GHz,约32TB/s。这个数字片上SRAM通常能扛住一部分,但片外HBM带宽大约也只有几百GB/s到TB/s级别,差了一个数量级。所以实际操作里必须做数据切分、权重驻留、算子融合,把片外访存压到最低,才有可能让有效算力接近峰值算力。
我在一次项目里就吃过这个亏。最初设计拍板用了256x256的大矩阵阵列,觉得算力数字好看。结果软件层算子库做矩阵分块时,发现片上缓存根本放不下一个稍大的特征图分块,不得不在阵列和SRAM之间反复来回搬运数据,实际性能只有峰值的四成不到。后来把阵列改小、缓存加大,性能和能效反而都上去了。
2. 硬件侧的三大件:计算阵列、片上内存、数据通路
硬件设计往细了说是一堆模块,但架构上真正决定AI芯片性能上限的就三样东西:计算阵列怎么做、片上内存怎么分层、数据通路怎么组织。这三块的设计决策直接决定了软件侧能写出什么样的优化代码,反过来软件侧的调度需求又会约束硬件实现。我在这里展开讲一下我在设计过程中的取舍思路。
2.1 计算阵列:走脉动还是走SIMD
计算阵列的主流方案大致分两类:一类是脉动阵列(Systolic Array),代表是大家熟知的TPU风格,权重和激活按固定节奏流过阵列,每个PE只和邻居通信,控制简单、能效高,但对数据流形态要求严格;另一类是SIMD风格的向量阵列,每个PE有独立寄存器堆和控制逻辑,灵活性强,适合不规则算子,但控制复杂度和面积开销更高。
我的选择标准其实很朴素:看目标模型的计算形态是否规整。纯卷积和Transformer里的矩阵乘,脉动阵列非常划算,一条权重流水线就能让阵列满负荷运转;但如果你的目标模型里大量存在动态shape、稀疏、多分支的操作,脉动阵列的固定数据流就会成为限制,这时候SIMD反而更合适。我还做过一种折中方案:主MAC阵列用脉动结构,旁边配一块小的向量单元处理激活、归一化、attention里的softmax这类非规则算子,实际效果很好,芯片面积增加不到10%,但对软件侧的收益是巨大的——算子不用再费劲映射到矩阵阵列上。
2.2 片上内存层次与分配策略
片上内存是AI芯片的生命线。我一般把它分成三层:寄存器堆、局部SRAM(Per-Lane/PE Private)、共享SRAM(Global Buffer/Shared Memory)。寄存器堆在每个PE里面,速度最快、容量最小;局部SRAM给每个PE或每个计算簇私有,存算子级中间结果;共享SRAM是整个芯片的核心数据中转站,所有PE都从它读数据,设计容量通常是几百KB到几十MB。
共享SRAM的设计容量和bank数量是硬件设计阶段争议最多的点。容量太大,面积和漏电功耗压不住;容量太小,算子分块算法施展不开。我的经验是,共享SRAM容量至少要能放下一个典型卷积层的输入通道分块加权重分块加输出分块。举个例子,一个3x3卷积,输入是64通道的56x56特征图,输出也是64通道,在FP16下,输入特征图分块大约384KB,权重大约144KB,输出大约384KB,加起来接近912KB。所以一个面向中高端场景的AI芯片,共享SRAM设计在1MB到2MB才算比较宽裕,再配合bank交织实现多PE同时无冲突访问。
2.3 数据通路和数据搬移的工程权衡
数据通路是让数据按时按点到达PE的那套互联逻辑。片上互联简单做就是总线,高性能做通常用crossbar或NoC(片上网络)。AI芯片里最怕的是"数据到了但PE在等""PE算完了但结果写不回去"这种流水线气泡,所以数据通路的设计目标是保证确定性低延迟。
我在一个多核AI芯片项目里做过一个决策:核间通信用硬件信箱(Mailbox)加共享SRAM的简单方案,而不是上复杂的缓存一致性协议。原因是AI芯片的核间通信模式是"生产者-消费者"高度确定的,每个核按固定节奏产出中间结果,如果走完整的缓存一致性,硬件复杂度和验证成本会大幅上升,但实际收益很小。这算是一个从软件调度角度反推硬件设计的例子——软件层可以把同步点排得很规整,硬件就不需要那么通用的机制。
3. 软件侧的三横层:指令集、编译器与算子库
AI芯片软件栈通常分三横层:最底下是指令集和运行时,中间是编译器,最上面是算子库和框架适配层。很多硬件团队容易忽略软件栈的重要性,结果芯片流片回来发现没有好用的编程入口,算法工程师根本用不起来。我在项目里一般把软件工作量看得和硬件一样重,提前就并行开发。
3.1 指令集设计要兼顾硬件简洁和软件表达力
指令集是软硬件交互的接口。AI芯片的指令集和CPU差异很大,不需要通用寄存器那套复杂语义,而是要贴近数据流硬件:显式表达数据从哪里搬到哪、计算单元执行什么形状的矩阵操作、同步点在哪。我的设计偏好是采用"非阻塞启动、以事件同步"的风格:计算指令发出后不等它完成,继续发下一条搬数指令,由硬件事件队列管理依赖关系。
指令编码的位宽要精打细算。我见过一个项目把指令设计成512位宽的超长指令字,一条指令描述整个矩阵操作的所有参数,看起来表达力很强,但编译器调度特别困难,指令存储体积也占了很多面积。后来改成了多指令流的思路,一条短指令128位,配合硬件状态寄存器逐步展开计算,控制逻辑更简单,调试起来也直观很多。
需要强调的一个实际经验是,指令集最好在软件编译器开工前就冻结,否则两边团队反复改接口定义,时间成本极高。软硬件协同的意思不是同步无限讨论,而是在一个确定的时间点冻结规格,之后只允许微调。
3.2 编译器要做的三件核心优化
AI芯片编译器理论上是把深度学习计算图映射到芯片指令流。听起来和传统编译类似,但实际差异很大。传统面向CPU的编译器是把高级语言翻译成指令序列,AI编译器面对的是一张图,要做的是布局与调度优化。我在内部常说,AI编译器更像一个"给数据安排动线的调度器",而不是"翻译代码的翻译官"。
三个核心优化里,我认为最优先的是算子融合。融合的本质是减少中间结果的片外访存。最典型的是Conv+BN+ReLU融合,或者Attention里QKV投影+缩放+Softmax融合。一次融合能省掉的访存量经常是几倍到十几倍的关系,效果远比单纯调高计算单元利用率显著。第二是数据布局转换,把NCHW转换成NHWC或者分块布局(tiled layout),让访存连续,缓存命中率更高。第三是张量分块循环调度——把大的张量计算切分成适合硬件缓存尺寸的小块,循环层级的顺序直接影响数据复用距离。
3.3 算子库与运行时调度的落地细节
算子库是算法工程师直接接触的那层。一个成熟的AI芯片算子库不仅仅是一堆算子的集合,更关键的是为不同shape的输入选择最优实现的机制。我习惯在每个算子内部实现多个版本,例如针对小batch的轻量实现、针对大输入切分的高吞吐版本、针对特定shape的专用版本,然后通过静态规则加运行时profiling来选择。
运行时的任务调度在AI芯片里容易被忽略。推理任务其实是一个有严格依赖关系的DAG,每个算子的输入输出都在内存的不同区域。运行时调度器要做的不仅是按拓扑序把算子逐个提交,还要负责内存池的复用和细粒度的流水线重叠。我在一个部署项目中就发现运行时对内存的管理对整体性能影响极大,不设内存池时每次请求都重新分配Buffer,峰值内存占用高了20%,延迟也多了接近30%,换成预分配复用池之后,内存占用和延迟都明显降下来。
4. 软硬协同验证:仿真平台、FPGA原型与性能回填
软硬件设计做得再漂亮,最终都要过验证这一关。AI芯片的验证比传统芯片多了一个维度:不只是验证硬件逻辑对不对,还要验证硬件和软件模型之间的一致性判断准确不准。软硬协同验证是AI芯片项目里最容易被低估、最耗时的阶段。
4.1 先做指令级虚拟原型而不是直接RTL仿真
我的经验是,不要在硬件RTL还没稳定的时候就让软件团队干等。先做一个指令级虚拟原型(Instruction Accurate Simulator),软件团队在这个模拟器上做编译器和运行时开发,等RTL冻结后再做周期级仿真对齐。虚拟原型的运行速度比RTL仿真快几个数量级,能让软件栈提前几个月完成主体开发。
这个虚拟原型的精度问题值得说。指令级模型通常不精确到周期,只保证功能和指令语义一致,所以用来开发编译器和算子库是够的。但性能调优、判断某条优化是否真的能提升芯片利用率,就必须依赖周期级仿真。我常常安排一个"性能回填"流程:先用虚拟原型跑出功能正确的结果,再用周期级仿真跑关键模型,把指令数、等待时间、访存冲突这些数据统计出来,回填给编译器团队做调度权重修正。
4.2 FPGA原型和Emulation平台的取舍
中大规模AI芯片一般在流片前会做FPGA原型或硬件仿真(Emulation)。FPGA原型的特点是速度快、能跑比较完整的软件栈,但容量有限,一颗AI芯片的MAC阵列规模往往远超FPGA资源,需要多次综合拆分。Emulation平台容量大、精度高,但速度慢,一天可能只能跑几十帧图。
我的取舍策略是以FPGA跑功能验证和应用开发,以Emulation跑性能验证。FPGA样机上我会跑真实的模型推理,验证整体流程从编译到运行是通的,同时也能做边缘场景测试;Emulation则用来测关键算子的临界时序,比如多核并发访存的冲突行为、权重驻留调度对整体吞吐的影响,这些在FPGA上复现不出来。这个双轨并行方案帮我规避过好几次重大设计失误,比如某次发现核间同步等待时间比预期长三倍,直接在流片前改掉了硬件同步模块设计。
4.3 软硬件联合调试的痛点
芯片调试比普通软件开发难得多,因为寄存器视角和软件视角是割裂的。我们极靠前就设计了片上调试模块,包括轻量级硬件性能计数器、片上trace缓冲、内存读写监控点。软件团队可以通过调试接口读取这些数据,看到某个时刻计算阵列在干什么、访存总线繁忙到什么程度、哪些指令在等待依赖。
拿一个实际调试案例说,某个Transformer模型部署出来性能比预期低了30%,单看软件栈各种常规优化已经做得很足。后来通过硬件trace发现,矩阵单元做QK^T计算时需要连续读取K矩阵的不同行,而K矩阵在内存里的布局和硬件bank的映射发生了巨大冲突,导致访问延迟被拉高。找到根源后就改了一个数据布局转换函数,性能立刻恢复到预期。这种问题没有硬件可观测能力,软件团队排查一周都不一定想到。
5. 性能调优的关键手段:从访存优化到量化部署
硬件和软件全部就绪后,真正决定一款AI芯片口碑的是在真实模型上的性能表现。我在好几个项目里都经历过"指标达标但实际模型跑不快"的尴尬。这里分享几条我在性能调优阶段一定优先做的事情。
5.1 内存搬移尽量隐藏在计算背后
AI芯片优化目标很简单:让计算单元尽量不闲着。但实际操作很难,因为计算依赖的数据必须准时到位。常规手段是做双缓冲或多缓冲流水线。比如矩阵乘过程里,计算单元算第N块的同时,DMA搬第N+1块进片上SRAM,等计算单元完成第N块,第N+1块数据已经就位,流水线无气泡。
双缓冲的真正难点在于切分粒度的控制。切得太细,流水级数多、指令调度复杂,同步开销增大;切得太粗,片上缓存放不下足够的缓冲块,下游又在等数据。我一般用"目标延迟毫秒级、片上SRAM容量、单算子总访存量"这三个约束联立解出一个切分粒度,第一版粗调后通过仿真跑出数据再细调。这个环节极其依赖软硬件同事坐在一起做联合调参,两边各退一步往往结果最好。
5.2 量化部署中的硬件配合与软件校准
AI芯片通常支持INT8/FP16等多种精度。量化部署时,硬件与软件的配合很关键。硬件给出SIMD整数计算单元、整数累加器、转FP的精度转换模块;软件层负责把权重从FP32校准到INT8,并把激活值的动态范围统计准。任何一个环节出错,精度就会崩。
我的实际经验是,量化对AI芯片的加速收益不是线性的。INT8的峰值算力理论上是FP16的两倍,但实际跑起来受制于数据搬运,瓶颈往往不在计算而在带宽,所以有时你把精度降到INT8,有效吞吐提升可能只有1.3倍,而不是2倍。这时候真正带来收益的是"能否省一半的记忆体带宽"。如果芯片在INT8模式下访存位宽减半,那加速比才会更接近理论值。这个细微差别很多人容易忽略,导致硬件设计时没有优化INT8模式下的数据通路。
5.3 多核并行与算子切分的配合
大规模AI芯片基本都是多核架构,把一张大模型图切分到多个核上执行是性能提升的另一条主线。我的做法是"数据并行+算子并行"混合切分:对计算量大的卷积和矩阵乘,按通道或按序列维度拆分到多核;对依赖关系强的前后算子,用流水线方式排布到不同核上。
这个切分方案必须让编译器能够感知核间通信的开销。一次跨核同步的代价如果比单核计算时间还长,那切分就得不偿失。我一般会在编译器里把跨核通信建模成一种"虚拟算子",给它一个成本模型,调度器在选择切分方案时把通信成本纳入优化目标。通过这种机制,软件团队不用记住每个硬件的通信延迟数字,编译器会自动平衡计算与通信。实测之后,性能提升稳定跑在1.8倍以上,而且切分方案可解释、可回放。
6. 一次性能回填后的架构修正:一个复盘样本
为了把前面讲的软硬件协同设计讲透彻,我复盘一个真实经历过的项目。某款AI芯片目标跑计算视觉多模型,初期设计时团队拍板让四核共享一个大的SRAM池,编译器负责动态分割,这样软件写起来省事。结果做FPGA原型时跑真实业务,发现多核并发访问共享SRAM时bank冲突特别严重,实际带宽只有设计峰值的55%,压缩了整体吞吐接近四成。
排查过程是,先从软件侧做数据布局对齐,效果有提升但有限。又把SRAM从一个大池改成两两核绑定一个较小的SRAM组,同时保留跨组访问路径但降低其优先级。这个改动让共享冲突明显减少,编译器策略也简单了,核内数据局部性天然变好。最终整体性能从接近预测值的60%提升到了85%以上,验证了"给软件减负就是给硬件提效"这句话在AI芯片里的价值。
这个案例的教训被我固化成了设计规范:第一,缓存和SRAM尽量按计算簇划分,少做全局大池;第二,跨簇共享路径纯粹用于边界数据交换,不要让它承担主线数据搬运;第三,任何内存架构设计都要有配套的编译器访存模型,让软件能在设计阶段就知道访问Pattern的好坏,而不是等流片后才发现。
我在实际项目中反复体会到的一句话:AI芯片是典型的系统级工程,硬件设计必须把软件调度的现实约束当作输入,软件设计也必须把硬件实现的物理限制当作约束,两边站在同一条流水线上才能做出真正跑得快的产品。性能问题如果不在设计阶段就开始闭环,后期靠软件优化弥补的成本会非常离谱。
最后分享一个小细节:设计评审的时候,我习惯让软件团队在硬件架构文档上标出"这段逻辑会让编译器很难受"的注释,同样硬件团队也在软件设计文档上标出"这个优化依赖某种硬件特性但硬件实现代价高"的标记。这个看似简单的双向标注机制,后续避免了很多无谓的方案冲突,也帮我更准确地判断每次设计改动的真实收益。