MXU矩阵单元：TPU脉动阵列的工作原理

可视化演示

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

脉动阵列概述：给矩阵乘法装上流水线

回到那家汉堡流水线店。老板发现单靠一个大厨炒一锅菜太慢，于是把做汉堡拆成：放面包、加肉饼、加酱汁、加蔬菜、打包。每个工位只做一步，所有工位同时开工，一个面包胚从第一个工位走到最后一个工位，后厨就能批量出货。MXU之所以叫“矩阵乘法单元”，就是因为它把数学里的矩阵乘法拆成了无数个“乘加小步”，让数据自己流经工位，而不是让一个计算核心反复取数。

矩阵乘法是神经网络最基本的运算——全连接层、卷积层、注意力权重计算，本质上都是二维矩阵相乘。MXU把这种运算映射到二维脉动阵列上：所有PE按网格排列，每个PE负责一个输出元素的一部分累加工作。它不像CPU那样有复杂控制逻辑，而是用“数据驱动计算”的方式，让每个PE在数据到达时自动运算。

脉动阵列最大的优势是数据复用和低数据搬运。权重先加载进阵列并锁存，激活值像传送带上的面包胚一样流经每个PE，每个PE只需要在本地取权重、乘上流入的激活值、累加部分和，然后传给下一个PE。整个过程中不需要频繁访问寄存器堆或主存，能耗和延迟大幅降低。

就像一个汉堡包的“总成本”由所有工位累加的成本组成，矩阵乘法的每个输出也由所有中间乘积累加而来。理解了这一步，后面所有流水线细节都顺理成章。

数据流入流出：左侧进配方，顶部进食材

流水线后厨的规则很精确：配料架上的“配方卡”必须在开工前挂到每个工位上方，面包胚则从顶部传送带一个个滑下来。工位师傅不会主动去取料，只等食材自己送到面前。MXU的数据流与此几乎一模一样。

脉动阵列有三种数据流：权重流、激活值流、部分和流。权重从阵列左上角开始注入，沿着每一行逐列向右流动，到达目标PE后锁存下来；激活值从阵列顶部进入，逐行向下流动；每个PE把当前流过的激活值和本地锁存的权重相乘，结果累加到“部分和”上，再传给下面的PE。数据流是“脉动”的核心——数据像心跳一样按固定节奏在阵列中移动。

具体来说，在第 $t$ 个时钟周期，激活值 $x_t$ 进入第一行PE；下一个周期，这个激活值流到第二行，同时新的 $x_{t+1}$ 进入第一行。每个PE在收到激活值时执行一次乘加，并把部分和向下传递。这样整个阵列每个周期都在满负荷工作，不会像传统处理器那样停下来等数据。

就像汉堡工位上，师傅看到面包胚滑下来，就按配方撒料，再把半成品推给下一位师傅；下一周期下一位师傅继续加料，直到完成全部工序。数据流方向对了，整个阵列就“活”了起来。

PE单元设计：一个只会“乘加”的小工位

流水线工位上的师傅不需要会做满汉全席，他只需要记住“肉饼+面包+酱汁”然后机械重复。PE单元就是这样一个“偏科”师傅：它不会跳转、不会分支，只做一件事——乘加运算，但做得极快。

一个PE至少包含三部分：乘法器、加法器和寄存器（锁存器）。乘法器把流入的激活值和本地的权重相乘；加法器把乘积累加到部分和上；寄存器把权重锁存在原位，防止下一个周期数据覆盖。这些PE按二维网格排列，每个PE只与相邻PE通信，所以布线极短、功耗很低。

权重数据在进入阵列后，先通过“数据锁存”让每个PE的权重寄存器保存固定值。之后无论激活值怎么流动，权重都不会变，直到下一次任务重载。这种设计让PE单元非常小巧，可以在一个芯片上集成几万甚至几十万个PE。例如TPU v1的MXU就是一个256×256的PE阵列，单周期能执行65,536次乘加。

工位师傅手里的“配方卡”就是权重寄存器，它被图钉钉住不会飘走；只要面包胚（激活值）一来，师傅就乘以菜量（权重），记到小票（累加器）上。PE越简单，后厨就能排列越多的工位，产能自然越高。

矩阵分块策略：大矩阵怎么塞进小阵列

如果一次要做的汉堡数量超过传送带容量，后厨不会停摆，而是把订单拆成几批，每批做一部分，再拼成完整出货。TPU拿到超大矩阵时同样要“分块”。

MXU阵列尺寸是固定的（例如128×128或256×256），而现实中的矩阵往往更大。分块策略就是把大矩阵按阵列的行数和列数切成小块，小块依次送入MXU计算。比如一个1024×1024的矩阵在256×256的MXU上会切成4×4=16块。

分块不仅解决尺寸问题，还能提高数据复用：每个权重块留在阵列中，多次参与不同激活值块的计算，避免重复搬运权重。实际中还要考虑“边界块”（块小于阵列尺寸时造成浪费）和“块间累加”（相邻块的输出需要相加才能得到最终结果）。

