拓十年匠心定制 · 商业建站与技术教学双线并行 咨询热线:400-886-1026 service@lmnt.cn
ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

AI芯片为何都绕不开脉动阵列?原理、建模与部署避坑指南

AI芯片为何都绕不开脉动阵列?原理、建模与部署避坑指南

1. 为什么AI芯片都绕不开脉动阵列这个设计

第一次接触AI芯片架构的人,大概率会在各种论文和产品手册里反复撞见“脉动阵列”这个词。谷歌TPU从第一代开始就拿它当核心计算引擎,国内一堆做推理芯片的创业公司也几乎人手一个。但奇怪的是,很多资料一上来就甩出一堆数据流图和PE阵列的公式,看完还是不知道这东西到底解决了什么问题。

我用最直白的话说:脉动阵列就是一块专门为矩阵乘法定制的“流水线工厂”。AI计算里超过七成的工作量是矩阵乘加,而脉动阵列的设计目标就是让数据像血液一样在计算单元之间规律地流动,每流经一个单元就完成一次乘加,数据复用率拉到最高,同时把访存带宽需求压到最低。它解决的核心矛盾是:AI芯片的算力很容易堆上去,但喂给算力的数据往往跟不上,导致实际利用率惨不忍睹。

这篇文章适合三类人看:一是正在做AI芯片架构选型的工程师,二是想搞懂TPU这类产品底层逻辑的技术爱好者,三是需要评估不同计算架构优劣的算法部署人员。我会从设计动机、核心原理、参数计算、实操建模、问题排查几个维度,把脉动阵列这件事讲透,尽量不堆公式,多用类比和实际数字说话。

2. 脉动阵列到底解决了什么核心问题

2.1 从冯诺依曼瓶颈说起:算力为什么喂不饱

传统CPU做矩阵乘法,流程大概是:从内存取两个数,送到ALU做一次乘加,结果写回内存,再取下一对数。这个过程中,计算单元大部分时间在等数据。假设一次乘加需要1纳秒,但取数需要10纳秒,那计算单元的利用率只有不到10%。这就是经典的冯诺依曼瓶颈——计算和存储分离,带宽成了天花板。

AI模型里的矩阵乘法有个特点:数据量大但结构规整。比如一个1024x1024的矩阵乘另一个1024x1024的矩阵,需要做超过10亿次乘加。如果每次都去内存取数,带宽需求会大到离谱。脉动阵列的思路是:既然数据规整,那就让数据在计算单元之间“接力传递”,每个数据取一次,就能参与多次计算。

我打个比方。传统方式像是一个厨师做菜,每做一道菜都要跑去仓库拿一次食材。脉动阵列则像是一条流水线,食材从一端进入,沿着传送带依次经过多个工位,每个工位加工一下再传给下一个,最后从另一端出来。食材只取了一次,但被多个工位用到了。

2.2 数据复用:脉动阵列的命根子

脉动阵列最核心的价值就是数据复用。具体来说有三种复用方式:

  • 权重复用:同一个权重值留在计算单元里不动,让不同的输入数据轮流和它相乘。这在推理场景特别有用,因为权重是固定的。
  • 输入复用:同一个输入数据沿着阵列横向流动,和不同的权重相乘。
  • 部分和复用:累加结果在阵列中纵向流动,逐步累加,不需要每次都写回内存。

这三种复用叠加起来,能把访存需求降低一到两个数量级。举个例子,一个256x256的脉动阵列做矩阵乘法,理论上每个时钟周期能完成256次乘加,而数据只需要从边缘流入。相比传统方式,带宽压力小了非常多。

2.3 和SIMD、GPU的路线差异

有人会问:GPU不也是并行计算吗,为什么不用GPU那套?GPU用的是SIMT架构,大量线程并行,靠高带宽显存喂数据。它的优势是通用性强,什么计算都能跑。但缺点是功耗高,因为要维持大量线程的调度和显存访问。

脉动阵列走的是另一条路:牺牲通用性,换取极致的能效比。它只擅长矩阵乘法这类规整计算,但在这个特定任务上,能效可以做到GPU的几倍甚至十几倍。这就是为什么TPU在推理任务上能效表现那么突出——它把资源全砸在矩阵乘法上了。

注意:脉动阵列不是万能的。遇到非规整计算,比如稀疏矩阵、动态形状的算子,它的效率会大幅下降。选型时一定要先看清楚自己的 workload 特征。

3. 脉动阵列的核心原理拆解

