矩阵乘法加速：Transformer的心脏与引擎

各位同学，今天我们来聊聊Transformer里最核心、最吃算力的操作——矩阵乘法。说实话，我做了这么多年芯片设计，见过太多算法在理论上跑得飞快，一到硬件上就卡成PPT。而矩阵乘法，就是那个最容易让硬件“露怯”的地方。

Transformer里到处都是矩阵乘法。Attention机制里的Q、K、V计算，FFN层的两个线性变换，甚至Embedding层都离不开它。我统计过一个典型的BERT-Large模型，矩阵乘法占了总计算量的70%以上。你想想看，如果这个环节不加速，整个模型就别想跑起来。

为什么矩阵乘法在Transformer里这么重要？

说白了，Transformer的本质就是“通过矩阵运算做信息变换”。Attention机制里，我们计算Q和K的点积得到注意力分数，再和V相乘得到加权结果。每一步都是矩阵乘法。FFN层更是直接两个大矩阵相乘。

我举个例子，假设你的序列长度是512，隐藏维度是768。那么Q、K、V每个都是512×768的矩阵。Q×K^T这一步，就是512×768乘以768×512，结果是512×512。这还只是一层，一个Transformer模型动辄12层、24层。算下来，一次推理就要做几十次甚至上百次大矩阵乘法。

嗯，这里要注意：矩阵乘法的计算复杂度是O(n³)。序列长度一增加，计算量就爆炸式增长。这也是为什么现在大家都在研究长序列Transformer，但硬件跟不上——矩阵乘法加速没做好。

Systolic Array（脉动阵列）原理与设计

好，问题来了：怎么加速矩阵乘法？

传统CPU做矩阵乘法，是一个元素一个元素地算，慢得让人抓狂。GPU虽然快，但功耗高、面积大。我们需要一种专门为矩阵乘法设计的硬件结构——脉动阵列。

脉动阵列这个名字很形象。想象一下心脏跳动，血液被泵到全身。脉动阵列也是这样，数据像血液一样在计算单元之间“脉动”流动，每个单元只做简单的乘加运算，但组合起来就能高效完成矩阵乘法。

我最早接触脉动阵列是在2017年，当时在做一个AI加速器项目。说实话，第一次看到这个结构时觉得太巧妙了——用最简单的计算单元，通过巧妙的连接方式，实现了极高的计算效率。

下面我画了一张脉动阵列的结构图，帮你理解它的核心逻辑：

每个PE（Processing Element）就是一个乘加单元。权重数据从左向右流动，输入数据从上向下流动。每个PE把当前权重和输入相乘，再加上从上方传来的部分和，结果传给下一个PE。这样，整个阵列就像一条流水线，数据有节奏地流动，效率极高。

脉动阵列的三种数据流模式

脉动阵列虽然结构固定，但数据怎么流动却大有讲究。我把它分成三种模式：Weight Stationary、Input Stationary、Output Stationary。说白了，就是“谁不动，谁动”的问题。

1. Weight Stationary（权重固定）

这种模式下，权重数据预先加载到每个PE里，然后就不动了。输入数据从上方流入，部分和从左向右传递。每个PE把固定的权重和流动的输入相乘，累加结果。

我个人习惯在模型推理场景用这种模式。为什么？因为权重是固定的，可以提前加载好。推理时只需要输入数据流动，省去了权重搬运的开销。我在做BERT加速器时，就是用的Weight Stationary，效果不错。

2. Input Stationary（输入固定）

反过来，输入数据固定不动，权重数据流动。每个PE保存一份输入，权重从左向右流过，每个PE把权重和本地输入相乘，结果向下传递。

这种模式适合训练场景。训练时输入数据（比如一个batch的样本）相对固定，权重需要频繁更新。Input Stationary让权重流动，正好匹配训练时权重不断变化的特点。

3. Output Stationary（输出固定）

这种模式比较特殊。部分和结果固定在一个PE里，权重和输入都流动。每个PE收到权重和输入后，乘积累加到本地保存的部分和上。

嗯，这里要注意：Output Stationary适合做矩阵乘法中的“外积”操作。我在做卷积加速时用过这种模式，但Transformer里用得不多。不过了解它还是有必要的，说不定哪天就用上了。

数据流模式	固定数据	流动数据	适用场景	我的评价
Weight Stationary	权重	输入、部分和	推理	最常用，效率高
Input Stationary	输入	权重、部分和	训练	适合权重更新频繁的场景
Output Stationary	部分和	权重、输入	外积计算	Transformer里用得少

脉动阵列的优缺点分析

说了这么多好处，脉动阵列有没有缺点？当然有。我踩过的坑可不少。

优点：

• 高吞吐量：每个时钟周期，所有PE都在工作。没有闲置单元。
• 低功耗：数据只在相邻PE之间传递，不需要全局总线。功耗比GPU低得多。
• 可扩展：想要更大规模的矩阵乘法？多加几个PE就行。结构不变。
• 局部性：数据复用率高，减少了内存访问次数。

缺点：

• 灵活性差：脉动阵列是为固定大小的矩阵乘法设计的。矩阵大小变了，效率就下降。我曾经遇到过一个模型，矩阵大小不是PE阵列的整数倍，结果有一半的PE在空转。
• 数据加载开销：虽然计算时数据流动很高效，但初始加载数据需要时间。如果矩阵太小，加载时间比计算时间还长。
• 控制复杂：需要精确控制数据流动的时序。稍微有点偏差，结果就全错了。
• 面积开销：每个PE都需要本地存储和互联逻辑。PE数量一多，芯片面积就上去了。

好了，这一章的内容就到这里。矩阵乘法加速是Transformer硬件加速的基石，脉动阵列是目前最成熟的方案。下一章我们会深入脉动阵列的具体实现细节，包括PE的内部结构、数据调度策略，以及如何用Verilog写一个简单的脉动阵列。