把基础打牢

§1.1Transformer 一口气速通

我们要加速的"宿主"是一个 transformer decoder（GPT 系列、Llama 系列、Qwen 系列 全都一样）。一个 transformer block 长这样——把它背下来，后面所有 trick 本质上都是改这八步：

# x: [B, S, D]   D = hidden dim, S = 当前序列长度
# 一层 transformer block 的前向，省略 dropout / 残差缩放
def block(x, theta):
    h = rmsnorm(x, theta.norm1)
    q = h @ theta.Wq               # [B, S, D]
    k = h @ theta.Wk               # [B, S, D]
    v = h @ theta.Wv               # [B, S, D]
    # reshape into heads: [B, n_heads, S, head_dim]
    q, k, v = split_heads(q, k, v)
    # 加 rotary / ALiBi 等位置编码
    q, k = apply_rope(q, k)
    # 因果 attention
    attn = softmax(q @ k.transpose(-1,-2) / sqrt(head_dim) + causal_mask) @ v
    attn = merge_heads(attn)
    y = x + attn @ theta.Wo        # 残差
    z = y + ffn(rmsnorm(y, theta.norm2), theta.W_up, theta.W_down)
    return z

生成阶段把 $L$ 层 block 像穿糖葫芦一样串起来，最后一层接一个 x @ theta.W_emb.T 得到 vocab 上的 logits，再 sample 一个新 token， 追加到输入末尾，循环——这就是自回归。

几个关键尺寸（以 Llama-3-70B 为例）

形状直觉 · 一个 token 在网络里走一圈要碰几次矩阵

把一个新 token 的 hidden state $h \in \mathbb{R}^D$ 经过一层 block，需要：

• 3 次 $D \to D$ 的投影：$W_q, W_k, W_v$；
• 1 次 $D \to D$ 的输出投影：$W_o$；
• 2 次 attention matmul：$qK^\top$ 和 $\text{softmax}(...)V$；
• FFN：1 次 $D \to d_{ff}$、1 次 $d_{ff} \to D$（SwiGLU 实际是 3 次）。

乘上 $L = 80$ 层，每个 token 大约要碰 $\;80 \times (4 \cdot D^2 + 2 \cdot D \cdot d_{ff}) \approx 70\text{ GFLOP}$。 在 H100 上算力 989 TFLOPS，理论 14 µs / token—— 但实际 decode 每 token 要 ~5 ms。差了 350×。 这 350× 跑哪去了？§1.2–§1.3 揭晓。

§1.2推理的两阶段 · Prefill vs Decode

把推理拆开看，自回归 LLM 有两个截然不同的工作模式：

# Prefill: 一次把整段 prompt 喂进网络
def prefill(prompt_ids, model):
    x = embed(prompt_ids)                         # [B, S, D]
    for layer in model.layers:
        x, kv = layer(x, kv_cache=None)
        save_to_kv_cache(layer.id, kv)            # 把 K/V 存起来
    return x[:, -1]                               # 只用最后一个位置算 logits

# Decode: 一次只喂一个 token，复用历史 KV
def decode_step(last_token_id, model):
    x = embed(last_token_id)                      # [B, 1, D]
    for layer in model.layers:
        kv_past = load_from_kv_cache(layer.id)    # 把全部历史 KV 读回来
        x, kv_new = layer(x, kv_past=kv_past)
        append_to_kv_cache(layer.id, kv_new)      # 多挂一个 token 的 K/V
    return sample(x[:, -1] @ model.W_emb.T)

把这一段刻进脑子里："decode 阶段的延迟 $\approx$ 把整个模型权重从 HBM 读一次的时间。" 因此 7B 模型在 H100 上每个 token 的下限大约是 14 GB / 3 TB/s ≈ 4.7 ms——对应 ~210 tok/s。 所有"per-token 量化"、"speculative decoding"、"MoE per-token 稀疏激活"，目的都只有一个—— 减少"每个 token 必须从 HBM 拖进来的字节数"。

§1.3Roofline · 哪里慢，一张图说清

Williams & Patterson 在 2009 年 CACM 上画过这张图—— Roofline。横轴是 arithmetic intensity $I = \dfrac{\text{FLOPs}}{\text{Bytes from memory}}$， 纵轴是实际可达到的 FLOP/s。 你的内核在这张图上一定坐落于两条线之下：

两条线相交的那个 $I^*$ 称为ridge point。Hopper FP16 上 $I^* \approx 330$；FP8 上 $I^* \approx 660$； Blackwell FP4 上 $I^* \approx 1300$。 这意味着算力越大、ridge 越高、越多的内核会"掉"到斜坡那一侧（memory-bound）。

怎么手算一个内核的 I

Decode 的 I 永远在个位数附近，无论怎么调；这是为什么 vLLM 拼命做 continuous batching——多 batch 等于把同一份权重分摊给更多 token， 分子放大、分母不变。但 batch 越大，KV cache 越占显存，这条路有上限。

§1.4显存账单 · 谁吃了我的 80GB

把"显存里到底装了什么"列清楚，是所有后续优化的对账单：

推理时

训练时再加上

# 心算工具：估一个 dense LLM 的训推显存
def estimate_mem(params_B, L, D, H_q, H_kv, d_h, S, B, bytes_per=2):
    # 1) Weights (推理也要)
    W = params_B * 1e9 * bytes_per
    # 2) KV cache per sequence
    KV = 2 * H_kv * d_h * bytes_per * L * B * S
    # 3) Adam state (训练; fp32 m,v + fp16 master copy)
    Adam = params_B * 1e9 * (4 + 4 + 2)
    # 4) Activations per fwd (训练; 用 FlashAttention 时 attn 部分≈0)
    Act_train = B * S * L * D * 2 * bytes_per
    print(f'weights:    {W/1e9:.1f} GB')
    print(f'KV (B={B}, S={S}): {KV/1e9:.1f} GB')
    print(f'Adam:       {Adam/1e9:.1f} GB  (训练)')
    print(f'Act:        {Act_train/1e9:.1f} GB  (训练)')

# Llama-3-70B 训练 B=1 S=8192:
# weights 140 GB · KV 0.6 GB · Adam 700 GB · Act 21 GB
estimate_mem(70, 80, 8192, 64, 8, 128, 8192, 1)