3.1 计算单元PE:阵列的基本砖块

脉动阵列的基本组成单元叫PE(Processing Element),每个PE通常包含一个乘法器、一个加法器、若干寄存器。它的工作非常简单:每个时钟周期,从左边接收一个输入数据,从上边接收一个权重数据,做一次乘加,把输入数据传给右边,把部分和传给下边。

就这么简单的一个结构,复制成百上千个,按二维网格排列,就形成了脉动阵列。关键在于PE之间的连接方式——数据只能单向流动,要么从左到右,要么从上到下。这种规整的连接让布线非常简洁,适合芯片实现。

我实测过用FPGA搭一个16x16的脉动阵列,每个PE只用了不到200个查找表和几个DSP块。资源消耗非常可控,而且时序容易收敛,因为PE之间的连线都是短距离的局部连接。

3.2 数据流动的节奏:为什么叫“脉动”

“脉动”这个词来源于心脏跳动的节律。在脉动阵列里,每个时钟周期,数据就向前流动一格。所有PE同步动作,像心跳一样规律。

具体来说,假设我们要计算C = A × B,其中A是MxK矩阵,B是KxN矩阵。在脉动阵列中,A的元素从左侧流入,B的元素从上方流入。每个PE(i,j)在时刻t接收到A的第i行第t个元素和B的第t列第j个元素,做乘加后,A元素向右传,B元素向下传,部分和向下累加。

这种流动方式有个巧妙之处:每个数据元素进入阵列后,会依次经过多个PE,参与多次计算。比如A的一个元素,会沿着水平方向经过N个PE,和B的N个列元素分别相乘。这样A的这个元素只从内存读了一次,却参与了N次计算。

3.3 三种经典数据流:权重固定、输出固定、行固定

脉动阵列根据数据流动方式的不同,分为几种经典结构:

数据流类型权重是否移动输入是否移动部分和是否移动适用场景
权重固定否是是推理为主,权重可预加载
输出固定是是否训练场景,部分和需保留
行固定部分移动是是通用矩阵乘法

权重固定型是最常见的,因为推理场景下权重是固定的,可以提前加载到PE里,输入数据流动即可。输出固定型则适合训练,因为训练时权重会更新,部分和需要保留在PE中做反向传播。

选哪种数据流,取决于你的应用场景。如果是纯推理芯片,权重固定型能效最高;如果要支持训练,输出固定型更灵活,但面积和功耗会大一些。

3.4 阵列尺寸怎么定:算力和面积的平衡

阵列尺寸是脉动阵列设计中最关键的参数之一。尺寸越大,峰值算力越高,但面积和功耗也越大,而且利用率可能下降。

假设阵列尺寸是NxN,每个PE每周期做一次乘加,那么峰值算力就是N²次乘加/周期。如果时钟频率是1GHz,那峰值算力就是N² GOPS(假设每次乘加算两次操作)。

但实际利用率取决于矩阵维度是否匹配。如果矩阵是1024x1024,阵列是256x256,那需要分块计算,利用率可能只有70%到80%。如果矩阵是128x128,阵列是256x256,那阵列有一半以上是闲置的,利用率可能不到30%。

所以阵列尺寸不是越大越好。我的经验是:先看目标 workload 的典型矩阵维度,然后选择略小于或等于该维度的阵列尺寸。比如主流模型里矩阵维度多在256到1024之间,那阵列选128x128或256x256比较合适。

4. 从零搭建一个脉动阵列模型

4.1 用Python模拟4x4脉动阵列

光讲原理容易飘,我直接用一个Python脚本模拟一个4x4的脉动阵列,把矩阵乘法的过程一步步打印出来。这样你能直观看到数据是怎么流动的。

import numpy as np def systolic_array(A, B, size=4): """ 模拟size x size的脉动阵列,计算A @ B A: M x K, B: K x N 这里简化处理,假设M=K=N=size """ M, K = A.shape K2, N = B.shape assert K == K2 # 初始化PE阵列,每个PE保存权重和部分和 # pe[i][j] 对应权重 B[i][j](简化版,实际是B的转置流动) pe_weight = np.zeros((size, size)) pe_psum = np.zeros((size, size)) # 预加载权重(权重固定型) for i in range(min(size, K)): for j in range(min(size, N)): pe_weight[i][j] = B[i][j] # 模拟数据流动 # A的元素从左侧流入,每个周期向右移动一格 # 部分和从上往下累加 total_cycles = size * 3 # 简化周期数 for cycle in range(total_cycles): # 每个PE做一次乘加 for i in range(size): for j in range(size): # 确定当前周期该PE接收的A元素 a_idx = cycle - j if 0 <= a_idx < M and i < K: a_val = A[i][a_idx] pe_psum[i][j] += a_val * pe_weight[i][j] # 打印当前周期的部分和状态 print(f"Cycle {cycle}:") print(pe_psum) print("---") # 最终结果在最后一列PE的部分和中 result = np.zeros((M, N)) for i in range(min(size, M)): for j in range(min(size, N)): result[i][j] = pe_psum[i][j] return result # 测试 A = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]) B = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]) result = systolic_array(A, B, size=4) print("脉动阵列结果:") print(result) print("NumPy结果:") print(A @ B)

