MXU矩阵乘单元：TPU的核心计算引擎

可视化演示

展示256x256 PE脉动阵列数据流、数据复用模式及MXU与GPU张量核架构性能对比。

矩阵乘法之重

你可能会问：为什么一座快递分拣中心，非得专门建一个巨型车间来处理“包裹面单匹配”这件事？因为在AI的世界里，“矩阵乘法”就是那个无处不在的包裹面单匹配。神经网络里的每个神经元，做的本质工作就是“把一堆输入数字各自乘上一个权重，再全部加起来”。几十亿个神经元堆在一起，你算一算，一次推理要做多少次乘加？答案是几十亿次起步，训练时更是要来回算上几万亿次。可以说，AI计算的九成以上，都是在做矩阵乘法这件事。

传统计算机面对这种海量乘法时，会撞上一堵墙——“内存墙”。冯·诺依曼结构里，计算单元和内存是分开的，CPU得先把数据从内存里“搬”到计算单元，算完再“搬”回去。搬数据的速度，远跟不上算数据的速度。就像快递分拣中心如果只有一个人，他需要反复跑到仓库深处去翻找面单对照表，那他一天也分拣不了几个包裹。

MXU的破局思路是“空间换时间”。既然搬数据慢，那就把要用的数据一次性摊开，摆在一个巨大的工位阵上，让数据在工位之间像流水一样自己流动起来。谷歌的第一代TPU里，MXU直接配了256×256个PE（处理单元）工位，也就是65536个工人同时开工。这相当于把原来只有一条传送带的小作坊，升级成了2.56万个工位同步运转的超级工厂。

回到我们的比喻：矩阵乘法如此重要，就像面单匹配是分拣中心的灵魂。MXU的存在，就是把这个灵魂动作极致地工厂化——让每一件包裹经过每个工位时，工人只需做一个简单动作，剩下的交给流水线本身。

脉动阵列大阅兵

现在让我们走进MXU的核心——脉动阵列。想象你站在一个巨大的厂房里，2.56万名工人（PE）站成一个方方正正的棋盘，每人面前一台像计算器一样的小机器。有趣的是，每个工人不需要看任何屏幕或操作台，他们只需要做两个动作：把收到的两个数字乘起来，再加到自己本子上记的那个数上。就这么简单。工人之间没有复杂的电话线，只靠一条条短短的滑槽（寄存器直连）把包裹递给左边、右边或前方的邻居。

这就是“脉动阵列”名字的由来——数据不是乱哄哄地各自乱跑，而是像血管里的血液一样，随着心脏（时钟）的每一次跳动，整齐地向前涌动一格。每个时钟周期，每个PE完成一次“乘法+累加”（MAC操作），然后把结果传给下一个PE。这里的“下一次”，不是随机任务，而是精心安排的流水线接力。

流水线的妙处在于“时间上重叠”。第一个数据刚走完第一列，第二个数据就已经开始进入第一列了。就像包裹流水线上，第一个包裹还在第3个工位时，第100个包裹已经滑进了第1个工位。这样一来，2.56万个PE几乎每个周期都在满负荷运转，而不是像传统架构那样，大部分时间都在“等待下一个数据”。

脉动阵列的威力，不在单个工人有多快，而在“阵列”本身——2.56万个工人步调一致地重复简单动作，汇成一股每秒千万亿次的算力洪流。这就是它和普通处理器的本质区别：后者是少数几个全能冠军轮流上阵，前者是2.56万个专才永不疲倦地同步开火。

权重钉在工位上

快递分拣中心有一个巧妙的设计：每个工位的小抽屉里，预先塞好了一张“面单对照表”——这不是谁临时塞的，而是开工前，整个厂房的管理系统根据今天要处理的包裹批次，统一给每个工位配好的。工人们开工后，只需要低头看一眼抽屉里那张固定不变的小抄，不需要去仓库翻任何手册。这张小抄，在MXU的世界里，就是“权重”。

为什么要这样做？因为在一个计算批次中，权重矩阵会被反复使用。如果用传统方式，每个PE每做一次乘法，都要跑到内存里去取一次对应的权重值，那2.56万个工位就是2.56万个同时冲向仓库的快递员，仓库的吞吐能力瞬间被挤爆。这就是“内存墙”的具体体现。MXU的做法是：先把整个权重矩阵一次性载入每个PE的寄存器阵列——每个PE的寄存器里，稳稳地存着它负责的那一个权重值。

于是，在推理过程中，权重“纹丝不动”，而输入数据（激活值）像传送带上的包裹一样，一列一列地流经每个工位。每流经一个工位，就完成一次“当前权重×当前输入”的乘法。这样，同一份权重被反复使用了成千上万次，而取权重的次数，被压缩到了最小——只在最开始装载一次。这就是“权重固定”数据复用的精髓。

