Tensor Parallelism in LLM Inference

Background

§2.1 LLM Inference Computation

• Llama-2 architecture 简述（layer 结构、attention、KV cache）
• Prefill vs decode phases
• Arithmetic intensity $I$ 定义
• Roofline model: $T_{\text{compute}}^{\text{ideal}}=\text{FLOPs}/min(\pi ,{\beta }_{\text{mem}}\cdot I)$
• Utilization $\eta$: 实现 gap，actual = ideal / $\eta$
• 自然映射：prefill 高 $I$ → compute-bound；decode 低 $I$ → memory-bound

§2.2 Tensor Parallelism and Collective Communication

• Megatron-style TP：column-then-row partition → 每 layer 1 次 AllReduce
• Ring AllReduce algorithm
• Hockney $\alpha$-$\beta$ model: $T_{\text{comm}}=\alpha +M/{\beta }_{\text{net}}$（per collective）
• Latency $\alpha$ 作为实现 gap（与 $\eta$ 对偶）
• Comm Roofline 视角：$\text{BW}(M) = 1/(\alpha/M + 1/\beta_\text{net})$，与 §2.1 compute roofline 结构同构

§2.3 Modified Amdahl's Law

• Standard Amdahl: $S(p)=1/(s+(1-s)/p)$，常数 $s$
• 在我们的系统里失败：comm 开销随 $p$ 增长（§2.2 的 $T_{\text{comm}}(p)$）
• Modified Amdahl:

$$T(p)=\underset{\text{§2.1}}{\underbrace{\frac{T_{\text{compute}}^{\text{ideal}}}{p\cdot \eta }}}+\underset{\text{§2.2}}{\underbrace{N_{\text{coll}}[\alpha +\frac{2(p-1)M}{p{\beta }_{\text{net}}}]}}$$

• 两个 implementation gap 参数：$\eta, \alpha$
• §4 用这个公式做 a priori 预测、参数标定、双 roofline 可视化

archive

2.3

• TP 并行加速了矩阵乘法，但是引入了通信

• 一次矩阵乘法如何被并行：输入切分，输出需要 allreduce

• 具体讨论一次 allreduce 的 traffic 和 cost

  一次 allreduce 单向 traffic：

  $$ \text{Traffic}(B, T, d, n) = 2(n-1)\cdot \frac{V}{n} = \frac{2(n-1)}{n} \cdot 2 B T d \text{ bytes}. $$

  一次 allreduce 时间：

  $$ T_\text{AR}(n, V) \approx 2(n-1) \cdot (\alpha + \frac{m}{\beta}) = 2(n-1)\alpha + \frac{2(n-1)}{n} \cdot \frac{V}{\beta} $$

  对两项分别进行讨论。

  先讨论 V 固定 n 变化的时候，两项分别怎么变。

  再讨论 n 固定 V 变化的时候，两项分别怎么变。 大 V 是 bandwidth bound，小 V 是 latency bound。

• prefill 和 decode 阶段，一次 allreduce 时间

  用 PCIe 的实际数据 $\beta =32$~GB/s, with $\alpha \sim 1,\mu$。

  有两个变量 n V。先简单讨论两项和 n V 关系。latency项只和 n 有关系。bandwidth 项和 n V 都有关系。

  然后给出 V 的公式 2BTd。

  这节要讨论 prefill 和 decode，prefill T=L decode T=1。固定好 T 以后，那么所有在变的量只有 B 和 n。

  讨论 prefill 时（固定 T），B 和 n 变化带来的影响。

  讨论 decode 时（固定 T），B 和 n 变化带来的影响。

2.4

• 开头介绍三个经典指标

• 经典模型的两个问题

• 用新公式对端到端时间进行建模

  $$ T_p = \frac{(1-f) T_1}{p} + f T_1 + T_\text{comm},\ T_\text{comm} = \sum_{i=1}^{K} T_\text{AR}(p, V_i), $$

• 算出等效 f

  $\tilde{f}=f_1+\frac{p}{p-1}\cdot \frac{T_{\text{comm}}}{T_1}$

  $\tilde{f}=f_1+\frac{2Kp\alpha }{T_1}+\frac{2KV}{T_1\beta }$

• 把新的 Tp 分别应用在 prefill 和 decode 两个阶段

  主要是看 p 变化的时候，第一项和第三项怎么变。

  一次 prefill 时，先考虑 Tcomm 怎么变。

  首先，Tcomm bandwidth bound，直接忽略 latency。剩下的 bandwidth 项，和 p 和 V 有关系。

  假定 V 不变，那么 Tcomm 随 p 变大逐渐饱和。

Section 4

§4.1 User-facing performance

Data: wall-clock, end-to-end

指标: TTFT, TPOT, Total latency, Throughput（一张表覆盖 9 个 config）

Punchline: TTFT 随 $p$ 近线性下降；TPOT 随 $p$ 子线性，$B=1$ 甚至 regress → prefill 和 decode 行为不对称，引出后续 phase-split 分析

§4.2 Standard Amdahl's law falls short

Data shift: 一次性切换到 GPU active time on least-busy rank（之后 §4.2-4.5 一致使用）

论证: 用 $p=2$ 数据拟合 serial fraction $s$，外推 $p=4$ → 预测与实测系统性偏离

