在 Advanced Attention 章节过后，我们将视角放在全局的 LLM forward 计算过程中：

  • Prefill 阶段：模型处理输入 Prompt，计算 QK 矩阵，存储 KV-Cache
    • 在这个阶段对于 Full-Attention 来说，需要完成多个 o(D2)o(D^2) 的大矩阵乘法来计算 Attention Score
    • 因此是计算密集型的操作
  • Decode 阶段：模型开始实现自回归解码

在本章节，我们将重点放在两个部分上：

  • FLOPs 分析: 在 Transformer 的每一个步骤上，模型花费的计算复杂度是多少? 哪一个地方是最为耗时的步骤?
  • Decoding Methods: 在 Attention 章节我们详细介绍了若干 efficient attention 的方法，主要是在 prefiling 阶段（当然也有在 decoding 阶段）进行模型优化，在这里，我们详细介绍两个比较主流的 efficient decoding 算法, speculative decoding 和 mutli-token-predictions.

LLM Training FLOPs

Forward Pass FLOPs

我们首先从最简单的 Attention & LLM 出发，计算单次 Attention 计算中，不同组件的计算复杂度耗时。

Basic: Counting Foundations

一般来说，一次浮点运算（加减乘除）被定义为 1 FLOP，我们来看常见的两种数学计算:

  • 向量的点积计算 Dot Products for u,v∈Vu,v \in \mathcal{V}:
    • 向量的维度为 [d,1][d,1], 则需要进行 dd 次乘法和 d−1d-1 次加法，得到一个标量值
    • FLOPS: 2d−1≈2d2d - 1 \approx 2d
  • 矩阵乘法计算 Matrix Multiplications (GEMM):
    • 不考虑矩阵乘法在计算复杂度上的优化，只考虑最简单的 GEMM 朴素矩阵乘法
    • 对于 A([p,q]×B([q,s])=C([p,s]))A([p,q] \times B([q,s]) = C([p,s])),. 计算复杂度是 ps(2q−1)≈2pqsps (2q-1) \approx 2pqs

Multi-Head Attention Compute

不考虑 Batch-size, 我们来看 Attention 计算部分，发生了什么矩阵计算:

  • Input: X∈[L,d]X \in [L , d]
  • Q, K, V Projection:
    • Q=XWQ∈[L,d]Q = X W_Q \in [L, d], WQ∈[d,d]W_Q \in [d, d]
    • K=XWK∈[L,d]K = X W_K \in [L, d], WK∈[d,d]W_K \in [d, d]
    • V=XWV∈[L,d]V = X W_V \in [L, d], WV∈[d,d]W_V \in [d, d]
    • Compute 计算: 3×Ld(2d−1)=6Ld2−3Ld3 \times Ld(2d-1) = 6Ld^2 - 3Ld
    • 对于 MHA 来说，每一个头投影的矩阵是 Qi=XWQ,i∈[L,dh]Q_i = X W_{Q, i} \in [L, d_h]
      • 最终，每一个 head 的 output [L,dh][L, d_h] 拼接成 [L,d][L, d] 进入 FFN Layer
      • Compute 计算: 3×h×(Ldh(2d−1))=6Ld2−3Ld3 \times h \times (L d_h (2d-1)) = 6Ld^2 - 3Ld
  • Atten Scores Compute
    • 对于每一个头，模型会输出对应的 Attention Score 矩阵 [L,L][L, L], 具体是 QKTQ K^T 这两个矩阵相乘，Compute: h×L2(2dh−1)=2L2d−L2hh \times L^2(2d_h - 1) = 2L^2 d - L^2h
  • Scaling: 除以 dk\sqrt{d_k} 的操作:
    • Compute: h×L2×1=hL2h \times L^2 \times 1 = h L^2
  • Stable Softmax:
    • 因为 Softmax 存在 e 指数幂的运算，数值稳定性在数值较大处往往会溢出，因此实际工程上往往采用 Stable Softmax 的运算技巧，找到最大的数，然后分子分母同时除以 eamax⁡e^{a_{\max}}（数学上等价），保证数值稳定性
    • Compute: 对于一次 softmax 运算，需要有 LL 次减法, LL 次除法, L−1L-1 次加法: 3L−13L-1
      • 注意，这里没有只考虑的四则运算，如果只考虑时间复杂度，这里的 FLOPS 应该是 kLk L
      • Softmax 是针对每一个 query 对它全部可见的 key 去做的，在不考虑掩码的情况下，应该是 h×L×(3L−1)h \times L \times (3L-1)
  • Weighted Values & Output Projections
    • Oi=Ai×Vi∈[L,dh]O_i = A_i \times V_i \in [L, d_h]
    • Compute: h×Ldh(2L−1)=2L2d−Ldh \times L d_h (2L-1) = 2L^2 d - Ld
    • 最终 O∈[L,d]O \in [L,d], 过一个 Wo∈[d,d]W_o \in [d, d] Compute: Ld(2d−1)=2Ld2−LdL d (2d -1) = 2Ld^2 - Ld
