动画图解 LLM 系统 · 2026-10-20

分布式推理：数据怎样走过每张卡

权重留在哪里，KV 怎样增长，GPU 之间究竟传什么

→ 逐步讲解 · ▶ 连续播放 · 拖动时间轴回看 · F 全屏

先认清三样东西

权重常驻，KV 增长，激活往前走

两张卡各有自己的 HBM。权重是模型参数，KV cache保存已有 token 的注意力状态。

计算时，SM 从本卡 HBM 读权重与 KV。圆点表示读取的数据；原本的权重和 KV 仍留在 HBM。

中间结果叫 activation。本层用完即可释放或复用；跨卡的下一层需要它时，再通过互联传过去。

图中块大小、时长、两卡布局均为教学示意。实际缓存切分取决于 attention 类型和并行方式。

两个阶段 · 1

Prefill：整段提示词一起进入

提示词是 t₁…t₄。每层对这批 token 做投影和因果 attention，同层内有大量 token 可以一起计算。

工作按并行策略分给 GPU；矩阵乘的 token 维度更大，更容易复用读入的权重。

每张卡把自己负责的 K/V 写入本卡缓存。最后一个提示位置的输出用于预测第一个生成 token。

4 个 token 为教学示意；长提示也可以分块 prefill，具体瓶颈取决于形状、序列长度和硬件。

两个阶段 · 2

Decode：每轮只推进一个新位置

对每个活跃请求，这轮输入一个新 token。它经过所有层；每层只生成当前位置的 Q/K/V。

当前 Q 读取已有 KV，新 K/V 追加到缓存。过去 token 的各层激活不需要全部重算。

小 batch 的 decode 常花很多时间读权重与 KV。算完输出、采样得到下个 token，再进入下一轮。

图中圆点示意 HBM 读取和新 token 流动；多请求 batching 可提高权重复用，长上下文也会加大 KV 读取成本。

工程推演

先写下形状，才能判断通信是否必要

符号	形状 / 单位	它决定什么
T	本轮总 token 数；prefill 可为 B×S	矩阵乘和激活的规模
H / F	hidden width / FFN width	权重和每 token 计算量
p / b	TP rank 数 / 每元素字节	每卡切片和传输字节
X / W	[T,H] / [H,F]，Y=XW	乘加约 2THF FLOPs

例：BF16 的 b=2；MiB=2²⁰ bytes。rank 是一个参与通信的进程，下面每个 rank 对应一张 GPU。

Tensor parallel · 按输出维度切

不同输出通道：合在一起叫拼接

把 W 的输出维度切成 W₀、W₁，两张卡读相同的 x，各算一部分输出通道。

GPU 0 得到 y₀，GPU 1 得到 y₁。它们是不同位置的结果，不是同一位置的两份贡献。

如果下一步需要完整 y，就用 all-gather 拼接。若后续算子能直接吃分片，也可以继续保持分片。

采用 y = xW 记法，输出通道切分即 column parallel；示意维度 8。

Tensor parallel · 按输入维度切

同一输出的两份贡献：必须相加

现在切的是输入维度。GPU 0 算 x₀W₀，GPU 1 算 x₁W₁；两边都在为相同的输出位置做贡献。

[2, 5] 与 [3, 1] 逐元素相加得到 [5, 6]。all-reduce 做归约求和，再让每张卡拿到结果。

Transformer 常把“列切分”与“行切分”配对，中间保持分片，减少不必要的 all-gather。

数字为手算示例；若下一步继续消费分片，也可用 reduce-scatter 等布局。TP 不会让所有 KV 必然等分：少量 KV heads 可能被复制。

工程推演

两层 MLP 配对切分，中间不必拼回完整张量

Zᵣ = φ(X W₁,ᵣ)　　Yᵣ = Zᵣ W₂,ᵣ　　Y = Σᵣ Yᵣ

张量	完整形状	每个 rank
X	[T,H]	相同的 [T,H]
W₁ / Z	[H,F] / [T,F]	[H,F/p] / [T,F/p]
W₂ / Yᵣ	[F,H] / [T,H]	[F/p,H] / [T,H]

逐元素激活 φ 可以独立作用在每个通道。第二层产生同一输出的部分和，最后一次 all-reduce 即可。门控 MLP 要把 gate 与 up 的同组通道一起切。

工程推演

手算一次两卡 MLP，看到“拼接”和“求和”的区别

X=[1,2]　W₁=[[3,5],[4,6]]　W₂=[[1,2],[3,4]]

