TPU：脉动阵列

可视化演示

展示系统处理流水线、实时性能指标曲线及多方案对比分析，包含数据流转、吞吐延迟监控与方案评估。

源于访存瓶颈

旧工厂里最刺眼的画面，是工人们来回奔跑的腿。拧一颗螺丝需要五秒，跑去仓库领料却要三十秒。我们把这种行为叫“冯诺依曼瓶颈”：在传统计算机里，计算单元（工人）和内存（仓库）是分离的，每算一步都要先跨过一条狭窄的通道把数据取回来。计算本身再快，也被这条来回取数的通道死死卡住——工人空有一双快手，却等不来物料。

AI 计算放大了这个矛盾。神经网络里的矩阵乘法动辄涉及百万千万个数据，可内存带宽的增长速度远远赶不上芯片算力的增长速度。结果就是：芯片的计算单元绝大部分时间在“空转等数据”，我们称之为“内存墙”。因此，聪明的芯片设计师不再急着把单个工人训练得更快，而是换了一个思路——为什么不让数据自己走起来，主动流到工人们手边呢？

这个思路的转变，正是脉动阵列的起点。它把“以计算为中心”变成“以数据流动为中心”，从根子上绕开了冯诺依曼瓶颈。从此，工人不必再奔跑，传送带替他们跑完所有的路。

阵列拓扑结构

厂房改造的第一件事，是重新规划工位布局。旧工厂的工位围着仓库转，新工厂把工人排成一个横平竖直的方阵：你站在任何工位上，只需伸手够得着左右和上下四个邻居——物料从左边工位传来，做完自己的活，又递给右边的工位。这就是脉动阵列的拓扑结构：二维 PE 网格。

每个工位有一个名字，叫 PE（Processing Element，处理单元）。它很“笨”，里里外外只装了一个乘法器和一个加法器，负责把两数相乘再加到已有结果上。但正因为它足够简单，才能在芯片上铺得密密麻麻。Google 第一代 TPU 里，这样的 PE 排成了一个 256×256 的大方阵，共 65536 个，每个都只与相邻单元通过极短的“局部总线”相连，全阵列没有任何一条贯穿全局的昂贵总线。

局部互连意味着什么？好比工位之间是几十厘米的小传送带，而不是横跨整座厂房的天桥。数据从一个 PE 传到隔壁 PE，几乎不花什么代价；而每一条长长的全局总线，都意味着高能耗、高延迟、占面积。阵列拓扑的智慧，就是把长距离通信消灭在摇篮里——所有搬运都发生在邻居之间，一步传递，成本极低。

脉动节拍机制

新厂房上线前，你给每个工位装了一只统一的节拍器。铃响一声，所有工人同时做一个动作；再响一声，再做一个动作。没有人抢跑，也没有人拖延——这就是“脉动”二字的来历：像心脏收缩一样，整个阵列在同一节拍下同步搏动，把数据一波一波地泵向远方。

一个节拍内，每个 PE 只做一件事：把刚传到手边的激活值 x，乘上自己工位上驻留的权重 w，再加到从上方传来的部分和 p 上，然后把新的部分和交给下方的邻居。用数学语言表达，就是：

节拍器一响，全厂上下同时呼吸。你可能担心这么多工位要同步，时钟信号怎么传得过来？办法是把时钟像水波一样从中心向外扩散，让所有 PE 在同一微秒内踏出同一步。精准的同步，保证了数据在单元间的接力不会乱——前一个工位刚松手，后一个工位刚好接住。

数据流模式

生产线上有三样东西可以动：零件盒、半成品、成品。聪明的厂长发现，不同的产品适合不同的物料安排。第一条线把零件盒固定在各工位，半成品沿产线流动；第二条线让半成品一直待在流水线上，工人来回调整；第三条线让成品在生产过程中逐级传递、边做边组装。脉动阵列的数据流，也分为这样三种经典模式，分别叫权重驻留、输入驻留、输出驻留。