后厨把大订单拆成多批面包胚，每批送上传送带，最后把批次结果拼接起来，顾客收到的还是同一批完美汉堡。分块的本质是“用固定尺寸的流水线，消化任意大小的订单”。

利用率分析：不要让工位闲着

流水线后厨最怕“半空”：传送带上只有零星几个面包胚，后面十个工位等着发呆，厨师们还全额领工资。MXU也有同样的苦恼：当矩阵维度小于阵列尺寸时，很多PE在“空转”。

MXU的利用率可以用实际参与计算的PE数除以总PE数来衡量。假设阵列是 $R \times C$，实际输入的矩阵维度是 $m \times n$。如果 $m\ge R$ 且 $n\ge C$，所有PE都被喂饱，利用率100%；如果矩阵很小，比如 $m=32,n=32$，在256×256的MXU上，只有约1.56%的PE在工作，其余都在“摸鱼”。

为什么小矩阵在小阵列上利用率低？因为脉动阵列的设计目标是让数据尽量就地复用，一旦输入数据量填不满阵列，新增的PE只会白白等待。这也是为什么Google在TPU中加入批处理和向量单元来调度小矩阵，把多个小矩阵拼接成一个大矩阵一起计算，或者把矩阵维度填充到阵列大小。

只有每批订单都填满整条传送带，后厨的每一分工位才算物尽其用；只来一个顾客时，老板只能看着流水线叹气。利用率优化，本质是“让阵列永远吃满”。

MXU演进：从v1到v4的升级之路

当年的后厨只有一条手工流水线，后来老板不断扩建，从单层厂房到多层自动化仓库，再到用机械臂替换部分工人，每一代都带来新的工序安排。Google的TPU从v1到v4，MXU也在不断演化。

TPU v1（2016年）是专用推理芯片，MXU是一个256×256的脉动阵列，峰值算力92 TOPS（int8），主打每一瓦电都用在推理上。TPU v2（2017年）转向训练，增加了可编程性，但单个MXU缩小到128×128，同时引入片上高带宽存储；TPU v3延续v2的架构思路，增加了算力和液冷。TPU v4（2023年）更强调大规模互联，把多个MXU组合成更大的“超级阵列”，还加入了稀疏化支持——跳过无用的零值，进一步提高有效算力。

虽然MXU的尺寸和数量在变，但脉动阵列的核心思想从未改变：用大量小PE、短数据路径和固定数据流换取高吞吐。每一代TPU都在平衡三件事：阵列规模（乘加器数量）、数据供给（存储带宽）和利用率（对多样性矩阵的适应性）。

从手工小作坊到工业级中央厨房，MXU始终是那个“把所有配料流水线化”的核心工序。变的是规模，不变的是脉动。

AI 应用场景

Python 演示

"""模拟一个二维脉动阵列完成矩阵向量乘法权重锁存在每个PE中，激活值从顶部流入，部分和逐行向下累加，最终输出结果。"""class PulseArray:    def __init__(self, rows, cols):        self.rows = rows        self.cols = cols        self.weights = [[0] * cols for _ in range(rows)]        self.accumulator = [[0] * cols for _ in range(rows + 1)]     def load_weights(self, weight_matrix):        for i in range(self.rows):            for j in range(self.cols):                self.weights[i][j] = weight_matrix[i][j]     def run(self, activations):        self.accumulator = [[0] * self.cols for _ in range(self.rows + 1)]        total_cycles = self.rows + len(activations) - 1        for t in range(total_cycles):            for i in range(self.rows):                k = t - i                if 0 <= k < len(activations):                    a = activations[k]                    for j in range(self.cols):                        self.accumulator[i + 1][j] += a * self.weights[i][j]        return self.accumulator[self.rows] def direct_matvec(weight_matrix, activations):    rows = len(weight_matrix)    cols = len(weight_matrix[0])    out = [0] * cols    for i in range(rows):        for j in range(cols):            out[j] += weight_matrix[i][j] * activations[i]    return out if __name__ == "__main__":    w = [[1, 2], [3, 4]]    x = [2, 5]    mxu = PulseArray(2, 2)    mxu.load_weights(w)    print("MXU结果:", mxu.run(x))    print("直接结果:", direct_matvec(w, x))

1. self.weights = [[0] * cols for _ in range(rows)]初始化每个PE的权重寄存器。
2. total_cycles = self.rows + len(activations) - 1计算流水线需要多少个时钟周期才能灌满并排空。
3. k = t - i当前周期第 $i$ 行PE接收到的激活值编号，体现逐行延迟。
4. self.accumulator[i + 1][j] += a * self.weights[i][j]每个PE执行乘加，并把部分和传给下一行。
5. return self.accumulator[self.rows]读取阵列底部的最终输出。

常见误区

面试题