跳转至

Transformer 的计算与数据流

Transformer 把序列建模的主计算组织为 Attention 与逐 Token MLP。对推理系统而言，最关键的不是只会画 Block，而是能从参数和序列长度推导 GEMM Shape、KV Cache 与 Prefill/Decode 的不同瓶颈。

Decoder-only Block

flowchart TB
    X[Input Hidden<br/>B × T × D] --> N1[Norm]
    N1 --> QKV[Q / K / V Projection]
    QKV --> ATT[Masked Self-Attention]
    ATT --> O[Output Projection]
    O --> ADD1[Residual Add]
    X --> ADD1
    ADD1 --> N2[Norm]
    N2 --> UP[Gate / Up Projection]
    UP --> ACT[Activation + Elementwise Gate]
    ACT --> DOWN[Down Projection]
    DOWN --> ADD2[Residual Add]
    ADD1 --> ADD2
    ADD2 --> Y[Output Hidden<br/>B × T × D]

不同模型会采用 Pre-Norm/Post-Norm、LayerNorm/RMSNorm、GELU/SwiGLU、不同位置编码和 Bias 设计，但硬件分析仍可归入四类：

  • 大矩阵乘：QKV/O Projection 与 MLP；
  • Attention 中的 QKᵀ 与 PV；
  • Reduction：Norm、Softmax；
  • Elementwise：Residual、Activation、RoPE 等。

Attention 的 Shape

设 Batch 为 \(B\)，序列长度为 \(T\)，隐藏维为 \(D\)，Query Head 数为 \(H_q\)，KV Head 数为 \(H_{kv}\)，Head Dim 为 \(D_h\)。

\[ Q \in \mathbb{R}^{B \times H_q \times T \times D_h},\quad K,V \in \mathbb{R}^{B \times H_{kv} \times T \times D_h} \]

Scaled Dot-Product Attention：

\[ \mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{D_h}} + \mathrm{Mask}\right)V \]

标准实现若显式物化 \(T \times T\) Score/Probability，会产生大量 HBM 读写。FlashAttention 的核心不是近似 Attention，而是用 Tiling 在片上 SRAM 中完成分块 Softmax 与累加，减少 HBM 往返。

MHA、GQA 与 MQA

结构 KV Head 数 KV Cache/读取量 特点
MHA \(H_{kv}=H_q\) 最大 每个 Query Head 有独立 K/V Head
GQA \(1 < H_{kv} < H_q\) 按 Head 比例下降 多个 Query Head 共享一组 K/V Head
MQA \(H_{kv}=1\) 最小 所有 Query Head 共享 K/V，质量与实现需验证

忽略 Padding、对齐和管理元数据时，单序列 KV Cache 近似为：

\[ M_{KV} \approx 2 \times L \times T \times H_{kv} \times D_h \times bytes \]

其中 2 表示 K 与 V，\(L\) 为层数。Batch/并发增加时再乘以活跃序列数。这个公式直接连接模型配置与显存容量、Decode 带宽。

Prefill 与 Decode 是两种 Workload

一次处理 Prompt 的多个 Token，矩阵的 M = B × T 较大，GEMM 与 Attention 更容易形成较大 Tile。长序列时 Attention 中间量与 HBM 流量突出，通常更偏计算密集或 IO-aware kernel 优化。

每个序列每步只生成一个新 Token，Projection/MLP 常表现为较小 M 的 GEMM/GEMV；每层还要读取历史 KV Cache。算术强度下降，权重与 KV 带宽、Kernel Launch 和调度开销更突出。

flowchart LR
    P[Prefill<br/>T 个输入 Token] --> K[建立每层 KV Cache]
    K --> D1[Decode Step 1<br/>1 个新 Token]
    D1 --> D2[Decode Step 2]
    D2 --> DN[Decode Step N]
    D1 -. 追加 K/V .-> K
    D2 -. 追加 K/V .-> K
    DN -. 追加 K/V .-> K

不要用 Prefill 的 TFLOPS 推断 Decode

两个阶段的矩阵 Shape、复用、并发和数据移动完全不同。必须分别记录 Prompt/Output 长度、Batch/并发、TTFT、每 Token 延迟和带宽指标。

MLP 为什么同样重要

以两层 FFN 为例，若中间维为 \(I\)：

\[ X_{B T \times D} W_{up,D \times I} \rightarrow H_{B T \times I} \]
\[ H_{B T \times I} W_{down,I \times D} \rightarrow Y_{B T \times D} \]

很多 LLM 的 \(I\) 是 \(D\) 的数倍，MLP 因此占据大量参数与 FLOPs。SwiGLU 类结构还包含 Gate/Up 两个投影。Decode 时，MLP 权重需要逐层被读取，权重量化和融合能否降低真实 HBM 流量尤其关键。

算子融合的边界

常见候选：

  • QKV Projection 合并；
  • Bias/Activation/Elementwise 融合；
  • Residual + Norm 融合；
  • RoPE 与 Q/K 变换融合；
  • Attention 的 Score、Mask、Softmax 与 PV 融合；
  • Quantize/Dequantize 与 GEMM Epilogue 融合。

判断融合收益时同时检查：

  1. 是否消除了 HBM 中间量或 Kernel Launch；
  2. 是否增加寄存器/Shared Memory 并降低 occupancy；
  3. 动态 Shape、Mask 与数据类型是否仍被支持；
  4. 数值顺序变化是否影响精度；
  5. 是否失去成熟库的更优矩阵实现。

Transformer 的最小性能账本

项目 Prefill Decode
Shape Batch、Prompt 长度 活跃序列数、当前上下文长度
主要计算 大 GEMM、长序列 Attention 小 M GEMM/GEMV、逐步 Attention
主要常驻数据 权重、临时 Activation 权重、KV Cache
主要移动 Attention 中间量、权重 每步权重与历史 KV 读取
主要优化 FlashAttention、融合、Tensor Core Batching、量化、KV 精度/布局、融合
关键指标 TTFT、Prefill tokens/s TPOT/ITL、Decode tokens/s

延伸阅读