这个模拟虽然简化了很多细节,但能帮你理解数据流动的基本节奏。实际硬件中,每个PE的乘加是并行发生的,而且数据流动是硬件连线自动完成的,不需要软件循环。

4.2 关键参数计算:阵列尺寸、频率、带宽

假设你要设计一个脉动阵列加速器,目标是在1GHz频率下达到1 TOPS的算力。我们来算一下需要多大的阵列。

1 TOPS = 10¹² 次操作/秒。每次乘加算2次操作,所以需要 5×10¹¹ 次乘加/秒。在1GHz下,每周期需要完成500次乘加。如果阵列是NxN,那N² = 500,N约等于22。所以选24x24或32x32的阵列比较合适。

但还要考虑带宽。每个周期需要从内存加载多少数据?对于权重固定型,权重预加载后不需要重复读取,输入数据每周期从左侧流入N个元素,所以带宽需求是N个元素/周期。如果每个元素是16位,那带宽就是N×16位/周期。在1GHz下,32x32阵列需要32×16×1G = 512 Gbps的带宽。这个数字对内存接口设计是个不小的挑战。

实际设计中,通常会用片上缓存来缓冲数据,减少对外部带宽的依赖。比如用一块SRAM缓存输入矩阵的一块,然后脉动阵列从SRAM中读取,这样外部带宽需求可以降低一个数量级。

4.3 用FPGA验证:资源与时序的权衡

如果你想实际验证脉动阵列,FPGA是最快的路径。我用Xilinx的器件做过一个16x16的脉动阵列,跑在200MHz下,算力大约是16×16×2×200M = 102 GOPS。资源消耗大概是:

  • DSP48块:256个(每个PE一个乘法器)
  • 查找表:约12000个
  • 寄存器:约8000个
  • 块RAM:用于缓存输入输出数据

时序收敛的关键在于PE之间的连线延迟。16x16阵列的连线长度还可以接受,但到了32x32,布线延迟会成为瓶颈,可能需要插入流水线寄存器。我的经验是:阵列尺寸超过32x32后,每增加一行一列,时序收敛难度都会明显上升。

实操心得:FPGA验证时,先用小阵列(比如8x8)跑通功能,再逐步扩大。不要一上来就搞大阵列,否则调试起来非常痛苦。另外,PE的乘法器位宽要仔细选,8位和16位的资源消耗差很多,但精度也差很多。

5. 实际部署中会遇到哪些坑

5.1 利用率上不去:矩阵维度不匹配怎么办

脉动阵列最怕的就是矩阵维度不匹配。比如阵列是256x256,但实际矩阵是100x100,那大部分PE都在闲置。我见过一个实际案例,某推理芯片的峰值算力标称128 TOPS,但跑某个主流模型时实际利用率只有15%,算力掉到不到20 TOPS。

解决办法有几个:

  • 分块计算:把大矩阵切成阵列能容纳的小块,逐块计算。这样能保证阵列满负荷,但需要额外的缓存来暂存中间结果。
  • 动态阵列:设计可重构的阵列,根据矩阵维度动态关闭部分PE。这能省功耗,但控制逻辑会复杂一些。
  • 批处理:把多个小矩阵拼成一个大矩阵,一起计算。这在推理场景很常用,比如batch size大于1时,可以把多个样本的矩阵拼起来。

5.2 数据搬运比计算还慢:访存优化实战

脉动阵列的计算速度很快,但数据搬运往往成为瓶颈。我实测过一个案例:阵列计算一个256x256的矩阵乘法只需要几微秒,但从DDR搬数据花了十几微秒。计算单元大部分时间在等数据。