“权重钉在工位上”这个设计，把最费时的数据搬运工作压缩到了极致。试想，如果每个包裹都要工人跑到仓库核对一次面单，那分拣中心早就瘫痪了。正因为小抄就在手边，工人们才能以每秒数十亿次的速度，行云流水地完成分拣——MXU 也正是靠这个绝招，在 2015 年就以远超同代 GPU 的能效，惊艳了整个 AI 界。

数据流动模式

权重视为“推不动”的固定工位，那数据呢？数据是“流动”的。在 MXU 的脉动阵列里，输入数据（激活值）不是一次性广播给所有PE，而是沿着行方向，一个周期一个周期地“接力推进”。每个PE从左侧邻居那里接过数据，用它做一次乘法累加，然后在下个周期把它传给右侧邻居。就像快递流水线上，包裹不会瞬间瞬移到最后一站，而是跟着传送带一格一格地走——但正因为每格都有工人处理，整条线并没有“等待”的时间。

数据流动的深层妙处，是“半成品”也沿着对角线方向自然传播。假设我们要计算 C = A×B，C 矩阵里每个元素是 A 的一行和 B 的一列的点积。在脉动阵列中，A 的各个元素沿行方向流动，B 的各个元素沿列方向流动，它们在每个 PE 处“相逢”，相乘后累加，累加的结果（部分和）被 PE 保留下来，等待和下一个到达的元素继续累加。最终，当最后一批数据流过之后，每个 PE 的累加器里，就静静躺着结果矩阵的一个元素。

为了让这幕配合完美，所有数据必须步调一致，这就是“脉动”的含义。整个阵列由全局时钟统一控制：每个周期，每个PE同时完成“接收左邻居的数据→乘上自己寄存器里的权重→加到累加器→把数据传给右邻居”。不允许谁快一步，也不允许谁慢一步。这种严格的同步，让数据流动像训练有素的仪仗队，整齐划一。

在数据流动的世界里，不存在“把数据从远处取来”的概念——数据永远来自身边的邻居，而结果永远积攒在自己的本子上。这种“让数据动起来，而不是让工人跑起来”的哲学，正是脉动阵列对冯·诺依曼结构最优雅的反叛。

硬件运算链

现在我们把镜头拉近，看看单个PE的内部构造。每个PE，本质上是一台微型计算站，里面有三个核心部件：一个乘法器（负责把两个数相乘）、一个加法器（把乘法和之前的累加值相加）、以及一组寄存器（负责暂存数据）。它们排成一条微型流水线：数据从左边的邻居进来，乘法器立刻把它和寄存器里的权重相乘，乘积送入加法器，与累加器里的旧值相加，新值存回累加器——一个周期内完成全部动作。

为了让这套流水线高速运转，MXU 在物理设计上做了两个关键决策。第一，寄存器直接嵌在PE旁边，数据从上一个PE传过来不需要经过复杂的总线，直接通过窄而短的导线“手递手”传递，延迟极低。第二，整个阵列由统一时钟驱动，所有PE的“心跳”完全同步，避免了异步设计带来的协调开销。正是这种高度规整的物理布局，让 2.56 万个PE可以在单个芯片上愉快地并肩作战，而不会因为布线问题互相打架。

MXU 的规模通常表达为 256×256——即 256 行、256 列，共 65536 个PE。这个数字不是随意定的，它和 TPU 的片上存储（SRAM）容量、内存带宽之间是精心匹配的。每个时钟周期，MXU 可以从左边的缓冲器取 256 个激活值，从下方的缓冲器取 256 个权重，完成 65536 次乘法累加。如果主频是 700MHz，那么峰值算力就是 65536×2×700M ≈ 92 TOPS——这里的"×2"是因为一次 MAC 算一次乘法和一次加法，共两次浮点操作。

回忆快递分拣中心：每个工人（PE）的工位都配有一台固定的小计算器（乘法器+加法器）和一本流水账本（累加器）。2.56万个工人，人人手边有账本，每秒钟计算器同步按下 7 亿次，整座工厂的吞吐量自然惊人。MXU 的全部秘密，就是把这套极其简单的工位，乘以 2.56 万份，再用时钟完美地同步起来。

利用率与演进

脉动阵列如此强大，但它并非“银弹”。MXU 有个挥之不去的烦恼：形状匹配。脉动阵列就像一个固定尺寸的棋盘——256×256，只能高效处理 256×256 的矩阵乘法。如果输入矩阵的某一维不是 256 的整数倍，就会出现“部分工位空转”的现象。比如要算一个 100×100 的矩阵乘法，会有超过 80% 的工位闲着。这就像分拣中心有 2.56 万个工位，但今天只来了 100 件包裹，绝大多数工人只能干瞪眼。

另一个瓶颈是数据搬运开销。权重固定解决了权重重复取用的问题，但输入数据（激活值）依然是“流”进来的——它必须先被从内存搬到芯片上的缓冲器里，才能喂给阵列。如果阵列算得太快，而数据供给跟不上，MXU 就会“等米下锅”，利用率大幅下降。谷歌的解决方案之一是“批处理”：一次喂一批数据进来，让权重被更多数据复用，摊薄单个数据的搬运成本。批处理越大，复用率越高，利用率越好——这也是为什么 TPU 特别适合大 batch 的推理和训练场景。

