""" ============================================================================= CS349E · 前置讲 P5 · 05_communication_in_llm_inference.py —— LLM 推理里的通信账：张量并行、MoE、KV 搬运 ============================================================================= 【这个脚本要回答的问题】 前四个脚本讲了网络和集合通信本身。这里把它们放回 LLM 推理，算几笔具体的账： 演示 1：张量并行（TP = 8，一个 H100 节点）跑 Llama-3.1-70B 的 decode：每步 160 次 all-reduce，每次只有几 KB 到几 MB。 通信占一步的多少？延迟 α 有多要紧？哪种 all-reduce 算法合适？ 演示 2：同一个模型的 prefill（4096 个 token）：每次 all-reduce 64 MB，通信和计算谁大？能不能藏起来？ 演示 3：TP 跨节点（16 块 GPU，2 个节点）会怎样——为什么张量并行通常不出节点。 演示 4：MoE 的专家并行：每层两次 all-to-all（分发、合并），节点内 vs 跨节点，BF16 vs FP8。 演示 5：解耦式推理里的 KV cache 搬运：一条提示的 KV 有多大，搬一次要多久，和 prefill 本身比。 【怎么读输出】全部是公式（α–β 模型 + 第 1 讲 llm_math 的 roofline），每次运行结果相同。计算时间是 roofline 下界 （理想的 100% 利用率），通信时间用 α–β 模型（α 是假设值，演示 1 给了两个 α 做对照）。所以比例只看量级。 【运行】python3 05_communication_in_llm_inference.py （不到 1 秒） 【依赖】numpy；从 ../lecture01 导入 console、hw_specs、llm_math；同目录的 comm_sim。 """ import os import sys HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 from hw_specs import H100 # noqa: E402 from llm_math import LLAMA31_70B, decode_step, prefill_time # noqa: E402 from comm_sim import IB_NDR, NVLINK4, t_formula # noqa: E402 M = LLAMA31_70B ARS_PER_LAYER = 2 # Megatron 式张量并行：注意力输出、MLP 输出各一次 all-reduce BF16 = 2 def ar_time(algo, p, nbytes, alpha, beta): """一次 all-reduce 的时间（交换机拓扑上的 α–β 公式，见 comm_sim.t_formula）。""" return t_formula(algo, p, nbytes, alpha, beta) # --------------------------------------------------------------------------- # 演示 1：TP decode # --------------------------------------------------------------------------- def demo1_tp_decode(p=8, ctx=4096): """ 每一层的 all-reduce 消息 = B × d_model × 2 字节（BF16 的激活）；每步 n_layers × 2 次。 计算 / 访存时间：llm_math.decode_step(n_chips = 8)，理想切分、**不含通信**（第 1 讲的下界）。 两个 α：2 µs（链路级的量级）和 20 µs（更接近真实 NCCL 小消息 all-reduce 的量级，见 03 脚本演示 3 的讨论）。 """ print_header(f"演示 1：张量并行 TP = {p}（一个 H100 节点）跑 {M.name} 的 decode（ctx = {ctx}）") rows = [] res = {} n_ar = M.n_layers * ARS_PER_LAYER for B in (1, 8, 64, 256): msg = B * M.d_model * BF16 step = decode_step(M, H100, B, ctx, n_chips=p) t_comp = step["t_step"] cells = [B, fmt_bytes(msg), fmt_time(t_comp)] for alpha in (2e-6, 20e-6): for algo in ("ring", "rd"): t = n_ar * ar_time(algo, p, msg, alpha, NVLINK4.bw) res[(B, alpha, algo)] = (t, t_comp) cells.append(f"{fmt_time(t)}（{t / (t + t_comp) * 100:.0f}%）") rows.append(cells) print(f"每步 {n_ar} 次 all-reduce；括号里是通信占'计算 + 通信'的比例（假设两者不重叠）。") print_table(["batch B", "每次 AR 的消息", "每步计算 / 访存（下界）", "α=2µs ring", "α=2µs 递归倍增", "α=20µs ring", "α=20µs 递归倍增"], rows) r1 = res[(1, 2e-6, "ring")] r1b = res[(1, 20e-6, "rd")] print_observations([ f"decode 时每次 all-reduce 只有 {fmt_bytes(M.d_model * BF16)}（B = 1）到 {fmt_bytes(256 * M.d_model * BF16)}（B = 256）——" "全部落在延迟主导的区间（01 脚本：NVLink 的半带宽点约 900 KB）。", f"B = 1 时，用 ring、α = 2 µs，{n_ar} 次 all-reduce 一共 {fmt_time(r1[0])}，已经是每步计算 / 访存时间 " f"{fmt_time(r1[1])} 的 {r1[0] / r1[1]:.1f} 倍：ring 每次要 14 步，每步一个 α。", "换成步数少的算法（递归倍增 3 步；实际引擎用的'one-shot / two-shot'自定义 all-reduce 也是同一个思路），" f"通信时间降到原来的约 {res[(1, 2e-6, 'rd')][0] / r1[0]:.0%}。vLLM、TensorRT-LLM、SGLang 都为小消息专门写了 all-reduce 内核。", f"α 的影响巨大：α = 20 µs 时，即使用递归倍增，B = 1 的通信也占了 {r1b[0] / (r1b[0] + r1b[1]) * 100:.0f}%。" "所以 decode 的张量并行越切越细并不总是越快——每多一次 all-reduce，就多一份固定的延迟。", "B 越大，计算 / 访存时间增长得比通信快（消息变大，但仍远低于半带宽点），通信的占比下降。", ]) # --------------------------------------------------------------------------- # 演示 2：TP prefill # --------------------------------------------------------------------------- def demo2_tp_prefill(p=8): """ prefill 一条 n_tok 个 token 的提示：每次 all-reduce = n_tok × d_model × 2 字节。 计算时间：llm_math.prefill_time(n_chips = 8)（compute-bound 的下界）。通信用 ring、α = 2 µs。 """ print_header(f"演示 2：同一个模型的 prefill（TP = {p}）——通信和计算谁大？") rows = [] res = {} n_ar = M.n_layers * ARS_PER_LAYER for n_tok in (512, 4096, 32768): msg = n_tok * M.d_model * BF16 pf = prefill_time(M, H100, n_tok, n_chips=p) t_c = n_ar * ar_time("ring", p, msg, NVLINK4.alpha, NVLINK4.bw) res[n_tok] = (t_c, pf["t_step"]) rows.append([n_tok, fmt_bytes(msg), fmt_time(pf["t_step"]), pf["bound"], fmt_time(t_c), f"{t_c / pf['t_step'] * 100:.0f}%"]) print_table(["提示长度", "每次 AR 的消息", "prefill 计算（下界）", "受限于", "通信（ring，α = 2 µs）", "通信 / 计算"], rows) r = res[4096] per_tok_comm = n_ar * 2 * (p - 1) / p * M.d_model * BF16 / NVLINK4.bw per_tok_lin = M.linear_flops_per_token() / p / H100.peak_flops("bf16") alpha_part = n_ar * 2 * (p - 1) * NVLINK4.alpha print_observations([ f"prefill 时消息变成几十 MB，进入带宽主导的区间：4096 个 token 时，每次 64 MB 左右，{n_ar} 次一共 {fmt_time(r[0])}，" f"是计算时间 {fmt_time(r[1])} 的 {r[0] / r[1] * 100:.0f}%。", f"比例随提示长度下降（{res[512][0] / res[512][1] * 100:.0f}% → {res[4096][0] / res[4096][1] * 100:.0f}% → " f"{res[32768][0] / res[32768][1] * 100:.0f}%）。拆开看：每个 token 的通信带宽时间约 {per_tok_comm * 1e6:.1f} µs、" f"线性层计算约 {per_tok_lin * 1e6:.1f} µs，两者都 ∝ token 数，比值 {per_tok_comm / per_tok_lin * 100:.0f}% 是一个常数" "（'每个 token 的 FLOPs / 通信字节' 对 '算力 / 带宽'，即第 1 讲的网络 roofline）；" f"另外每次 prefill 还有 {fmt_time(alpha_part)} 固定的 α，短提示时它占比很大；而注意力的计算 ∝ n²，长提示时计算增长得更快。", "这部分通信可以和计算重叠：把一层的计算切成几块，前一块的 all-reduce 和后一块的矩阵乘法同时进行" "（第 4 讲；How to Scale Your Model 第 3 章'通信与计算重叠'）。只要通信时间小于计算时间，理论上可以完全藏住。", ]) # --------------------------------------------------------------------------- # 演示 3：TP 跨节点 # --------------------------------------------------------------------------- def demo3_tp_cross_node(B=64, ctx=4096): """ 把 TP 从 8 扩到 16（两个节点）：计算 / 访存时间减半（理想切分），但 all-reduce 要跨节点。 跨节点的 all-reduce 用层次化算法的近似：节点内 RS + AG（NVLink）+ 节点间 ring（网卡，每块 GPU 处理 1/8）。 """ print_header(f"演示 3：TP 从 8 扩到 16（跨两个节点），{M.name}，decode B = {B}") msg = B * M.d_model * BF16 n_ar = M.n_layers * ARS_PER_LAYER rows = [] s8 = decode_step(M, H100, B, ctx, n_chips=8) c8 = n_ar * ar_time("ring", 8, msg, NVLINK4.alpha, NVLINK4.bw) rows.append(["TP = 8（一个节点）", fmt_time(s8["t_step"]), fmt_time(c8), fmt_time(s8["t_step"] + c8)]) s16 = decode_step(M, H100, B, ctx, n_chips=16) intra = 2 * (7 * NVLINK4.alpha + 7 / 8 * msg / NVLINK4.bw) inter = ar_time("ring", 2, msg / 8, IB_NDR.alpha, IB_NDR.bw) c16 = n_ar * (intra + inter) rows.append(["TP = 16（两个节点，层次化 AR）", fmt_time(s16["t_step"]), fmt_time(c16), fmt_time(s16["t_step"] + c16)]) print_table(["配置", "每步计算 / 访存（下界）", "每步通信", "合计（不重叠）"], rows) print_observations([ f"TP = 16 让每块 GPU 的计算 / 访存减半（{fmt_time(s8['t_step'])} → {fmt_time(s16['t_step'])}），" f"但每步的通信从 {fmt_time(c8)} 涨到 {fmt_time(c16)}：每次 all-reduce 都多了跨节点的一段（网卡的 α 和带宽都更差）。" + ("合计反而变慢了。" if s16["t_step"] + c16 > s8["t_step"] + c8 else "合计略快一点，但远没有翻倍。"), "所以实践中 TP 一般不超过一个 NVLink 域（8 块 H100；GB200 NVL72 把这个域扩到了 72 块），" "更大的模型或更多的 GPU 用流水线并行、专家并行或多副本来扩展（第 4 讲）。", ]) # --------------------------------------------------------------------------- # 演示 4：MoE 的 all-to-all # --------------------------------------------------------------------------- def demo4_moe_alltoall(tokens_per_gpu=256, d=7168, k=8, n_moe_layers=58): """ 一个 DeepSeek-V3 规模的 MoE（d_model = 7168，每个 token 选 top-8 个专家，58 个 MoE 层——取自其技术报告的配置）。 专家并行 EP：专家分布在 EP 块 GPU 上。每个 MoE 层： 分发（dispatch）：每个 token 的向量发给它选中的 k 个专家所在的 GPU（同一块 GPU 上的不用发）； 合并（combine）：专家的输出发回来。 每块 GPU 每层发出的字节 ≈ tokens × k × d × 字节数 × (EP − 1)/EP（假设专家均匀分布、选择均匀）。 节点内用 NVLink 450 GB/s；跨节点时，How to Scale Your Model 第 12 章指出 all-to-all 基本只能用网卡的带宽（50 GB/s）。 """ print_header(f"演示 4：MoE 的专家并行——每层 2 次 all-to-all（每块 GPU {tokens_per_gpu} 个 token，d = {d}，top-{k}）") rows = [] res = {} for ep, where, link in ((8, "节点内（NVLink）", NVLINK4), (32, "4 个节点（网卡）", IB_NDR)): for fmt_name, disp_bytes in (("BF16", 2), ("FP8（分发）", 1)): disp = tokens_per_gpu * k * d * disp_bytes * (ep - 1) / ep comb = tokens_per_gpu * k * d * 2 * (ep - 1) / ep t_layer = ((ep - 1) * link.alpha + disp / link.bw) + ((ep - 1) * link.alpha + comb / link.bw) res[(ep, fmt_name)] = t_layer rows.append([f"EP = {ep}，{where}", fmt_name, fmt_bytes(disp), fmt_bytes(comb), fmt_time(t_layer), fmt_time(t_layer * n_moe_layers)]) print_table(["配置", "分发的格式", "每层分发（每 GPU）", "每层合并（每 GPU）", "每层 2 次 all-to-all", f"× {n_moe_layers} 层"], rows) print_observations([ f"每个 token 要发给 {k} 个专家，all-to-all 的字节数正比于 top-k：每层每块 GPU 要发出几十 MB。", f"节点内（NVLink）每步的 all-to-all 合计约 {fmt_time(res[(8, 'BF16')] * n_moe_layers)}；跨 4 个节点（网卡）要" f" {fmt_time(res[(32, 'BF16')] * n_moe_layers)}——慢了约 {res[(32, 'BF16')] / res[(8, 'BF16')]:.0f} 倍。" "这就是 MoE 推理偏爱大的 NVLink 域（GB200 NVL72）的原因（04 脚本：all-to-all 吃的是对分带宽）。", "DeepSeek-V3 用 FP8 做分发（合并仍用 BF16），把分发的字节减半；还专门写了把 NVLink 和 InfiniBand 叠起来用的通信内核，" "并让通信和计算重叠（双 micro-batch 交替）。这些都是在和本演示里的这笔账较劲。", "这里假设专家的选择是均匀的；实际的负载不均衡（热门专家）会让最忙的那块 GPU 收到远多于平均的数据——" "all-to-all 的时间由最忙的那块决定（第 3、4 讲）。", ]) # --------------------------------------------------------------------------- # 演示 5：KV cache 的搬运 # --------------------------------------------------------------------------- def demo5_kv_transfer(p=8): """ 解耦式推理（第 5 讲）：prefill 在一组 GPU 上算完一条提示，把它的 KV cache 发给另一组做 decode 的 GPU。 KV 大小 = 提示长度 × 每 token 的 KV 字节（llm_math：Llama-3.1-70B，BF16，GQA 8 个 KV 头）。 传输：同一节点 NVLink（单卡 450 GB/s）；跨节点，一块网卡（50 GB/s）或节点的 8 块网卡并行（400 GB/s）。 对照：这条提示在 8 块 H100 上 prefill 的时间（roofline 下界）。 """ print_header(f"演示 5：解耦式推理里的 KV cache 搬运（{M.name}，每 token 的 KV {fmt_bytes(M.kv_bytes_per_token())}）") rows = [] res = {} for n_tok in (1024, 4096, 32768): kv = n_tok * M.kv_bytes_per_token() pf = prefill_time(M, H100, n_tok, n_chips=p)["t_step"] t_nv = NVLINK4.time(kv) t_1nic = IB_NDR.time(kv) t_8nic = IB_NDR.alpha + kv / (8 * IB_NDR.bw) res[n_tok] = (kv, pf, t_1nic, t_8nic) rows.append([n_tok, fmt_bytes(kv), fmt_time(pf), fmt_time(t_nv), fmt_time(t_1nic), fmt_time(t_8nic), f"{t_8nic / pf * 100:.0f}%"]) print_table(["提示长度", "KV 大小", "prefill（8×H100 下界）", "NVLink（节点内）", "1 块网卡", "8 块网卡并行", "8 网卡 / prefill"], rows) kv4, pf4, t1, t8 = res[4096] print_observations([ f"一条 4096 token 的提示，KV 就有 {fmt_bytes(kv4)}。用一块网卡搬要 {fmt_time(t1)}，" f"8 块网卡并行要 {fmt_time(t8)}，大约是 prefill 本身（{fmt_time(pf4)}）的 {t8 / pf4 * 100:.0f}%。", "提示越长，prefill 的注意力部分按 n² 增长，而 KV 只按 n 增长：长提示的搬运占比略小一些。", "KV 可以一层一层地发：第 l 层的 KV 一算完就开始传，和后面各层的计算重叠（DistServe、Mooncake 都这么做，第 5、7 讲），" "真正暴露在首 token 延迟里的只是最后几层。", "用了 GQA（8 个 KV 头而不是 64 个）的模型 KV 小 8 倍——模型结构直接决定了解耦式推理的网络账。" "用 FP8 存 KV 又能再减半（P3、第 7 讲）。", ]) if __name__ == "__main__": demo1_tp_decode() demo2_tp_prefill() demo3_tp_cross_node() demo4_moe_alltoall() demo5_kv_transfer()