rank 0 拿 W₁ 第 0 列：
Z₀ = ReLU(1×3+2×4) = 11
再乘 W₂ 第 0 行：
Y₀ = 11×[1,2] = [11,22]

rank 1 拿 W₁ 第 1 列：
Z₁ = ReLU(1×5+2×6) = 17
再乘 W₂ 第 1 行：
Y₁ = 17×[3,4] = [51,68]

all-reduce 得到 [62,90]；与单卡 ReLU(XW₁)W₂ 完全一致。若拼接 Y₀、Y₁，会得到错误的 4 维结果。

工程推演

调用通信前，把输入与输出契约写清楚

操作（p 个 rank）	rank r 输入	每卡最终得到
all-reduce SUM	[T,H] 的部分和	[T,H] 完整和
all-gather（末维）	[T,H/p] 的不同通道	[T,H] 拼接结果
reduce-scatter SUM	[T,H] 的部分和	和的第 r 份 [T,H/p]
all-to-all（变长）	按目标 rank 打包的 token	发给本 rank 的各份 token

all-gather 的输出顺序由 rank 决定；reduce-scatter 是先求和、再分片。PyTorch 的 list API 收集的是各 rank 张量，拼接到哪一维由你显式决定。

动画计算 · 4 卡 ring

Ring all-reduce：每张卡到底发送多少数据？

4 个 rank 把 8 MiB 分成 4 份。reduce-scatter 走 3 轮，all-gather 再走 3 轮；每轮发送 M/p。

Vsend = 2(p−1)M/p = 2×3×8/4 = 12 MiB

工程推演

大消息看带宽，小消息常先撞上延迟

t ≈ 2(p−1)α + [2(p−1)/p] · M / β

教学参数
p=4；每轮 α=3 μs
有效单向带宽 β=100 GB/s
大消息 M=8 MiB

代入
启动项 = 6×3 = 18 μs
字节项 = 12×2²⁰/10¹¹ ≈125.8 μs
合计约 143.8 μs

β 是该链路和消息规模下可用的有效带宽；不能直接拿所有链路汇总的规格代入。NCCL 也可能选择 tree 等算法，实际策略以 trace 与通信微基准为准。

工程推演

同一个 hidden width，prefill 与 decode 的消息差很多

T；H=4096，BF16	M=T×H×2	四卡 ring 估算时间
decode：T=1	8192 B = 8 KiB	18 + 0.123 ≈ 18.1 μs
decode：T=32	262144 B = 256 KiB	18 + 3.932 ≈ 21.9 μs
prefill：T=1024	8388608 B = 8 MiB	18 + 125.829 ≈ 143.8 μs

继续沿用上一页 α、β。T 变小后，传输字节项几乎消失，但每层通信的固定延迟仍在。低 batch decode 增大 TP 后，省下的计算可能抵不过这些等待。

Pipeline parallel

按层分卡：跨卡传的是激活

GPU 0 放前几层，GPU 1 放中间层，GPU 2 放后几层。每段的权重和对应 KV 留在本段。

token 先完成 GPU 0 的层，再把中间 activation 交给 GPU 1，接着交给 GPU 2。

单个 token 的各段存在依赖。不同 micro-batch 可以同时占据不同阶段，减少等待形成的空泡。

12 层、3 卡为教学示意。PP 主要减少每卡权重占用；吞吐和单请求延迟需要分别评估。

工程推演

流水线吞吐：先算填满和排空用了几拍

总时间 = (m+p−1)tstage　　利用率 = m/(m+p−1)

p 个等时 stage，m 个独立 micro-batch，每段用时 tstage；忽略传输，按前向流水执行。第一批走完要 p 拍，其余每拍出一批。

p=4，m=1
4 拍只完成 1 批
4 张卡共有 16 个“卡·拍”
有效工作 4 个 → 25%

p=4，m=8
11 拍完成 8 批
有效工作 32 / 总槽位 44
利用率 72.7%

这些 micro-batch 必须彼此可独立推进。同一请求的下一个 decode token 依赖前一个采样结果，不能直接当作独立 micro-batch 填满流水线。

动画推演 · 前向流水

四批依次进入，空泡不会凭空消失

每次前进一拍：新批次从 S0 进入，已有批次前进一个 stage。最后一批在第 7 拍离开，总利用率 16/28 = 57.1%。

Expert parallel

把 token 发给它选中的专家

