Transformer 解剖：从 Attention 到推理系统

第 18 章 Flash Attention 与分布式推理：内存层级与多卡协同

作者 杨艺韬 · 6,666 字 · 发布于 · 更新于

第六部分到这里讲了 KV Cache、量化、投机解码——但都是「在已有 attention 实现上做加法」的优化。这一章我们 改 attention 本身的实现——同样的数学，跑得快几倍——以及 把模型切到多张卡 的分布式策略。

这两件事看似无关，其实共享同一个底层逻辑：在内存层级上找最优点。Flash Attention 在「HBM ↔ SRAM」之间找最优；TP / PP / EP 在「单卡 ↔ 多卡 ↔ 多机」之间找最优。理解了一边，另一边就豁然开朗。

读完这章你能：

  • 画出 GPU 的内存层级（HBM、L2、SMEM、寄存器）和它们的容量 / 带宽；
  • 解释 Flash Attention 1/2/3 各代的核心创新；
  • 推导 online softmax 的数学正确性；
  • 设计 70B / 671B 模型的并行部署拓扑（什么时候 TP、什么时候加 PP、什么时候加 EP）；
  • 估算给定带宽下分布式推理的通信瓶颈。

18.1 GPU 内存层级：物理事实

先把基础事实摆出来。一张 H100 GPU 的内存层级：

flowchart TB
  REG["寄存器 (Registers)<br/>256 KB/SM × 132 SM ≈ 33 MB<br/>带宽最高（无官方公开值）<br/>SM 内私有"]
  SMEM["共享内存 SMEM (SRAM)<br/>~228 KB / SM × 132 SM ≈ 30 MB<br/>带宽 ~30 TB/s（全卡合计）<br/>SM 内共享"]
  L2["L2 Cache<br/>50 MB<br/>带宽 ~12 TB/s（实测值）<br/>所有 SM 共享"]
  HBM["HBM (Global Memory)<br/>80 GB<br/>带宽 3.35 TB/s<br/>整卡共享"]
  HBM --> L2 --> SMEM --> REG

关键事实：

  1. HBM 容量大、带宽相对小——80 GB，3.35 TB/s
  2. SMEM 容量小、带宽极大——单 SM 228 KB、约 230 GB/s，全卡 132 个 SM 合计约 30 TB/s（HBM3 的 9 倍）
  3. 越靠近计算单元越快越小——寄存器 > SMEM > L2 > HBM

L2 的 ~12 TB/s 不是 NVIDIA 官方数据：H100 Whitepaper 与 Hopper Tuning Guide 都没有公布 L2 带宽，这个数来自第三方 microbenchmark（Luo et al., 2024），随访问模式波动，只能当量级看。HBM 的 3.35 TB/s 才是官方规格。CUDA 专栏第 4 章用的是同一个来源、同一个数。

模型参数（140 GB Llama-70B）远超 HBM；即使是 KV Cache（几十 GB）也远超 SMEM。所以模型推理必须从 HBM 读数据，但能在 SMEM 里完成的计算尽量在 SMEM 里完成。

这就是 Flash Attention 的全部精髓。

18.2 朴素 attention 的内存访问账

回顾标准 attention 的实现（第 8 章）：

S = Q @ K.transpose(-2, -1) / sqrt(d)   # (N, N)
A = softmax(S, dim=-1)                   # (N, N)
out = A @ V                              # (N, d)

每一步在 GPU 上的数据流：

flowchart LR
  HBM1["HBM: Q, K, V"] --> COMP1["算 S = QK^T"]
  COMP1 --> HBM2["HBM 写 S 矩阵 N×N"]
  HBM2 --> COMP2["算 softmax(S)"]
  COMP2 --> HBM3["HBM 写 A 矩阵 N×N"]
  HBM3 --> COMP3["算 A @ V"]
  COMP3 --> HBM4["HBM 写 out 矩阵"]

每一步都要把数据写回 HBM，再从 HBM 读回来——HBM 访存量 Θ(Nd+N2)\Theta(Nd + N^2)，其中 N2N^2 那一项来自反复读写被物化的注意力矩阵。

N=8K 时，单层单个头一次 attention 光是把 S、A 两个矩阵各写一遍就是 2×N2×4(FP32 中间存)≈0.52 \times N^2 \times 4 \text{(FP32 中间存)} \approx 0.5 GB；乘上头数和层数（以 Llama-3 70B 的 64 个 Q 头、80 层计），一次 8K prefill 仅写入就约 2.7 TB，再算上读回，远超模型权重大小（140 GB）。

HBM 带宽 = attention 实际瓶颈。理论 GPU 算力够用，但因为内存访问太多，attention 跑不快。