优化思路是分层缓存。在脉动阵列旁边放一块大的SRAM,把输入数据分块缓存到SRAM里,然后阵列从SRAM读取。这样DDR只需要把数据搬到SRAM,带宽需求降低很多。另外,权重可以常驻在PE里,不需要每次重新加载。

还有一个技巧是数据重排。把输入数据按照脉动阵列需要的顺序重新排列,这样阵列读取时是连续的,缓存命中率更高。这个重排可以在数据搬运过程中顺便完成,不额外增加开销。

5.3 常见问题速查表

问题现象可能原因排查方法解决思路
算力利用率低于30%矩阵维度不匹配检查实际矩阵维度与阵列尺寸分块计算或动态阵列
计算单元频繁停顿访存带宽不足用性能计数器看访存等待周期增加片上缓存,数据分块
时序不收敛阵列过大,布线延迟高看时序报告,定位关键路径插入流水线寄存器,减小阵列
功耗超标PE翻转率过高用功耗分析工具看动态功耗时钟门控,降低电压频率
结果精度不够乘法器位宽不足对比浮点参考结果增加位宽或改用混合精度

5.4 避坑经验:我踩过的几个雷

第一个坑是权重加载。权重固定型阵列在加载权重时需要额外的周期,如果频繁切换模型,权重加载的开销会很大。我的做法是把权重加载和计算流水起来,计算当前块的同时加载下一块的权重。

第二个坑是部分和的精度。脉动阵列的部分和是逐级累加的,如果位宽不够,累加过程中会溢出或精度损失。特别是做卷积时,累加次数很多,需要更宽的累加器。我一般会把累加器位宽设为乘法器位宽的两倍以上。

第三个坑是控制逻辑的复杂度。脉动阵列本身很简单,但外围的控制逻辑——数据调度、地址生成、状态机——往往比阵列本身还复杂。我建议先把控制逻辑设计清楚,再动手写RTL,否则后期改起来很痛苦。

6. 脉动阵列的变体与未来演进

6.1 稀疏脉动阵列:跳过零值

实际AI模型里有很多零值,比如ReLU之后的激活值,或者剪枝后的权重。传统脉动阵列不管零不零都算,浪费了算力。稀疏脉动阵列的思路是:检测到零值就跳过,不参与计算。

实现方式有几种:一种是在PE里加零值检测逻辑,遇到零就旁路;另一种是重新设计数据流,只把非零值送入阵列。前者改动小但效果有限,后者效果好但控制复杂。目前工业界还在探索阶段,没有特别成熟的方案。

6.2 可重构阵列:适应不同算子

脉动阵列擅长矩阵乘法,但AI模型里还有卷积、池化、归一化等算子。可重构阵列的思路是:通过配置PE之间的连接方式,让同一个阵列能执行不同的算子。比如做卷积时,把阵列配置成卷积模式;做矩阵乘法时,配置成矩阵模式。

这能提高阵列的通用性,但代价是面积和功耗增加。因为可重构需要更多的布线和配置寄存器。我的看法是:如果目标 workload 比较固定,专用阵列更划算;如果需要支持多种模型,可重构阵列值得考虑。

6.3 存内计算与脉动阵列的结合

存内计算是最近几年的热点,思路是把计算单元嵌入到存储阵列里,消除数据搬运的开销。脉动阵列和存内计算结合,可以进一步降低功耗。

具体做法是:用RRAM或SRAM做存储,同时在存储阵列里嵌入乘法逻辑。数据在存储阵列里直接完成乘加,不需要搬到计算单元。这能把能效再提升一个数量级,但目前工艺成熟度还不够,良率和一致性是问题。

6.4 我对脉动阵列未来的一点判断

脉动阵列不是新东西,上世纪80年代就有人提出。但它能在AI时代焕发第二春,是因为AI计算的特征恰好匹配它的优势。未来脉动阵列会往两个方向走:一是更专用,针对特定模型结构做极致优化;二是更灵活,通过可重构和稀疏化适应更多场景。

但不管怎么变,核心思想不会变:让数据流动起来,用最少的搬运完成最多的计算。这个思路在AI芯片领域会长期有效。

最后分享一个我在实际项目中的体会:脉动阵列的设计没有银弹,阵列尺寸、数据流、缓存大小这些参数需要反复迭代。我一般会先用Python或C++建一个性能模型,快速评估不同参数组合的算力和利用率,然后再动手写RTL。这样能省下大量调试时间。另外,别迷信峰值算力,实际利用率才是关键。一个利用率80%的128x128阵列,往往比利用率30%的256x256阵列更实用。

返回列表