结论: standard Amdahl 假设 $s$ 与 $p$ 无关，但实际 comm 开销随 $p$ 增长 → 需要更精细的模型

§4.3 Modified Amdahl with explicit communication

公式: $$T(p) = \underbrace{\sum_k \frac{\text{FLOPs}k / p}{\eta_k \cdot \min(\pi, \beta\text{mem} \cdot I_k)}}{T\text{compute}} + \underbrace{T_\text{kv_copy}}{\text{bytes}/\beta\text{mem}} + \underbrace{N_\text{coll}\left[\alpha + \frac{2(p-1) M}{p \beta_\text{net}}\right]}{T\text{comm}}$$

两个 implementation gap 参数：$\eta$（compute utilization）和 $\alpha$（network latency），结构上对偶

§4.3 用理论值预测:

• 硬件参数：$\pi, \beta_\text{mem}, \beta_\text{net}$（A40 spec）
• 模型参数：FLOPs, bytes, message size $M$（从 Llama-2 7B 架构 + $B,L,p$ 算出）
• ${\alpha }_{\text{theo}}\approx 25,\mu \text{s}$（NCCL launch + PCIe round-trip + cross-NUMA 拆解）
• $\eta =1$（先理想化，§4.4 标定实际值）

→ 给出 a priori 预测，列表展示

§4.4 Kernel-level breakdown and parameter calibration

Phase split: prefill vs decode（NVTX-marked）

5 类 kernel: weight matmul, attention, communication, KV cache copy, other

关键 finding:

• Stacked bar chart（9 configs × 5 categories, decode）
• Comm 占比随 $p$ 增长，随 $B$ 缩小
• KV copy 在 $B=32,p=1$ 占 48%（HF eager artifact）

Per-GPU 不对称机制:

• Prefill 各 rank 同步 → spread ≈ 0
• Decode 各 rank drift → spread 几秒
• 机制：CUDA stream queue starvation + CPU launch jitter + NCCL busy-wait
• 选 least-busy GPU 的理由

参数标定（从实测数据）:

• ${\alpha }_{\text{actual}}$ 和 ${\beta }_{\text{net,actual}}$ 从 $p=2$ vs $p=4$ comm 时间反推（$\alpha$ 在差值里消掉，先得 $\beta$；回代得 $\alpha$）
• ${\eta }_k$ 从 (实测 $T_k$ × 理论 ceiling) / 理论 FLOPs 得到
• 对比 theoretical vs actual: ${\alpha }_{\text{theo}}\approx 25,\mu \text{s}$ vs ${\alpha }_{\text{actual}}\approx ?,\mu \text{s}$

§4.5 Double Roofline analysis

两张并排子图：

Compute Roofline:

• Ceiling: $min(\pi ,{\beta }_{\text{mem}}\cdot I)$
• 每个 (kernel, $B$, $p$) 一个点：理论 FLOPs / 实测时间
• 叠加 nsys 直接测的 Tensor Core util × $\pi$ → 两个独立估计互相验证

Communication Roofline（平滑曲线，不是折线）:

• $\text{BW}(M)=1/(\alpha /M+1/{\beta }_{\text{net}})$
• 渐近线: $M/\alpha$（latency-bound 斜线）和 ${\beta }_{\text{net}}$（bandwidth-bound 水平）
• 每个 AllReduce 一个点：$M$ vs $M/T_{\text{NCCL}}$
• 叠加 nsys 直接测的 PCIe TX/RX bandwidth → 独立验证

Punchline:

• Compute 在 ceiling 下面 1.5×（60-84% 利用率）— "fast enough"
• Comm 在 ceiling 下面 200-600×（< 1% 利用率, latency-bound）
• → TP scaling 瓶颈是 comm，不是 compute

§4.6 (Optional) Validation on quad-A100

实验: 把同一份 setup 跑在 Delta quad-A100（SXM4, NVLink NV4, 单 NUMA）

变化的参数: $\pi$ +2.1×, ${\beta }_{\text{mem}}$ +2.2×, ${\beta }_{\text{net}}$ +6×, NUMA penalty 消除

验证:

• §4.3 的模型公式不变，只换硬件参数 → 预测 quad-A100 行为
• 实测 → 与预测对比
• 预期：decode comm 占比从 63% 降到 ~7%, $B=1$ TP 不再 regress

论文价值: 把分析框架从 "fit & explain" 升级为 "predict & verify"

附录

如何切分 prefill decode

• 直接根据 kernel 区分，prefill decode 调用的是不同的 kernel
• 替换 model.forward，给 model.forward 前后加上 NVTX range
• 用 model.register_forward_pre_hook 和 register_forward_hook 注册 NVTX push/pop
• 不调用 model.generate，自己重写

新增内容

"We additionally capture per-kernel traces with NVIDIA Nsight Systems for each configuration; the processing pipeline used to extract per-phase, per-kernel-category timings is described in Section~4.4."

"我们观察到 prefill 阶段 weight_gemm 在多卡间存在 ~1% 的方差，这与节点上各 GPU 的 boost clock 在持续 compute-bound 负载下的浮动一致；attention 和 decode 阶段的 weight_gemm 由于落在 memory-bound 区域，不受 SM clock 浮动影响，因此跨卡时间一致。"