18.3 Flash Attention 1：分块 + Online Softmax

Flash Attention（Tri Dao et al., NeurIPS 2022）的核心想法：不要把 N×N 矩阵物化到 HBM，把整个 attention 算在 SMEM 里面。

但有个问题：N×N 矩阵在 N=8K、FP32 下约 270 MB，远超 SMEM 的 30 MB——直接放不下。

解决方案是 tiling（分块）+ online softmax：

  1. 把 Q、K、V 切成小块（tile）
  2. 每次只把一块 K、V 从 HBM 读到 SMEM
  3. 在 SMEM 里对当前 Q tile 算这一小块的 attention 部分
  4. 通过 online softmax 把多个 K/V 块的部分结果正确合并

Online Softmax 的数学

直接说：能不能把 softmax 「拆成多个块单独算然后合并」？

朴素的 softmax 公式：

softmax(xi)=exi∑jexj\text{softmax}(x_i) = \frac{e^{x_i}}{\sum_j e^{x_j}}

需要先扫一遍找 max⁡\max（数值稳定），再扫一遍算分母 ∑exj−max⁡\sum e^{x_j - \max}，再扫一遍算每个值的输出——三遍扫描 + 中间存 N 个值。

Online softmax 把前两遍（求 max、求分母）合并成一遍。维护两个变量 mm（当前最大值）和 ℓ\ell（当前分母）：

初始 m0=−∞m_0 = -\infty, ℓ0=0\ell_0 = 0。

每读一个新值 xix_i：

mi=max⁡(mi−1,xi)m_i = \max(m_{i-1}, x_i) ℓi=ℓi−1⋅emi−1−mi+exi−mi\ell_i = \ell_{i-1} \cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}

最终的 softmax 输出可以从 m,ℓm, \ell 推出来。关键点：每读一块新数据时，老的 ℓ\ell 需要被「rescale」（乘 emold−mnewe^{m_{\text{old}} - m_{\text{new}}}）以适应新的 max。

把它扩展到 attention 输出（O=A⋅VO = A \cdot V）：

每遇到一块新的 K(j),V(j)K^{(j)}, V^{(j)}：

  1. 算 S(j)=Q⋅K(j)TS^{(j)} = Q \cdot K^{(j)T}（在 SMEM 里）
  2. 找新 block 的 max：m(j)=max⁡lS:,l(j)m^{(j)} = \max_l S^{(j)}_{:,l}
  3. 更新全局 max：mnew=max⁡(mold,m(j))m_{\text{new}} = \max(m_{\text{old}}, m^{(j)})
  4. rescale 老的 output 累加器：O←O⋅emold−mnewO \leftarrow O \cdot e^{m_{\text{old}} - m_{\text{new}}}
  5. 算新 block 的部分 output 累加：O←O+eS(j)−mnew⋅V(j)O \leftarrow O + e^{S^{(j)} - m_{\text{new}}} \cdot V^{(j)}
  6. 更新分母：ℓnew=ℓold⋅emold−mnew+∑eS(j)−mnew\ell_{\text{new}} = \ell_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + \sum e^{S^{(j)} - m_{\text{new}}}

最后一步：O←O/ℓnewO \leftarrow O / \ell_{\text{new}}

数学上完全等价于一次性计算的标准 softmax + AV——但内存使用从 O(N²) 降到 O(N)（只需要存当前块）。

Flash Attention 1 的算法流程

flowchart LR
  subgraph OUTER["外层循环 over K, V blocks"]
    LOAD[Load 一块 K^j, V^j 到 SMEM]
    subgraph INNER["内层循环 over Q blocks"]
      QLOAD[Load Q_i 及其 O_i / m_i / l_i] --> COMP[Compute S_ij, 更新 max / sum]
      COMP --> ACCUM[rescale 并累加到 O_i]
      ACCUM --> WB[写回 O_i / m_i / l_i 到 HBM]
    end
    LOAD --> QLOAD
  end
  WB --> DONE[全部 K/V 块处理完：各 O_i 除以 l_i]

伪代码（简化，循环顺序按 FA1 论文 Algorithm 1：外层 K/V、内层 Q，与上图一致； FA2 把这两层对调，见下一节的「优化 2」。另一处简化：这里把「除以 ℓ\ell」留到最后统一做，FA1 原文每处理一块都归一化一次——推迟归一化正是 FA2 的改进之一，见下一节「优化 1」；缩放因子 1/d1/\sqrt{d} 也省略了）：