MXU 自身的演进也印证着这些思考。TPU v1（2015）是推理专用，MXU 搭载 256×256 阵列和 28MB 片上存储；TPU v2/v3 走向训练，MXU 改为“脉动阵列+向量单元”双引擎，并引入“数据流”（数据沿列方向流动）模式；TPU v4 则把超级计算机的互联拓扑融入 MXU 集群，让 4096 块芯片协同工作。每一代的宿敌没变——都是那堵“内存墙”，每一代的破局思路也没变——继续提高数据复用，减少搬运。

回顾全文，MXU 就是一座把“做乘法”这件事极致工厂化的超级分拣中心。它的聪明之处不是让单个工人变快，而是用权重固定、数据流动、脉动同步这三大设计，让 2.56 万个工人像一个人一样高效协同。理解 MXU，你就拿到了理解 TPU 乃至整个 AI 芯片设计哲学的金钥匙。

AI 应用场景

Python 演示

import numpy as np class PESystolicArray:    """模拟一个 NxN 的脉动阵列，权重固定模式。"""     def __init__(self, n, weight_matrix):        self.n = n        # 每个 PE 存放一个权重值        self.weights = np.array(weight_matrix, dtype=np.float32)        # 每个 PE 的累加器清零        self.accumulators = np.zeros((n, n), dtype=np.float32)     def run(self, a_matrix):        """a_matrix 形状为 (n, n)，按行流入阵列。"""        a = np.array(a_matrix, dtype=np.float32)        # 初始化累加器        self.accumulators.fill(0.0)         for t in range(2 * self.n - 1):            for i in range(self.n):                j = t - i                if 0 <= j < self.n:                    # PE(i, j) 从左方接收 a[i][j] 并乘上权重                    self.accumulators[i][j] += a[i][j] * self.weights[i][j]        return self.accumulators.copy()  class MXUEmulator:    """MXU 仿真器：分批处理大矩阵乘法。"""     def __init__(self, array_size=4):        self.array_size = array_size        self.total_mac = 0        self.idle_cycles = 0     def multiply(self, a, b):        """计算 A (m x k) 与 B (k x n) 的乘积，使用 2D 分块。"""        m, k = a.shape        k2, n = b.shape        assert k == k2        c = np.zeros((m, n), dtype=np.float32)        s = self.array_size         for i0 in range(0, m, s):            for j0 in range(0, n, s):                for k0 in range(0, k, s):                    a_block = a[i0:i0+s, k0:k0+s]                    b_block = b[k0:k0+s, j0:j0+s]                    if a_block.shape[0] < s:                        a_block = np.pad(a_block, ((0, s - a_block.shape[0]), (0, 0)))                    if a_block.shape[1] < s:                        a_block = np.pad(a_block, ((0, 0), (0, s - a_block.shape[1])))                    if b_block.shape[0] < s:                        b_block = np.pad(b_block, ((0, s - b_block.shape[0]), (0, 0)))                    if b_block.shape[1] < s:                        b_block = np.pad(b_block, ((0, 0), (0, s - b_block.shape[1])))                    pe_array = PESystolicArray(s, b_block.T)                    partial = pe_array.run(a_block)                    c[i0:i0+s, j0:j0+s] += partial[:min(s, m-i0), :min(s, n-j0)]                    self.total_mac += s * s * s        return c  def main():    m, n, k = 8, 6, 7    rng = np.random.default_rng(42)    a = rng.normal(size=(m, k))    b = rng.normal(size=(k, n))     emulator = MXUEmulator(array_size=4)    c_mxu = emulator.multiply(a, b)    c_np = a @ b     print("MXU 仿真结果:\n", c_mxu)    print("NumPy 参考结果:\n", c_np)    print("最大误差:", np.max(np.abs(c_mxu - c_np)))    print("总 MAC 次数:", emulator.total_mac)  if __name__ == "__main__":    main()

1. class PESystolicArray定义单个脉动阵列块，内部用二维数组模拟 PE 网格。
2. self.weights = np.array(...)把权重矩阵“钉”在每个 PE 的寄存器中，对应权重固定设计。
3. for t in range(2 * self.n - 1)模拟脉动时序：数据沿对角线逐拍流入，共需 2n-1 个周期。
4. self.accumulators[i][j] += a[i][j] * self.weights[i][j]每个 PE 完成一次乘累加，对应 MAC 操作。
5. class MXUEmulator外层仿真器，处理任意形状矩阵的分块计算。
6. for i0 in range(0, m, s)按 4x4 块遍历大矩阵，把大矩阵乘法拆成多个小方块，让每个块都适配脉动阵列尺寸。
7. np.pad当矩阵边缘不足 4x4 时补零，模拟真实硬件中“形状不匹配导致的部分空转”。
8. self.total_mac += s * s * s统计完成的总乘累加次数，用于评估硬件利用率。

常见误区

面试题