TPU 采用的是第一种——权重驻留：把神经网络训练好的权重矩阵像零件盒一样，提前“烧录”进每个 PE 的寄存器里，在整个计算过程中纹丝不动。需要移动的只有两样东西：激活值从阵列左侧注入，沿着每一行向右边的邻居逐拍传递；部分和从阵列顶部（通常是零）注入，沿着每一列向下累加。

为什么 TPU 偏爱权重驻留？在推理场景里，权重是固定不变的一道“模具”，而输入数据是源源不断流过的“毛坯”。模具安在工位上不动，毛坯不停地流过，每条产线就能持续不断地吐出成品。这种模式让数据流动的路径最短、规律最强，也最适合硬件资源有限的芯片。三种模式没有绝对优劣，但权重驻留恰好踩中了 AI 推理的节拍。

边界处理策略

流水线有起点，也有终点。厂房左侧设了上料斜坡，毛坯从那里送上传送带；右侧是出货口，成品在那里打包。脉动阵列同样需要精心打理自己的边界——处理不好，就会像传送带两端的物件堆积、掉落，乱成一团。

阵列的边界有三道工序。第一道是注入：每个时钟周期，要算的新激活值向量从左侧逐列推进，就像把毛坯一个接一个放上传送带；阵列顶部则不断注入 0，作为部分和的初始起点。第二道是零填充：当矩阵尺寸比阵列小、或计算刚开始、快要结束时，某些输入位置是空的，就用 0 把传送带上的空位填满，保证节奏不断拍。第三道是收集：底部每拍吐出一行部分和，这些数据还要经过重排缓冲器，把流水线上“错位”的输出归位，才能得到规规矩矩的结果矩阵。

边界处理还很考验耐心——一条流水线刚开动时，前几拍传送带上稀稀拉拉，产不出完整成品；快要停工前，也必须多跑几拍“空拍”，把线上遗留的半成品清空。这就是流水线固有的“填充期”和“排空期”，硬件工程师必须在阵列旁边备好缓冲，接住这些先乱后齐的输出。

访存优化收益

改造成流水线后，厂长算了一笔账：工人再也不用跑仓库了，取料的次数从“每个零件跑一趟”变成了“整批材料搬一次”。脉动阵列的收益，同样可以精确地算出来。

在一个 N 列的脉动阵列里，每个从左侧进入的激活值，会沿着它所在的行一路向右，依次经过 N 个 PE。它每经过一个 PE，就被这个 PE 读取并乘上一次。也就是说，同一个数据在片上被复用了 N 次，而这 N 次复用只占用了 N 次片上寄存器传递，完全不需要多次访问主存：

在传统架构里，每一次乘法都要先到主存取两个数；而在脉动阵列中，主存只在最左侧“投喂”数据流、在底部“回收”结果，中间的亿万次乘加全部靠片上流动完成。访存带宽的悬崖，就这样被数据复用这座桥轻轻跨了过去。传送带的一小步，是省下带宽的一大步。

矩阵乘法映射

最后一步，是把数学上的矩阵乘法，翻译成流水线的生产计划。假设要计算输出向量 \(\mathbf{y} = \mathbf{W}\mathbf{x}\)，其中 \(\mathbf{W}\) 是一个 K×K 的权重矩阵，\(\mathbf{x}\) 是输入向量。厂长一声令下，整条生产线按三个步骤运转起来。

第一步，烧录。把权重矩阵 \(\mathbf{W}\) 的第 \(i\) 行第 \(j\) 列元素 \(w_{ij}\)，预先写进阵列第 \(i\) 行、第 \(j\) 列的 PE 里。权重从此安家落户，不再移动。第二步，注入。把输入向量 \(\mathbf{x}\) 的分量们，按顺序从左侧送入每一行，让它们像一排士兵齐步走，沿行向右推进。第三步，收集。每个 PE 每拍把 \(w_{ij}\) 乘以流到眼前的 \(x\)，累加到从上方传下的部分和中；当数据走完整个阵列，底部吐出的就是完整的输出：