ComponentExact FLOPsLeading FLOPs
Q, K, V Projection6Ld2−3Ld6Ld^2-3Ld6Ld26Ld^2
Attention Scores2L2d−hL22L^2d-hL^22L2d2L^2d
ScalinghL2hL^2hL2hL^2
Stable Softmax3hL2−hL3hL^2-hL3hL23hL^2
Weighted Values2L2d−Ld2L^2d-Ld2L2d2L^2d
Output Projection2Ld2−Ld2Ld^2-Ld2Ld22Ld^2
Total MHA8Ld2+4L2d+3hL2−5Ld−hL8Ld^2+4L^2d+3hL^2-5Ld-hL8Ld2+4L2d8Ld^2+4L^2d
在实际 Casual Attention 的实践中，相当大一部分的计算可以被减半(full-size 的矩阵变成一个下三角矩阵):
  • 计算 attention 矩阵 -> 下三角矩阵
  • scaling: 下三角矩阵
  • softmax: 下三角矩阵
  • values: 下三角 矩阵

替换只需要 L2L^2 -> L2+L2\frac{L^2 + L}{2}

FFN, Normalizations, Output Costs

标准 FFN 可以写为 ([L,d][L, d] -> [L,d][L, d])

FFN⁡(X)=ReLU⁡(XW1)W2 \operatorname{FFN}(X)=\operatorname{ReLU}(XW_1)W_2

对于升维矩阵: Compute Lrd(2d−1)Lrd(2d - 1) (假设升维维度是 d -> rd) 对于降维矩阵: Compute Ld(2rd−1)Ld(2rd - 1) 不同的激活函数对应的 compute 不一样（比如 ReLU 0 FLOPs，SiLU 和 GELU 等引入较多的算术运算）

对于 一个标准的 Pre-RMSNorm 阶段，norm 会发生在:

H(l)=X(l)+MHA(RMSNorm(X(l))) H^{(l)} = X^{(l)} + \text{MHA}(\text{RMSNorm}(X^{(l)}))

X(l+1)=H(l)+FFN(RMSNorm(H(l))) X^{(l+1)} = H^{(l)} + \text{FFN}(\text{RMSNorm}(H^{(l)}))

每一次 Norm 都是对 [L,d][L, d] 做的，对最后一个维度做归一化。 对于一个单个 dd 维的向量，计算 RMSNorm 需要 4d+24d+2 次计算(先计算 mean square，再计算归一化后的值) 因此 计算量: L(4d+2)×(2n+1)L(4d+2) \times (2n+1) (在所有 blocks 的最后，还要加一个 RMSNorm)

LayerNorm 的计算会更加复杂

tokenize 的过程是一个 lookup 的查表过程，因此我们认为是一个 memory-bound 的过程，FLOPs 几乎为 0

