GPU系统拆解-08-Tensor Core、GEMM 与 FlashAttention：AI 计算热点为什么这样组织

本文是「GPU系统拆解」系列第 08 篇。
系列导读：GPU系统拆解-00-导读：从架构认知到推理系统的学习路线
上一篇：GPU系统拆解-07-PTX、SASS 与编译链：一段 CUDA 代码如何变成指令
下一篇：GPU系统拆解-09-Profiling 与性能定位：先找到瓶颈，再谈优化

这一篇要解决的是一个很核心的问题：为什么 AI 推理的主战场总是在 GEMM、Tensor Core、Attention kernel 这些地方。学懂这一篇之后，你再看 cuBLAS、CUTLASS、FlashAttention、TensorRT-LLM、vLLM，会更容易把它们放回同一条主线里理解。

1. 先给结论

  • 大模型推理的大头，通常不是 Python 或框架调度，而是 GEMM 和围绕它展开的数据搬运。
  • Tensor Core 的本质是专门做矩阵块乘加的高吞吐硬件，不是“更快的 CUDA Core”。
  • 想吃到 Tensor Core 的红利，前提是数据类型、布局、对齐和问题规模都要适合。
  • 高性能 GEMM 的关键不是公式，而是多层 tiling、数据复用和搬运-计算重叠。
  • Attention 虽然包含矩阵乘，但真正难的地方在于 softmax、中间结果管理和大量 I/O。
  • FlashAttention 的价值主要不是数学变了，而是它按 GPU 内存层次重组了计算顺序，减少了 HBM 读写。
  • prefill 更像 dense compute，decode 更容易暴露 KV cache 和带宽瓶颈。
  • 量化 在推理里重要，不只是为了省显存，也是为了更适配低精度高吞吐硬件路径。

2. 为什么 AI 推理离不开 GEMM

把一个 Transformer 或 LLM 拆开看，很多最重的计算都可以写成矩阵乘。

最典型的是线性层：

1
Y = XW + b

如果：

  • X 是 [M, K]
  • W 是 [K, N]

那么核心计算就是：

1
[M, K] x [K, N] = [M, N]

也就是标准 GEMM。

在 LLM 里，这种结构会反复出现：

  • Q = XWq
  • K = XWk
  • V = XWv
  • O = Attn(Q, K, V)Wo
  • MLP_up = XWup
  • MLP_down = HWdown

所以从硬件视角看，现代推理系统的主要工作其实可以概括成两件事：

  1. 尽量让大矩阵乘走到最强计算路径
  2. 尽量让数据搬运不要拖垮整体吞吐

这也是为什么 GPU 优化经常不是“把某个 API 用熟”，而是“把算子改写成更接近 GEMM 和 tile 的形状”。

3. Tensor Core 到底是什么

早期 GPU 主要依赖更通用的标量或向量计算路径，也就是大家常说的 CUDA Core 路径。它们很通用，但对深度学习里最常见的矩阵块乘加来说，不够专门。

Tensor Core 可以先粗略理解成：

专门为矩阵块乘加设计的高吞吐硬件路径。

它最适合的工作，不是零散标量运算，而是这类模式：

1
D = A x B + C

也就是一小块矩阵的乘加累积。

所以它和 CUDA Core 的区别，不只是“更快”，而是设计目标不同：

  • CUDA Core 强在通用性
  • Tensor Core 强在规则、密集、可分块的矩阵计算

这也是为什么深度学习 GPU 会越来越强调低精度矩阵吞吐，而不是只堆通用浮点能力。

4. 想用好 Tensor Core，需要满足什么前提

Tensor Core 不是“开了就快”。它对问题形状有明确偏好。

4.1 数据类型要合适

Tensor Core 通常更擅长这些类型：

  • FP16
  • BF16
  • TF32 的特定路径
  • INT8
  • 更低精度的量化类型

这背后不是“数值技巧”，而是硬件就是围绕这些低精度高吞吐路径设计的。推理里常见的混合精度和量化，本质上都是在主动贴近这条路径。

4.2 数据布局要合适

硬件喜欢规则、对齐、可分块的数据布局，不喜欢零散和碎片化访问。

所以工程里经常会看到：

  • weight packing
  • layout transform
  • pre-transpose
  • 某些维度要求按固定粒度对齐

这不是库作者故意把事情搞复杂，而是因为硬件真的有自己的偏好。

4.3 问题规模也要合适

如果矩阵太小、batch 太小、shape 太碎，即使数学上能走 Tensor Core，实际也可能吃不满。

原因通常是：

  • kernel launch 开销占比上升
  • tile 利用率下降
  • 数据复用不足
  • 寄存器和 shared memory 的投入回报不高

这也是为什么推理系统总在做：

  • batching
  • continuous batching
  • shape bucketing
  • request merging

从硬件角度看，它们的目标很一致：把任务组织成更适合 GPU 吞吐模式的形状。

5. 为什么 naive GEMM 一定不够快

最朴素的矩阵乘写法通常是：每个线程负责一个输出元素，然后沿着 K 维循环累加。