就这样，一条 256 列的生产线，每拍能同时完成 65536 个乘加运算。矩阵乘法没有惊天动地的魔法，有的只是把千万次独立的乘加，拆解成整齐划一的流水节拍，交给每一滴“血液”去完成自己的使命。数据流动之处，算力自然涌现。

AI 应用场景

Python 演示

class SystolicArray2D:    """权重驻留的二维脉动阵列：每个 PE 一拍做一次乘加。"""     def __init__(self, rows, cols):        self.rows = rows        self.cols = cols        self.w = [[0] * cols for _ in range(rows)]        self.psum = [[0] * cols for _ in range(rows)]        self.x = [[0] * cols for _ in range(rows)]     def load_weights(self, w):        self.w = [row[:] for row in w]     def clock(self, x_in):        """一个脉动拍：先算部分和，再右移输入。"""        new_psum = [[0] * self.cols for _ in range(self.rows)]        for i in range(self.rows):            for j in range(self.cols):                up = 0 if i == 0 else self.psum[i - 1][j]                new_psum[i][j] = up + self.w[i][j] * self.x[i][j]        self.psum = new_psum        for i in range(self.rows):            for j in range(self.cols - 1, 0, -1):                self.x[i][j] = self.x[i][j - 1]            self.x[i][0] = x_in[i]     def run(self, x_seqs):        """逐拍注入输入向量，返回每拍底部输出快照。"""        out = []        for t, xv in enumerate(x_seqs):            self.clock(xv)            out.append([self.psum[-1][j] for j in range(self.cols)])        return out  def reference(w, x):    """参考实现：第 j 个输出 = 权重矩阵第 j 列与输入向量的点积。"""    return [sum(w[i][j] * x[i] for i in range(len(w))) for j in range(len(w[0]))]  def main():    w = [        [1, 0, 1, 2],        [2, 1, 0, 1],        [1, 3, 1, 0],        [0, 2, 2, 1],    ]    x_seqs = [        [1, 0, 2, 1],        [0, 1, 1, 0],        [2, 1, 0, 2],        [1, 2, 1, 1],    ]    sa = SystolicArray2D(4, 4)    sa.load_weights(w)    history = sa.run(x_seqs)     print("=== 逐拍底部输出（流水线效果） ===")    for t, row in enumerate(history):        print(f"第 {t + 1} 拍: {row}")     print("\n=== 验证时序关系 ===")    for s in range(len(history)):        for j in range(4):            expect = 0            for i in range(4):                idx = s - 1 - i - j                if 0 <= idx < len(x_seqs):                    expect += w[i][j] * x_seqs[idx][i]            got = history[s][j]            assert abs(got - expect) < 1e-9, (s, j, got, expect)    print("所有拍与参考计算一致！")  if __name__ == "__main__":    main()

1. class SystolicArray2D:定义一个二维脉动阵列类，结构上等同厂房里的方阵工位。
2. self.w = [[0] * cols for _ in range(rows)]给每个 PE 准备一个权重寄存器，初始化全 0。
3. self.psum部分和数组，记录每个 PE 当前的累积结果，对应“半成品”。
4. self.x记录流经每个 PE 的激活值，对应“传送带上的毛坯”。
5. up = 0 if i == 0 else self.psum[i - 1][j]取上方传下的部分和；最顶行从 0 起步。
6. new_psum[i][j] = up + self.w[i][j] * self.x[i][j]核心公式 p_out = p_in + w×x，一拍一乘加。
7. for i in ... for j in range(self.cols - 1, 0, -1):从右往左倒序右移数据，避免覆盖未使用的旧值。
8. self.x[i][0] = x_in[i]把新的输入向量注入阵列最左列，好比毛坯送上产线。
9. out.append([self.psum[-1][j] ...])每拍采集底部一行结果，模拟出货口打包。
10. for s in range(len(history)):用推导出的时序公式逐项核验，保证模拟器没有“偷算”。

常见误区

面试题