MoE 层的 router 为 token 选择专家。专家权重分布在各卡，只有被选中的专家需要处理这个 token。

Dispatch 把激活送到专家所在卡；同一 token 选中多个专家时会发出多份。专家在本卡读取自己的权重。

Combine 把专家输出送回，再按 router 权重相加。跨卡流动的是激活与路由信息，专家权重留在原卡。

图示 top-2。EP 常涉及 all-to-all / 变长交换；负载不均衡时，最忙的专家卡会拖慢这一轮。

工程推演

EP 的字节数，先数“token–expert 分配”

dispatch payload ≈ T × k × H × b × ρ

T 个 token，每个选 k 位专家，hidden width H，每元素 b bytes，ρ 是落到远端的分配比例。Combine 返回同宽结果，payload 量级相近。

例：1024 token，top-2
H=4096，BF16，远端比例 ρ=0.75
1024×2×4096×2×0.75
= 12 MiB dispatch

再加回程
Combine 约 12 MiB
单层整组远端 payload ≈ 24 MiB
另有索引、路由权重、对齐与通信协议开销

这里算整组 GPU 的总发送量，不是每卡 24 MiB。不同 token 选中同一远端 rank 的多个专家时，支持去重的实现还可能只发一份激活。

工程推演

专家总工作量相同，最长的队列决定何时合并

教学路由负载	每卡 token–expert 数	每卡 1 μs/分配的简化时间
平衡	[256,256,256,256]	max = 256 μs
偏斜	[640,128,128,128]	max = 640 μs

计算利用率 ≈ Σ nᵣ / (p × max nᵣ)

偏斜时利用率为 1024/(4×640)=40%。先按 rank / expert 记录 token 数与耗时，再判断是路由偏斜、专家太少，还是网络拥塞。平均值会把最长队列藏起来。

工程推演

放得下靠总参数，算得快才看激活参数

每卡峰值 ≈ 权重分片 + KV + 激活/通信工作区 + graph pool

假设一个 MoE 有 64B 总参数、每 token 激活 8B；权重 BF16 理想均分到 4 卡：每卡仍需 64×10⁹×2/4 = 32 GB 权重。

32 GB 十进制约为 29.80 GiB。若每卡 KV 8 GiB、工作区 4 GiB，合计约 41.80 GiB；40 GiB 卡放不下。

总参数中的共享层可能采用不同切分或复制，因此“理想均分”先作下界。逐 token 激活少只减少部分计算和读取，不会自动释放未选专家的常驻权重。

工程推演

把数学关系翻译成一次真实 collective

import torch.distributed as dist
# 每个 rank 的 y_partial.shape 都是 [T, H]
# y_partial = activation(x @ w1_local) @ w2_local
dist.all_reduce(y_partial, op=dist.ReduceOp.SUM)
# 结果留在原 tensor；各 rank 得到完整 y
y = y_partial + bias  # 共享 bias 在求和之后加一次

各 rank 必须对同一 group 按兼容的顺序调用 collective，shape、dtype 和元素数一致。某张卡跳过一层、早退或走不同分支，会使其他卡卡在通信处。

GPU 实验需先设置本 rank 的 CUDA device，并以 NCCL 初始化 process group；CPU 实验先验证切分的数值等价。

工程推演

先在高速互联域内做 TP，再评估跨节点方案

频繁的小消息
每层 TP 归约对延迟敏感。用目标 token 数对应的消息大小跑 all-reduce 测试，比只看最大带宽更有用。

相对少的较大消息
PP 在 stage 边界传激活；EP 取决于路由和 all-to-all 拥塞。跨节点布局要把慢链路放进成本模型。

# nccl-tests 已编译、4 个本机 GPU 的示例
./build/all_reduce_perf -b 8K -e 8M -f 2 -g 4

8 KiB–8 MiB 对应本课的 decode / prefill 消息范围。记录延迟曲线与正确性；多节点运行方式随 MPI / 集群配置变化。

完整生命周期

一条请求：有固定归属，也有临时跨卡流动

常驻：随层、头或专家划分

权重只在加载或迁移时搬。KV 随请求增长，保留在负责该层 / KV head 的设备上。

请求结束，归还 KV blocks，供下一条请求使用。

每轮：计算与通信交替

本卡读 HBM → 执行算子 → 必要的 TP 归约 / PP 传激活 / EP 分发 → 下一层。

临时 activation 和通信 buffer 可以复用。

Data parallel 则复制一组完整模型或模型分片组，把不同请求交给不同副本，提高服务吞吐。