它能算对，但性能很差，问题主要有四个：

  • 每个线程都反复从 global memory 取数，远距离访问太多
  • 相邻线程使用了很多相同数据，却没有显式组织复用
  • 没有形成 block -> warp -> instruction 的分块层次
  • 完全吃不到 Tensor Core 这类专用矩阵路径

所以 naive GEMM 的教学意义，不是告诉你“矩阵乘怎么写出来”，而是让你看清：

高性能矩阵乘的核心，不是公式变了，而是数据组织方式彻底变了。

6. 高性能 GEMM 的真正骨架

现代高性能 GEMM 的主线很稳定，本质上就是围绕数据复用建立多层 tiling。

可以先把它想成这样一条路径：

1
2
3
4
5
6
global memory
  -> shared memory
  -> register fragments
  -> Tensor Core / MMA
  -> accumulator registers
  -> global memory

这条路径背后的目标有三个：

  1. 把远距离数据搬运次数降下来
  2. 让同一块数据被更多次复用
  3. 让最热的数据停留在更快的层级

6.1 为什么一定要 tiling

如果不分块，A 和 B 的很多元素会被重复从显存读取。

一旦引入 tiling：

  • 一个 block 先把 A_tile 和 B_tile 搬到 shared memory
  • 同一块数据可以被一组线程反复使用
  • 再把更小的片段喂到寄存器和 Tensor Core

这样显存流量和片上复用就都变得更可控。

6.2 不是一层 tile，而是多层 tile

真正高性能的 GEMM 通常至少有三层结构：

  • block tile
  • warp tile
  • mma instruction tile

对应的思维方式也变了：

  • 不是“一个线程算一个元素”
  • 而是“一组线程协同维护一个输出 tile”

这也是为什么后面学 CUTLASS 时会感觉它像一套模板化的分层映射系统。

6.3 shared memory 和寄存器各自解决什么问题

shared memory 的意义不只是“快”，更准确地说，它是程序员显式管理的数据复用缓冲区。

它主要负责：

  • 让线程协作加载数据
  • 把分散的显存访问变成结构化的局部访问
  • 为后续寄存器和 Tensor Core 提供中转层

寄存器则更像最热的数据工作区：

  • 保存 fragment
  • 保存累加器
  • 承接最频繁的中间结果读写

所以高性能路径不是简单的“多用 shared memory”，而是让不同层级各司其职。

6.4 为什么搬运和计算要重叠

如果流程总是：

  1. 先把数据搬完
  2. 再开始算
  3. 再搬下一批
  4. 再继续算

那么 GPU 大量周期都会浪费在等待上。

更好的方式是双缓冲或多 stage pipeline：

1
2
计算 tile t
同时预取 tile t+1

这就是为什么从 Ampere 到 Hopper，async copy、barrier、pipeline、TMA 这些特性越来越重要。它们不是边缘优化，而是在帮你把“搬数据”和“算数据”组织成完整流水线。

7. WMMA、mma、wgmma 应该怎么理解

这一部分不用一开始就死背接口名，但要抓住代际演进方向。

  • WMMA：更偏高层、偏学习友好的 warp 级矩阵乘加接口
  • mma：更贴近底层指令路径，控制更细，性能空间更大
  • wgmma：更大协同粒度、更强吞吐，体现 Hopper 一代矩阵路径的进一步演进

真正需要建立的认识只有一个：

Tensor Core 编程不是静止不变的。架构升级后，最优矩阵路径也会变化。

所以你不能把某一代的优化经验机械搬到所有 GPU 上。

8. 为什么 Attention 也绕不开 GEMM 思维

标准 attention 的核心结构是：

1
2
3
S = Q x K^T
P = softmax(S)
O = P x V

其中至少包含两个重计算阶段：

  1. Q x K^T
  2. P x V

所以只要问题规模合适，attention 的大头计算也天然会靠近 GEMM 路线。

但它比单纯线性层更难优化，因为中间还夹着：

  • mask
  • scale
  • softmax
  • layout 变化
  • 中间结果管理
  • KV cache 相关访问

也就是说，attention 不是“两个干净的大 GEMM”，而是：

矩阵乘 + 若干访存敏感操作 + 很重的中间态管理。

9. 为什么 naive attention 会吃大亏

最直观但低效的 attention 做法是：

  1. 读 Q、K
  2. 计算 S = QK^T
  3. 把整个 S 写回显存
  4. 再读回 S
  5. 做 softmax
  6. 再写回 softmax 结果
  7. 再读回它和 V 相乘

这个流程最大的问题不是算得不对，而是 I/O 太重。

尤其在长序列下，中间 score 矩阵会很大。问题很快就从“乘法快不快”变成“中间结果能不能承受频繁物化和搬运”。

所以从 GPU 视角看，naive attention 的核心痛点通常是：

materialize 了太多中间结果，导致 HBM 流量炸掉。

10. FlashAttention 的核心价值是什么

FlashAttention 最值得学习的地方，不是它用了什么“新数学”，而是它首先是一个 I/O-aware 算法。

它真正关心的是：

  • 哪些中间结果必须写回 HBM
  • 哪些可以留在片上存储
  • 能不能按 tile 分块处理 Q/K/V
  • 能不能边算边做归一化，而不是先生成完整 score 矩阵