# Q, K, V in HBM；O, m, l 也在 HBM，按 Q 块分片
# Q: (N, d), tile size B_r;  K, V: (N, d), tile size B_c
for j in range(N // B_c):                    # 外层: K, V 块
    K_j, V_j = load to SMEM
    for i in range(N // B_r):                # 内层: Q 块
        Q_i = load Q tile to SMEM
        O_i, m_i, l_i = load 该 Q 块的累加量      # 首轮为 zeros / -inf / 0
        S_ij = Q_i @ K_j.T                    # in SMEM
        m_new = max(m_i, S_ij.max(dim=-1))
        # rescale 老 output
        O_i = O_i * exp(m_i - m_new)
        # 累加新 block 贡献
        l_new = l_i * exp(m_i - m_new) + exp(S_ij - m_new).sum(dim=-1)
        O_i += exp(S_ij - m_new) @ V_j
        write back O_i, m_new, l_new           # 每个 Q 块的累加量要写回
# 全部 K/V 块处理完后，各 Q 块再统一除以 l_i

注意最后两行——正因为外层是 K/V，每个 Q 块的 O_i / m_i / l_i 在整趟循环里要反复 读出写回，这正是 FA2 要对调循环的动机。

Flash Attention 的收益

这里要把两个不同的量分清楚，很多材料把它们混为一谈：

  • 显存占用：确实从 O(N2)O(N^2) 降到了 O(N)O(N)——N×N 的注意力矩阵根本不落地，只留 mm、ℓ\ell 两个 O(N)O(N) 的统计量。这是长上下文（如 128K）训练和推理得以落地的重要前提之一。
  • HBM 访存量：没有降到 O(N)O(N)。K/V 会被每个 Q block 重读一遍，论文给出的 IO 复杂度是 Θ(N2d2/M)\Theta(N^2 d^2 / M)（MM 是 SRAM 容量），朴素实现是 Θ(Nd+N2)\Theta(Nd + N^2)。两者都随 N2N^2 增长，Flash Attention 省下的是一个 M/d2M/d^2 量级的常数因子——论文在 A100 上实测 HBM 访问最多少约 9 倍（原文 Fig. 2），即十倍量级。
上下文长度 朴素 HBM 访存量（相对 4K） Flash Attention（相对同长度朴素，量级示意）
4K 1× ~1/10
8K 4× ~1/10
32K 64× ~1/10
128K 1024× ~1/10

也就是说：Flash Attention 把整条曲线往下平移了一个十倍量级，但曲线的斜率没变。论文摘要里报的端到端加速是分任务的：BERT-large（序列 512）比 MLPerf 1.1 的训练记录快 15%、GPT-2（1K）3×、Long-Range Arena（1K–4K）2.4×——三组的基线各不相同，不宜横向比；一般而言序列越长收益越明显，因为长序列下 attention 在总时间里占比更大。它没有改变 attention 的算法，只是把内存访问模式从「物化到 HBM」改成「在 SMEM 里完成」。

18.4 Flash Attention 2：减少非矩阵乘开销

Flash Attention 2（Dao, 2023）在 V1 基础上做了几个细节优化：

优化 1：减少非矩阵乘 op

GPU 的 Tensor Core 算矩阵乘极快，其他操作（exp、scale、reciprocal）却慢得多。FA2 论文的例子是 A100：FP16/BF16 矩阵乘峰值 312 TFLOPs，非矩阵乘的 FP32 运算只有 19.5 TFLOPs（H100 SXM 的 BF16 稠密矩阵乘峰值是 989 TFLOPs）。FA1 里的 exp 和 rescale 操作占了不小比例。

FA2 通过重排算法（比如把逐块「除以分母 ℓ\ell」的归一化推迟到最后统一做一次，逐块的 emold−mnewe^{m_{\text{old}} - m_{\text{new}}} rescale 仍然要做），把更多时间花在矩阵乘上、把非矩阵乘的次数压到最少。

优化 2：循环顺序对调

FA1 的外层循环是 K/V，内层是 Q——每次进入内层都要 reload Q（K/V 在 SMEM 里复用，Q 多次加载）。

FA2 把循环对调：外层 Q、内层 K/V。Q 加载一次，K/V 块多次加载——更适合 GPU 的访问模式。

优化 3：work partition 优化

GPU 上每个 SM（Streaming Multiprocessor）独立工作。FA1 的 grid 只铺在 (batch × head) 两条轴上——一个 thread block 包办一个 head 的全部 attention。batch × head 小于 SM 数时（长上下文推理、小 batch 训练都会遇到），大量 SM 直接闲置。

循环顺序换成外层 Q 之后，每个 Q 块天然是一个独立的 thread block，grid 变成 (Q 块数 × head 数 × batch)——序列长度本身成了一条并行轴，长序列反而更容易填满 GPU。

综合下来，FA2 论文在 A100 上实测比 FA1 快约 2×。vLLM 在 CUDA 上默认的 attention 后端就是 FlashAttention：v0.8.5 在 Hopper 上默认 FA3、其余 GPU 默认 FA2（vllm/attention/utils/fa_utils.py 的 get_flash_attn_version）。

18.5 Flash Attention 3：Hopper 特化

Flash Attention 3（Shah et al., 2024）针对 Hopper 架构（H100）做了进一步特化。Hopper 引入了几项新硬件能力：

  1. TMA（Tensor Memory Accelerator）：异步内存搬运——计算和数据搬运可以重叠
  2. WGMMA（Warp Group Matrix Multiply Accumulate）：异步矩阵乘——SM 级别的 pipeline
  3. FP8 支持：原生 FP8 矩阵乘

FA3 用三种新技巧吃干净 Hopper 的能力：

技巧 1：异步流水线（async pipeline）

把「load 下一个 K/V 块」和「计算当前块」并行执行。当一个 warp 在算 GEMM 时，另一个 warp 在用 TMA 加载下一块——TMA 不占用 SM 计算资源，纯异步。论文还让 softmax 与 GEMM 重叠：两个 warpgroup 乒乓调度（pingpong），一个做 softmax 时另一个做 GEMM；同一 warpgroup 内也把相邻两块的 GEMM 与 softmax 流水起来。

技巧 2：Warp specialization

把 SM 内的 warp 分成两组：

  • Producer warps：专门负责加载数据（用 TMA）
  • Consumer warps：专门负责计算（用 WGMMA）

两组通过 SMEM 上的「pipeline」交换数据——produce 写入、consume 读取——形成软件级别的硬件 pipeline。

技巧 3：FP8 支持

FA3 额外提供一条 FP8 路径（不是默认，需要显式启用）。论文报告 FP8 版本能逼近 1.2 PFLOPS，且因为用了 block quantization + incoherent processing，数值误差比基线 FP8 attention 低 2.6 倍。

FA3 在 H100 上达到 740 TFLOPs（相对 FP16 理论峰值 989 TFLOPs，75% MFU——同一篇论文里 FA2 在 H100 FP16 上的基线是 35%）。这是单 attention kernel 能达到的工程极限之一。

flowchart LR
  V1["FA1（A100）<br/>Tiling + Online Softmax<br/>基线"] --> V2["FA2<br/>+ work partition<br/>+ 减 non-matmul<br/>A100 上约 2× FA1<br/>H100 FP16 上 ~35% MFU"]
  V2 --> V3["FA3（H100）<br/>+ async pipeline<br/>+ warp specialization<br/>+ FP8<br/>740 TFLOPs / ~75% MFU"]

图里的百分比必须连着硬件看：FA2 的 ~35% 和 FA3 的 ~75% 都是 H100 FP16 口径（Shah et al., FlashAttention-3, 2024，arXiv:2407.08608），而「FA2 比 FA1 快约 2×」是 FA2 论文在 A100 上的实测（§18.4）——跨硬件的数字不能直接相除。CUDA 专栏第 17 章用的是同一组口径。

18.6 现在转向分布式：为什么单卡不够

到这里我们解决了「单卡 attention 怎么跑得快」的问题。但回到现实：Llama-3 70B 是 140 GB，单 H100 的 80 GB 显存放不下。671B 模型按 BF16 计约 1.3 TB，要 17 张 H100 才装得下权重。

这就是分布式推理要解决的问题：怎么把超大模型切到多张 GPU 上。

主流的三种切分维度：

flowchart TB
  MODEL[一个大模型<br/>太大单卡放不下]
  MODEL --> TP[Tensor Parallel TP<br/>把大矩阵切成小矩阵 多卡协同]
  MODEL --> PP[Pipeline Parallel PP<br/>不同 GPU 跑不同层 流水线]
  MODEL --> EP[Expert Parallel EP<br/>MoE 专属 不同专家分到不同 GPU]

每种并行解决不同的问题，工程上常常叠加使用：8 张 GPU 跑 70B 用 TP=8；跑 671B MoE 则在同一批卡上叠 TP + EP（EP 从数据并行那一维切出来，18.10 节会讲清楚这个关系）。下面分别讲。

18.7 Tensor Parallelism (TP)：切大矩阵

TP 的核心想法：把大矩阵乘按某个维度切到多 GPU，每卡算一部分，然后做 all-reduce 合并。

举例：FFN 第一层 W1∈Rd×4dW_1 \in \mathbb{R}^{d \times 4d}（典型 d=8192，4d=32768）。把 W1W_1 沿列切 4 份：

W_1 = [W_1^{(1)} | W_1^{(2)} | W_1^{(3)} | W_1^{(4)}]

GPU 1 持有 W1(1)∈Rd×dW_1^{(1)} \in \mathbb{R}^{d \times d}，以此类推。

输入 x∈Rdx \in \mathbb{R}^d（每张 GPU 都有一份）：

  • 每卡算 yi=x⋅W1(i)y_i = x \cdot W_1^{(i)} 得到 yi∈Rdy_i \in \mathbb{R}^d
  • 拼起来（concat）得到 y∈R4dy \in \mathbb{R}^{4d}（被切到 4 张卡上）

然后 W2∈R4d×dW_2 \in \mathbb{R}^{4d \times d} 沿行切 4 份。每卡算它那部分：

  • 每卡算 zi=yi⋅W2(i)∈Rdz_i = y_i \cdot W_2^{(i)} \in \mathbb{R}^d
  • 4 张卡的 ziz_i 加起来（all-reduce sum）得到完整的 z∈Rdz \in \mathbb{R}^d

这样一对 GEMM（x⋅W1⋅W2x \cdot W_1 \cdot W_2）被切到 4 张卡，每卡算量减 4 倍，最后只需要一次 all-reduce 通信。

flowchart LR
  subgraph "GPU 1"
    G1[x] --> G1_1["× W_1^(1)"] --> G1_2["× W_2^(1)"] --> Z1[z_1]
  end
  subgraph "GPU 2"
    G2[x] --> G2_1["× W_1^(2)"] --> G2_2["× W_2^(2)"] --> Z2[z_2]
  end
  subgraph "GPU 3"
    G3[x] --> G3_1["× W_1^(3)"] --> G3_2["× W_2^(3)"] --> Z3[z_3]
  end
  subgraph "GPU 4"
    G4[x] --> G4_1["× W_1^(4)"] --> G4_2["× W_2^(4)"] --> Z4[z_4]
  end
  Z1 & Z2 & Z3 & Z4 --> AR[all-reduce sum]
  AR --> Z[完整 z]

TP 在 Attention 上

Multi-Head Attention 用 TP 也很自然：按 head 切。h 个头切到 N 卡，每卡算 h/N 个头：

  • WQ,WK,WVW_Q, W_K, W_V 按列切（h 个头独立计算）
  • 每卡独立算 attention（不需要通信）
  • WOW_O 按行切，结果 all-reduce sum

GQA 下 KV head 数比 Q head 数少，TP 维度最好能整除 KV head 数；TP 大于 KV head 数时，vLLM 的做法是把每个 KV head 复制到 TP/KV 张卡上各存一份（vllm/model_executor/layers/linear.py 中 QKVParallelLinear 的 num_kv_head_replicas），靠冗余存储而不是通信。

TP 的通信开销

每个 Block 需要 2 次 all-reduce（attention 后一次、FFN 后一次），每次传输的 hidden state 是 (B,T,d)(B, T, d)。Decode 时这个张量小得惊人——batch=1、d=8192、BF16 只有 16 KB，在 NVLink 域内完全是延迟主导，单次在微秒级。但要乘以 2×L2 \times L：80 层就是 160 次，加起来仍然远小于一个 decode step 的时长——TP 在 NVLink 域内接近免费（与《vLLM 推理内核深度解析》第 14 章的结论一致；想量化自己的机器，可用 nccl-tests 的 all_reduce_perf 测 16 KB 的延迟，再乘 2×L2 \times L）。

TP 在同机内（NVLink）下高效——一台 8 GPU 服务器内可以做 TP=8。但跨机就不行了——InfiniBand NDR 400 Gbps（单向 ≈50 GB/s）按同一口径（单向）只有 NVLink（单向 450 GB/s）的 1/9，而且跨机链路的单跳延迟比 NVLink 高一个量级；那个「乘以 160」的放大结构会把这点延迟放成决定性开销。

18.8 Pipeline Parallelism (PP)：切层

如果模型大到 8 卡 NVLink 也装不下（比如 DeepSeek-V3 的 671B），需要跨机扩展——PP 是这个场景的解决方案。

PP 的想法：把不同层放到不同 GPU。比如 80 层的 Llama-3 70B，前 20 层放 GPU 1，21-40 层放 GPU 2，41-60 层放 GPU 3，61-80 层放 GPU 4。

数据流：token 进入 GPU 1，经过 20 层，把结果传给 GPU 2，再经过 20 层，依此类推。

flowchart LR
  X[输入 token] --> GPU1[GPU 1<br/>Layers 1-20]
  GPU1 -.传递 hidden state.-> GPU2[GPU 2<br/>Layers 21-40]
  GPU2 -.传递.-> GPU3[GPU 3<br/>Layers 41-60]
  GPU3 -.传递.-> GPU4[GPU 4<br/>Layers 61-80]
  GPU4 --> OUT[输出 logits]

Pipeline Bubble

PP 看似简单但有个根本问题：在 GPU 1 处理 batch i 时，GPU 2、3、4 都在等——大部分时间空转。

解决方案：micro-batching——把一个大 batch 切成几个小 micro-batch，让流水线流起来：

flowchart TB
  T1[t=1: GPU1 处理 mb1]
  T2[t=2: GPU1 处理 mb2; GPU2 处理 mb1]
  T3[t=3: GPU1 处理 mb3; GPU2 处理 mb2; GPU3 处理 mb1]
  T4[t=4: GPU1 处理 mb4; GPU2 处理 mb3; GPU3 处理 mb2; GPU4 处理 mb1]
  T5[t=5: GPU1 idle; GPU2 处理 mb4; GPU3 处理 mb3; GPU4 处理 mb2]
  T6[t=6: GPU3 处理 mb4; GPU4 处理 mb3]
  T7[t=7: GPU4 处理 mb4]
  T1 --> T2 --> T3 --> T4 --> T5 --> T6 --> T7

可以看到流水线开头和结尾都有 idle GPU——这叫 pipeline bubble。bubble 比例 = (Npp−1)/Nmicro(N_{\text{pp}} - 1) / N_{\text{micro}}。要减小 bubble 必须增大 micro-batch 数——但这又增加了内存压力。

PP 适合场景

PP 的特点：

  • 通信少：每两个 PP 阶段之间只传一次 hidden state，数据量小
  • 跨机友好：可以跨 NVLink 域、跨数据中心
  • bubble 限制效率：在线低延迟场景下利用率低

适合：

  • 训练（大 batch、bubble 摊薄）
  • 离线推理大批次
  • 跨机扩展超大模型

不适合：

  • 在线低延迟推理（单 batch 时 bubble 严重）

18.9 Expert Parallelism (EP)：MoE 专属

EP 是 MoE 模型专门的并行方式：不同 GPU 持有不同的专家，token 按路由结果发到对应 GPU。

第 12 章我们讲过 MoE 的 All-to-All 通信——EP 的通用做法就是它的具体实现（注意 vLLM v0.8.5 的 EP 并不调 all_to_all，而是每个 MoE 层一次 AllReduce，见《vLLM 推理内核深度解析》第 14 章 §14.4.3）。

DeepSeek-V3 的部署示例（技术报告中预填充阶段用 EP32）：256 个路由专家分到 32 张卡，每张卡 8 个专家。一个 token 的 K=8 路由可能选了「分布在不同 GPU 上的」8 个专家——All-to-All 把 token 发到目标 GPU、计算完后 All-to-All 收回来。

flowchart LR
  GPU1[GPU 1<br/>专家 1-8] -.All-to-All.-> GPU2[GPU 2<br/>专家 9-16]
  GPU2 -.-> GPU3[GPU 3<br/>专家 17-24]
  GPU3 -.-> GPU32[...GPU 32<br/>专家 249-256]
  GPU1 -.-> GPU32

EP 的通信代价

All-to-All 是「N×N 通信」——N 张 GPU 同时给所有其他 GPU 发数据。带宽要求极高。

对 DeepSeek-V3 这种 256 专家 + 32 EP 的部署，每个 MoE 层要做两次 All-to-All（dispatch 把 token 发到目标专家、combine 把结果收回来）——V3 有 58 个 MoE 层，一步 decode 就是 116 次 All-to-All。

硬件这一侧的落差是关键：同机 NVLink 在 H100 上是单卡双向合计 900 GB/s（单向 450 GB/s），跨机 InfiniBand NDR 是单向 400 Gbps ≈ 50 GB/s——同口径差 9 倍。所以专家摆在同机内还是跨机，通信开销能差一个数量级，这也是「EP 组尽量不跨 NVLink 域」这条经验的来源。

DeepSeek 为此开源了 DeepEP——专门针对 MoE dispatch/combine 的 All-to-All 通信库；训练侧则用 DualPipe 把通信和计算重叠，把通信 bubble 压到很小。

18.10 三种并行的组合：3D 并行

实际工程几乎从不只用一种并行——TP、PP 再加上数据并行（DP），就是常说的「3D 并行」（这个词原指 DP × TP × PP）；MoE 模型还要叠上 EP，下面会看到 EP 并不是再乘上去的一维。

举例：一个 671B MoE 模型摊在 256 张卡上（示意拓扑，不是任何一家的官方配置）：

  • TP=8（同机 8 卡）：attention 和每个专家内部的矩阵都按 TP=8 切
  • PP=4：61 层按 4 段流水线拆，每段约 15 层
  • DP=8：8 路数据并行
  • EP=8：MoE 层把上面那 8 路 DP 重新解释成 EP=8，每组 32 个专家

总卡数 = TP × PP × DP = 8 × 4 × 8 = 256。

这里有个容易搞错的点：EP 不是在 TP/PP/DP 之外再乘的第四个维度。世界大小始终是 TP × PP × DP；MoE 层只是把 DP（有时连同 TP）这几路 GPU 换一种切法——同一批卡，非 MoE 层做数据并行、MoE 层做专家并行。把四个数直接相乘去算卡数，会算出一个根本不存在的集群。

flowchart TB
  M[671B MoE 模型] --> TP[TP=8 切 attention 和 FFN 矩阵]
  TP --> PP[PP=4 切 61 层]
  PP --> EP[EP=8 切 256 个专家]
  EP --> SCALE[256 GPU 部署]

每种并行解决不同的问题：

并行 解决什么 通信 适合
TP 单层放不下 频繁 all-reduce（跨机时慢） 同机 NVLink
PP 整模型放不下 阶段间传 hidden state（快） 跨机
EP MoE 专家放不下 All-to-All（中等） 同机 + 跨机

设计部署拓扑的经验法则：

  1. 先 TP 占满同机 NVLink：8 卡 NVLink 服务器内 TP=8
  2. TP 不够再 EP：MoE 模型用 EP 把专家分到多机
  3. 再不够再 PP：超大模型 + PP 跨机
  4. 最后 DP：吞吐扩展用数据并行

18.11 通信带宽决定一切

分布式推理的核心制约不是算力，是通信带宽。把所有相关带宽放一起看：

通信类型 带宽 用于
HBM ↔ SM 3.35 TB/s 单卡内权重读写
NVLink (H100) 900 GB/s（双向合计） 同机 GPU 间
NVLink 5（GB200 NVL72） 1.8 TB/s 机柜级 NVLink 域内 GPU 间
InfiniBand NDR（每 GPU 一张 400 Gbps 网卡） 50 GB/s（单向） 跨机
Ethernet 100G 12.5 GB/s（单向） 跨机（更便宜，更慢）

差异巨大。按同一口径（单向）比，NVLink 比 InfiniBand 快约 9×，比 Ethernet 快约 36×。这就是为什么：

  • 同机 NVLink 域内可以做 TP=8 而不显著影响 latency
  • 跨机 TP 几乎不可行——每次 all-reduce 都要被慢通信拖死
  • 跨机用 PP 或 EP——它们的通信量更小、对带宽不那么敏感

构建大模型推理集群时，硬件选型的核心是网络拓扑：

  • 一个机器内 8 张 H100 经 NVSwitch 全互联（最快的同机方案）
  • 多机之间用 InfiniBand 组成胖树（fat-tree）等无阻塞网络
  • 新一代的 GB200 NVL72 把 72 张 GPU 全部拉进一个 NVLink 域——对上层软件近似于一台「大 GPU」

这种硬件投入是 frontier model 推理的「门票」——开源模型再好，没有几亿美元的硬件配套，也跑不出 OpenAI / Anthropic 同等的服务质量。

18.12 一个完整的部署案例

把这一章所有内容串起来，看一个假想部署的量级估算：

任务：部署 Llama-3 70B 服务千万 DAU 的中文 chat 应用，要求 TPOT < 50 ms、TTFT < 1s。

Step 1：模型量化

  • 用 AWQ INT4 量化：140 GB → 35 GB
  • KV Cache 用 INT8：每 token 砍半
  • 单 H100（80 GiB）放下权重 + 30 GiB KV Cache 池

Step 2：单卡 attention 优化

  • 用 Flash Attention 3（H100 原生支持 TMA / WGMMA）
  • attention 这一段的墙上时间相比朴素实现是数倍的差距，长 prompt 下更明显

Step 3：推理引擎选择

  • 用 vLLM（自带 PagedAttention + Continuous Batching）
  • 启用 EAGLE-2 投机解码，假设在这个并发下单请求加速约 2.5×（示意取值：EAGLE-2 论文报告的 3.05–4.26× 是单请求口径，batch 一大收益会被摊薄，见第 17 章）

Step 4：单卡能服务多少

  • INT4 70B 单 token Decode：~10 ms（HBM 读 35 GB / 3.35 TB/s ≈ 10.5 ms）
  • 加 EAGLE-2 投机：按 Step 3 假设的 2.5× 折算，等效约 4 ms/token（估算，非实测）
  • TPOT < 50 ms ✓
  • KV Cache 池 30 GiB，INT8 KV 下每用户 8K context 占 1.25 GiB → 单卡约 24 个并发

Step 5：横向扩展

  • 假设峰值 QPS=10000、每个会话平均活跃 30 s，则同时在线约 300,000 个会话
  • 单卡 24 并发 → 需要约 12,500 张 H100
  • 按第 11 章用过的 $2.5 / GPU·hour 云价：12,500×720×2.5≈12{,}500 \times 720 \times 2.5 \approx $22M / 月

Step 6：成本优化

  • 启用 Prefix Caching（共享 system prompt + 文档）→ 命中的那部分 Prefill 算力直接省掉
  • 冷会话的 KV Cache offload 到 CPU 内存 → 把 HBM 让给活跃会话，同卡并发再上一档
  • 假设这两项合起来能把卡数压掉三成左右（示意取值，实际取决于前缀命中率和会话冷热分布），即约 8,750 卡、$16M / 月——仍然是「只有大厂玩得起」的数字

Step 7：跨机扩展

  • 实际峰值用 PD 分离架构：把 Step 6 之后的约 8,750 卡拆成 Prefill 集群约 2,000 卡（H100，算力优先）+ Decode 集群约 6,750 卡（H200，HBM 带宽优先）
  • 两个集群之间用 InfiniBand（NDR 400 Gbps，新机型可上 XDR 800 Gbps）传 KV Cache

这种量级的部署大致是头部模型服务商面对的日常（各家并未公开具体拓扑）。开源团队、小公司想达到同等服务能力，要么砸几亿美元自建集群，要么用 OpenAI / Anthropic 的 API。

本章小结

Flash Attention 部分

  1. GPU 内存层级：HBM 大慢、SMEM 小快——attention 必须吃尽 SMEM。
  2. 朴素 attention 的 HBM 访存量是 Θ(Nd+N2)\Theta(Nd+N^2)——长上下文场景下 attention 是瓶颈。
  3. Flash Attention = tiling + online softmax——把 attention 算在 SMEM 里：显存占用降到 O(N)O(N)（不物化 N×N），HBM 访存量则是 Θ(N2d2/M)\Theta(N^2d^2/M)，仍随 N2N^2 增长，省的是一个十倍量级的常数因子。这两个量别混。
  4. Online softmax 数学等价——一遍扫过 K/V 就得到 attention 输出，数学上与朴素 softmax 完全一致。
  5. FA1 → FA2 → FA3 持续优化：work partition、async pipeline、warp specialization、FP8 支持。FA3 在 H100 上 75% MFU。

分布式推理部分

  1. 三种并行：TP（切大矩阵）、PP（切层）、EP（切 MoE 专家）。
  2. TP 适合同机 NVLink：8 卡内 TP=8 是常见配置。跨机 TP 几乎不可行。
  3. PP 适合跨机扩展：通信少、bubble 限制。在线低延迟用得少。
  4. EP 是 MoE 专属：All-to-All 通信，需要专门优化。
  5. 所谓 3D 并行：原指 DP × TP × PP；MoE 再叠加 EP 是超大模型部署的常见组合，但卡数 = TP × PP × DP——EP 是 MoE 层把 DP 那几路换一种切法，不是再乘上去的第四维（§18.10）。
  6. 通信带宽是关键：NVLink（NVL72 这类机柜级 NVLink 域更快）> InfiniBand > Ethernet——选硬件就是在选拓扑。

第六部分到这里完结。我们用 5 章把 LLM 推理系统从「为什么慢」（两阶段）到「怎么快」（KV Cache、量化、投机解码、Flash Attention、分布式）全讲清楚了。

下一章是终章——第 19 章 Transformer 之后。Transformer 已经九年没变过骨架了，但研究界一直在探索替代方案：Mamba 等线性复杂度架构、Hybrid 混合架构、Diffusion-based 生成。我们会沿着这条路看一眼未来。

延伸阅读

  • Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022——FA1。
  • Dao, FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning, ICLR 2024——FA2。
  • Shah et al., FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision, NeurIPS 2024——FA3。
  • Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019——Tensor Parallelism 奠基。
  • Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019——Pipeline Parallelism 奠基。
  • Lepikhin et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020——Expert Parallelism。
  • DeepSeek-V3 Technical Report——DualPipe + PD-Sep 工程实战。
  • Megatron-LM GitHub: https://github.com/NVIDIA/Megatron-LM
  • vLLM 分布式文档: https://docs.vllm.ai/en/latest/serving/parallelism_scaling.html