4、深度学习算子解析（下）：矩阵乘（GEMM）与全连接层、归一化（BatchNorm/LayerNorm）算子、注意力机制（Attention）算子
好，咱们接着往下聊。上一节我们把卷积、池化这些基础算子掰开揉碎讲了一遍。这一节，咱们来啃几块更硬的骨头——矩阵乘、归一化，还有注意力机制。说实话，这几块才是现代AI芯片设计里真正考验功力的地方。
核心观点：矩阵乘是算力瓶颈，归一化是精度命脉，注意力是架构灵魂。这三者构成了Transformer时代芯片优化的主战场。
4.1 矩阵乘（GEMM）与全连接层：芯片的“算力黑洞”
先说说矩阵乘。GEMM（General Matrix Multiply）是所有深度学习算子里计算密度最高的一个。我做过一个统计：在BERT这类Transformer模型中，GEMM占了总计算量的70%以上。说白了，芯片能不能跑得快，就看GEMM优化得好不好。
全连接层本质上就是矩阵乘加偏置。公式很简单：Y = X * W + b。但真正落地到芯片上，问题就来了——数据怎么排布？访存怎么优化？
我个人习惯把GEMM的优化分成三个层次：
- 计算层次：用分块（tiling）把大矩阵切成小片，塞进片上SRAM。我见过不少团队一上来就搞大矩阵乘法，结果访存带宽成了瓶颈，算力根本跑不满。
- 数据层次：矩阵的布局（layout）很关键。行主序还是列主序？要不要做数据重排？我在项目中遇到过，同样的GEMM，换了个数据排布方式，性能直接翻倍。
- 指令层次：现代AI芯片都有专门的矩阵乘指令（比如NVIDIA的Tensor Core，或者各种NPU的脉动阵列）。编译器能不能自动识别并生成这些指令，是性能的关键。
// 一个简单的GEMM分块示例（伪代码）
for (int i = 0; i < M; i += TILE_M) {
for (int j = 0; j < N; j += TILE_N) {
for (int k = 0; k < K; k += TILE_K) {
// 将A[i:i+TILE_M][k:k+TILE_K]和B[k:k+TILE_K][j:j+TILE_N]加载到片上
// 执行小矩阵乘累加
C_tile += A_tile * B_tile;
}
}
}
实战技巧：分块大小不是越大越好。我建议根据芯片的SRAM大小和带宽来算——让数据在片上待得越久，访存开销就越小。但分块太大，又会浪费片上空间。这个平衡点，得靠实际跑profile来调。
4.2 归一化算子：BatchNorm与LayerNorm
归一化算子，说白了就是让数据分布“规整”一点。BatchNorm和LayerNorm是两种最常见的方案。它们做的事情很像，但应用场景完全不同。
BatchNorm：对每个特征通道，在batch维度上做归一化。训练时维护均值和方差的滑动平均，推理时直接用。嗯，这里要注意——推理时BatchNorm可以融合到前面的卷积或全连接层里，省掉一次额外的计算。我在做芯片部署时，这个融合优化是必做的。
LayerNorm：对每个样本，在特征维度上做归一化。Transformer里用的就是它。为什么不用BatchNorm？因为Transformer的输入长度是变化的，batch维度上统计不稳定。你想想看，如果一句话10个词，另一句话100个词，在batch维度上做归一化，短的样本会被长的样本“带偏”。
| 特性 | BatchNorm | LayerNorm |
|---|---|---|
| 归一化维度 | Batch（N） | 特征（C/H/W） |
| 训练/推理差异 | 大（需维护滑动平均） | 小（直接计算） |
| 适用场景 | CNN、固定batch size | RNN、Transformer、变长输入 |
| 芯片优化点 | 推理时可融合 | 需计算均值和方差，访存密集 |
避坑指南：我曾经在部署一个语音模型时，发现推理结果和训练结果对不上。查了两天才发现——BatchNorm在推理时用的滑动平均参数，和训练时的实际分布有偏差。解决方案是：在训练完成后，用一小批数据重新校准一下BatchNorm的统计量。这个坑，踩过的人都知道。
4.3 注意力机制（Attention）算子：Transformer的“心脏”
注意力机制，尤其是自注意力（Self-Attention），是近年来最重要的算子之一。它的核心公式是：
Attention(Q, K, V) = softmax(Q * K^T / sqrt(d_k)) * V
这个公式看着简单，但计算起来可不轻松。Q、K、V都是矩阵，中间要算一个巨大的注意力分数矩阵（Q * K^T），然后做softmax，最后再乘V。这个过程中，有两个地方特别耗资源：
- Q * K^T：这是一个矩阵乘，但结果矩阵的大小是序列长度×序列长度。如果序列长度是1024，这个矩阵就有100万个元素。我见过一些芯片，算力够，但显存带宽被这个中间结果给撑爆了。
- Softmax：需要对每一行做指数运算和求和。指数运算在硬件上很慢，而且容易溢出。很多芯片会用一个近似计算单元来做softmax，但精度会损失。
怎么优化？我分享几个实战经验：
- Flash Attention：把Q、K、V分块，在片上完成整个注意力计算，避免写回中间结果。这个技术现在很火，我建议每个做芯片的人都去读一下原始论文。
- 稀疏注意力：不是每个位置都需要关注所有位置。比如在图像里，一个像素只需要关注它附近的像素。用稀疏矩阵乘代替稠密矩阵乘，能省不少算力。
- 量化：注意力分数矩阵可以用INT8甚至INT4来存。我在一个项目里试过，把注意力分数从FP16降到INT8，精度损失不到0.1%，但显存占用直接减半。
// 标准自注意力计算（PyTorch风格）
def attention(Q, K, V):
d_k = Q.size(-1)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
attn_weights = torch.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, V)
return output
我的建议：如果你在芯片上实现注意力算子，不要直接照搬上面的代码。把Q、K、V的分块大小和芯片的SRAM大小对齐，把softmax的指数计算用查找表（LUT）近似，把矩阵乘用脉动阵列加速。这样一套组合拳下来，性能能提升3-5倍。
4.4 本章知识体系总览
为了让你更直观地理解这几个算子的关系，我画了一张图。这张图展示了从输入到输出的数据流，以及每个算子在其中的位置和作用。
从这张图里你能看到，GEMM是计算的主力，归一化负责稳定训练，注意力则是整个流程的“大脑”。在芯片设计时，这三者的优化策略完全不同——GEMM要堆算力，归一化要减访存，注意力要平衡计算和带宽。
好了，这一节的内容就到这儿。这几个算子，每一个都值得你花时间去深挖。尤其是注意力机制，未来几年它还会是AI芯片优化的核心战场。希望今天的分享能帮你少走一些弯路。
一句话总结：矩阵乘是算力的“发动机”，归一化是精度的“稳定器”，注意力是模型的“指挥中心”。三者在芯片上的映射优化，决定了AI芯片的最终性能天花板。
想要更多资料，联系可可