What is leading?

  • Four Projections: 8nLd28 n L d^2
  • FFN: 4rnLd24 r n L d^2
  • Score / Value Products (Attention Score 矩阵): 4nL2d4 n L^2 d
  • Vocab Size: 2LdV2 L d V

Architectural Variants

Gated FFNs

在 Gated FFNs 中，有三个矩阵计算和一个 SiLU 激活函数:

  • 三个矩阵计算，假设中间维度是 rdr d, 则 leadingflops 就是 6rLd26 r L d^2
  • SiLU 需要 4 次算术运算和一次指数运算

GQA

假设 hidden dim 维度是 dd, query width: dd, key/value width: gdh\frac{g d}{h}

  • Q Projection: 本质还是多头注意力，FLOPS 没有变化
  • K Projection: [L,d]×[d,gd/h][L,d] \times [d, gd/h]: compute: 2Lgd2h\frac{2Lgd^2}{h}
    • MHA 中，这一项是 2Ld22Ld^2，因为 g<hg < h, 所以这一项低于 MHA 的计算量!
  • V Projection 同理

对于 Pair FLOPs，本质上每个 query 还是拿到一份 K 矩阵和 V 矩阵（只是存在复用，有 projection），因此对于 attention 计算的部分，计算量不存在变化。

GQA 的优质主要在于显存的高效利用 && KV-Cache 的节省，从 FLOPs 就是 projection 部分的 flops 被节省了。

RoPE

与传统 Position Embedding 不同，RoPE 不直接与输入 XX 相加，而是在 Q、K Projection 之后，对 Q 和 K 的特征进行旋转。我们首先来回顾下 ROPE 是如何 apply 一个矩阵乘法:

θi,t=t⋅ωi=t⋅10000−2(i−1)d \theta_{i, t} = t \cdot \omega_i = t \cdot 10000^{-\frac{2(i-1)}{d}}

R(θi,t)=(cos⁡θi,t−sin⁡θi,tsin⁡θi,tcos⁡θi,t) R(\theta_{i, t}) = \begin{pmatrix} \cos\theta_{i, t} & -\sin\theta_{i, t} \\ \sin\theta_{i, t} & \cos\theta_{i, t} \end{pmatrix} Rt=(R(θ1,t)0⋯00R(θ2,t)⋯0⋮⋮⋱⋮00⋯R(θd/2,t)) R_t = \begin{pmatrix} R(\theta_{1, t}) & 0 & \cdots & 0 \\ 0 & R(\theta_{2, t}) & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & R(\theta_{d/2, t}) \end{pmatrix}

对于每一对维度，rope 需要做 (本质是一个 [2,2]×[2,1]→[2,1][2,2]\times [2,1] \to [2,1] 的矩阵乘法运算)：

  • 4 次乘法
  • 2 次加减法

因此，对于 hh 个 head, dh/2d_h / 2 个 pair，每个 pair 做 6 个 FLOPs, total compute:

h×(dh/2)×2×6×L=6Ld h \times (d_h / 2) \times 2 \times 6 \times L = 6Ld

这里没有具体统计 θi,t\theta_{i,t} 的计算以及相关三角函数和指数函数的计算过程。

RoPE during MHA
  • 每一个 token 都做 RoPE (LL 个 token)
  • 每一个 token 需要对 Q 和 K 矩阵做 RoPE (×2\times 2)
  • 两两组成一个维度 dh/2d_h / 2
  • 多个头 dd
  • 每一对维度FLOPs: 66

因此，乘起来就是 6Ld6Ld

RoPE during GQA

GQA 的优化在于 K 的数量是 gg 个而不是 hh 个。因此 FLOPs: 3L(h+g)dh=3(1+gh)d3L(h + g)d_h = 3(1 + \frac{g}{h})d

Speculative Decoding

Multi-Token Prediction

DeepSeek: https://arxiv.org/pdf/2412.19437