3、计算瓶颈分析：矩阵乘法（GEMM）与注意力机制（Attention）的计算特性与瓶颈
聊到大模型推理，大家第一反应往往是「显存不够」。但说实话，我这些年调优下来，发现真正卡脖子的地方，很多时候是计算本身。尤其是矩阵乘法（GEMM）和注意力机制（Attention），这两个家伙几乎吃掉了模型推理的绝大部分时间。
今天我们就来扒一扒，它们到底是怎么「吃时间」的，瓶颈又在哪里。
3.1 矩阵乘法（GEMM）：大模型的心脏
GEMM 就是 General Matrix Multiply，说白了就是两个矩阵相乘。Transformer 里的每个线性层、每个投影，本质上都是 GEMM。你想想看，一个 7B 模型，光前向传播就要做几千次矩阵乘法。
它的计算特性其实很单纯：
- 计算密集：一次 M×K 乘以 K×N，需要 2×M×N×K 次浮点运算。这个数字大得吓人。
- 访存密集：每次计算都要把矩阵从显存搬到计算单元。带宽不够，计算单元就得干等。
我个人习惯把 GEMM 的瓶颈分成两类：
| 瓶颈类型 | 典型场景 | 表现 |
|---|---|---|
| 计算瓶颈 | 大矩阵相乘（M、N、K 都很大） | GPU 计算单元跑满，但访存带宽还有余量 |
| 访存瓶颈 | 小矩阵相乘（比如单头注意力里的 QK^T） | 计算单元在等数据，利用率很低 |
我在项目中遇到过最典型的例子：用 FP16 跑 4096×4096 的矩阵乘法，计算单元利用率能到 80% 以上。但换成 128×128 的小矩阵，利用率直接掉到 10% 以下。为什么会这样？因为小矩阵的访存开销占比太大了。
核心结论：GEMM 的瓶颈不是固定的。它取决于矩阵尺寸、数据类型、硬件架构。调优的第一步，就是搞清楚你的 GEMM 到底被什么卡住了。
3.2 注意力机制（Attention）：计算与访存的双重挑战
Attention 比 GEMM 复杂一些。它由四个步骤组成：
- QKV 投影：三个 GEMM，把输入映射到 Q、K、V
- QK^T 计算：又一个 GEMM，计算注意力分数
- Softmax：对每行做归一化
- 加权求和：Softmax 结果乘以 V，还是 GEMM
你看，Attention 本质上就是「GEMM + Softmax + GEMM」的组合。但它的瓶颈点很特殊。
3.2.1 计算特性分析
Attention 的计算量跟序列长度 L 是平方关系。具体来说：
- QK^T 的计算量：2 × L × L × d_k（d_k 是头维度）
- Softmax 的计算量：L × L（指数运算 + 求和 + 除法）
- 加权求和的计算量：2 × L × L × d_v
当 L 从 1024 涨到 4096，Attention 的计算量翻了 16 倍。嗯，这里要注意，这还没算显存开销。
3.2.2 访存特性分析
Attention 的访存模式很「不友好」：
- QK^T 的结果矩阵：L×L，需要写回显存，然后马上读出来做 Softmax
- Softmax 的中间结果：同样 L×L，又要写回显存
- 加权求和：再读一次 V 矩阵
我曾经调试过一个 32K 序列长度的推理任务，发现 Attention 阶段有 70% 的时间花在「把数据搬来搬去」上，真正做计算的不到 30%。这就是典型的访存瓶颈。
避坑指南：我曾经在优化长序列推理时，花了两周时间调 GEMM 的 tile size，结果发现瓶颈根本不在计算，而在 Attention 的中间结果读写上。后来改用 Flash Attention，直接把端到端延迟砍了一半。
3.3 知识体系与核心逻辑
下面这张图是我自己总结的，把 GEMM 和 Attention 的计算瓶颈梳理清楚了：
3.4 实战中的瓶颈判断方法
光知道理论不够，得能动手判断。我一般用三个指标：
- 计算利用率：看 GPU 的计算单元有没有跑满。用 nvidia-smi 看 sm 利用率，低于 50% 基本就是访存瓶颈。
- 带宽利用率：看显存带宽有没有跑满。用 nsys 或者 Nsight Compute 看 dram__bytes_read 和 dram__bytes_write。
- 算术强度：计算量除以访存量。算术强度越高，越偏向计算瓶颈。
小技巧：我个人习惯在 profiling 时，先看 kernel 的「occupancy」和「memory bandwidth」。如果 occupancy 高但带宽低，说明计算瓶颈；反过来，带宽高但 occupancy 低，就是访存瓶颈。这个判断方法我用了好几年，基本没翻过车。
3.5 一个具体的例子
假设我们跑一个 7B 模型的推理，输入序列长度 2048，batch size 1。用 Nsight Compute 抓一下 profile：
Kernel: gemm_nt_64x64x64
Duration: 1.23 ms
Compute (SM) Throughput: 45.2%
Memory Throughput: 78.1%
Arithmetic Intensity: 12.3 FLOPs/Byte
Kernel: attention_forward
Duration: 3.45 ms
Compute (SM) Throughput: 22.1%
Memory Throughput: 91.3%
Arithmetic Intensity: 3.8 FLOPs/Byte你看，GEMM 的算术强度 12.3，计算利用率 45%，说明它处于「计算和访存混合瓶颈」的状态。而 Attention 的算术强度只有 3.8，内存带宽跑到了 91%，明显是访存瓶颈。
这时候优化 Attention 比优化 GEMM 收益更大。我当时就是看到这个数据，才决定上 Flash Attention 的。
3.6 小结
GEMM 和 Attention 是推理延迟的两大「元凶」。GEMM 的瓶颈取决于矩阵尺寸，Attention 的瓶颈则更多来自访存模式。判断瓶颈不能靠猜，得用 profiling 工具看数据。
下一节我们会聊具体的优化手段，包括算子融合、Flash Attention、量化等。但在此之前，我建议你先跑一遍 profiling，搞清楚自己的模型到底被什么卡住了。方向对了，优化才有效果。
公众号:蓝海资料掘金营，微信deep3321