所以它的核心收益主要来自：

  • 减少 HBM 读写
  • 避免完整 materialize 大中间矩阵
  • 让 attention 的执行顺序更贴近 GPU 的内存层次

你在这部分不必一开始就背完整的 online softmax 推导，但一定要记住一句话：

FlashAttention 的本质不是“把 attention 公式算得更花”，而是“把 attention 重组得更符合 GPU 的内存系统”。

11. 从 GEMM 到 FlashAttention，本质都是计算-存储共同设计

很多人学 CUDA 时容易把优化理解成“写更复杂的 kernel”。

这不够准确。

更本质的说法是：

高性能 GPU 程序从来不是单纯的计算设计，而是计算模式和存储层次一起设计。

在 GEMM 里，这种思路表现为：

  • tiling
  • shared memory staging
  • register blocking
  • Tensor Core 路径
  • pipeline

在 Attention 里，这种思路表现为：

  • 避免 materialize
  • 按块处理 Q/K/V
  • 减少无意义的中间写回
  • 尽量融合后处理

对 AI infra / 推理 来说，这条思路非常重要，因为它决定了你看问题时不会只盯着 FLOPS，而会一起看带宽、缓存、shape 和数据流。

12. 为什么 Tensor Core 很重要，但又不是全部答案

Tensor Core 最擅长解决的是：

  • 大规模、规则、密集的矩阵乘加
  • 合适精度和布局下的高吞吐 dense compute

但它不会自动帮你解决这些问题：

  • KV cache 随机读写
  • 小 batch 或碎 shape 带来的低利用率
  • layernorm、softmax、elementwise 等访存敏感算子
  • 过多的 kernel launch
  • 中间结果物化导致的 HBM 压力
  • 多请求调度和系统层开销

所以面试里如果只会说“推理快是因为 Tensor Core 很快”，其实不够成熟。

更好的说法是：

Tensor Core 对 dense matmul 很关键，但推理系统整体性能还强烈依赖带宽、缓存、KV 管理、shape 组织和调度策略。

13. 为什么 prefill 和 decode 的硬件特征很不一样

这两个阶段的差异，是后面理解 vLLM、PagedAttention、TensorRT-LLM 的前提。

13.1 prefill 更像 dense compute

prefill 往往一次处理整段 prompt，很多计算更接近完整的大矩阵乘。

它的特点通常是：

  • 并行度高
  • 更容易形成大 shape
  • 更容易把 Tensor Core 吃满

所以它更接近 compute-heavy 的场景。

13.2 decode 更容易暴露 memory system 问题

decode 每次只生成少量 token，但要持续读取：

  • 模型权重
  • 历史 KV cache
  • 当前 token 相关状态

这会带来几个结果：

  • 单次有效计算粒度更小
  • 访存压力更突出
  • shape 更碎
  • 更难把 Tensor Core 吃满

所以 decode 常常不是“算不动”，而是“喂不动”。

14. 为什么量化会和推理硬件深度绑定

量化在推理里重要，不只是模型更小。

从 GPU 视角看，它至少有三层价值：

14.1 降低带宽压力

位宽更低，意味着：

  • 同样带宽能搬更多数据
  • 同样显存能放更多参数
  • cache 和片上存储的相对容量变大

14.2 更容易提高有效吞吐

如果硬件对低精度类型有更强支持，那么单位时间内能完成的有效运算也会更高。

14.3 更适合真实部署目标

在线推理不只看单请求延迟，还看：

  • 吞吐
  • 显存成本
  • 节点利用率
  • 部署成本

这也是为什么各种 serving 系统和推理引擎都会把量化放在很核心的位置。

15. 这一篇必须记住的几句话

  • AI 推理的大头，通常是 GEMM 和围绕它展开的数据搬运。
  • Tensor Core 是专用矩阵块乘加路径，不是“更快一点的 CUDA Core”。
  • 高性能 GEMM 的关键是 tiling、复用和流水线，不是朴素逐元素思维。
  • Attention 的难点不只是乘法，还包括 softmax、中间结果管理和 I/O。
  • FlashAttention 的价值，本质上是减少 HBM 读写，让 attention 更贴近 GPU 内存层次。
  • prefill 更像 dense compute，decode 更容易暴露带宽和 KV cache 问题。
  • 量化的价值不只在模型变小，还在于它更适配低精度高吞吐硬件路径。

16. 精简版面试表达

如果面试官问为什么 AI 推理如此依赖 GPU 上的 GEMM 和 Tensor Core，一个更完整的回答可以是：

因为 Transformer 里的线性层、attention 核心乘法、MLP 大头都可以归结到矩阵乘。GPU 又天然擅长规则、密集、可分块的吞吐型计算，所以高性能实现会围绕 GEMM 做多层 tiling、shared memory 和寄存器复用，并尽量走 Tensor Core 路径。但推理系统不只是 matmul，attention、KV cache、softmax 和 decode 阶段的 memory pressure 也很重，所以真正的优化必须把算力和内存系统一起看。

系列导航