想清楚再优化

选并行方式，先问三个问题

模型加上目标并发的 KV，每卡放得下吗？ 放不下就需要切分或压缩。
卡之间传得快吗？ 高频 TP 通信依赖低延迟、高带宽互联；跨节点布局要量通信成本。
目标是更快首 token、更快逐 token，还是更大吞吐？ Prefill、decode 和不同 batch 可能需要不同配置。

增加 GPU 后，每卡计算更少，但通信与调度还在。看端到端时间，才能知道新增的卡帮了多少忙。

工程推演

两卡升级到四卡，为什么每轮反而慢了？

教学测量：同一 batch / context	TP=2	TP=4
各卡计算关键路径	7.0 ms	4.0 ms
未被计算覆盖的通信	1.0 ms	4.5 ms
调度与其他	0.5 ms	0.5 ms
端到端一轮	8.5 ms	9.0 ms

计算省 3.0 ms，暴露的通信多 3.5 ms，最终慢 0.5 ms。检查 TP=4 是否跨慢链路、消息是否太小、是否失去重叠；若目标是吞吐，也比较“两组 TP=2 副本”。

关键路径表只累计不能重叠的区段；不能把 profiler 中所有 kernel 和通信时长直接相加。

练习 1

8 卡 ring：别把整组总量当作单卡量

M=16 MiB，p=8。求每卡发送字节数。若 α=2 μs、β=50 GB/s，估算一次 all-reduce 时间。

先写出假设、单位与中间步骤，再展开核对。

练习 2

有 4 个 stage，要多少独立批次才到 80%？

各 stage 等时 2 ms；忽略通信。求最少 micro-batch 数 m、总时间和每批平均吞吐间隔。

先写出假设、单位与中间步骤，再展开核对。

练习 3

数值多出一倍 bias，通信错在哪里？

两卡都计算 y_partial = z_local @ w_local + bias，然后 all_reduce(SUM)。输出与单卡比较，总是多出一份 bias。该怎样修复？

先写出假设、单位与中间步骤，再展开核对。

综合案例 · 容量、通信、延迟

两台机器各 4 张卡，先筛掉放不下的方案

教学模型 64B 总参数、BF16；每卡 40 GiB。目标并发的 KV 已计入下表；权重按理想均分估算，共享层复制另查实现。

方案	每卡权重 + KV + 其他	逐 token 关键路径
TP=4，一机一副本	29.8 + 6 + 5 = 40.8 GiB	尚未满足容量
TP=8，跨两机	14.9 + 3 + 5 = 22.9 GiB	计算 4 + 通信 7 + 调度 1 = 12 ms
TP=4，PP=2	14.9 + 3 + 5 = 22.9 GiB	两段各 4 + 跨段 1 + 调度 1 = 10 ms

其他 5 GiB 含工作区与 graph pool。延迟为同一负载下的教学测量；表中的 KV 可等分是假设，需要用实际 KV heads 与复制策略核对。

练习 4 · 部署决策

逐 token 目标 11 ms，你先尝试哪种配置？

沿用上一页预算。选一个能放下且满足延迟目标的方案；若使用 PP=2，有 8 个独立 micro-batch，等时 stage 的理想利用率是多少？MoE 若出现 [640,128,128,128] 的专家负载，仅增加 EP 是否足够？

选择由给定数据决定。上线前再量长上下文峰值、p99 和目标到达率下吞吐；PP 的吞吐填充不能消除同一请求的 token 依赖。

动手实验

用小矩阵验证 TP，再算通信与流水线预算

python3 lab.py
# TP result = [62, 90]; ring = 12.00 MiB; PP = 72.73%

先把 W₁ 的某一列改掉，预测输出再运行。把 TP 从 4 改为 8，比较 ring 的固定延迟项；将流水线批次数改成 1，验证空泡为何增大。

下载 CPU 实验代码 · lab.py

Python 3 标准库即可运行。模拟输出用于检验推导；GPU 性能需要在目标硬件上测量。

继续深入

参考与延伸阅读

NVIDIA NCCL · Collective Operations：all-reduce、all-gather 与 reduce-scatter 的数据语义。
vLLM · Distributed Serving：TP / PP 的部署组合与节点布局。
vLLM · Expert Parallel Deployment：EP 分发、后端与服务配置。
PyTorch · Distributed communication：collective API、process group 与 tensor 契约。

接着看：Tensor parallel · 流水线并行 · CPU overhead