LLM Serving System的计算与通信建模
目录
符号表示与基础参数
基础参数
| 参数 | 符号 | 示例值 | 含义 | 对性能的影响 |
|---|---|---|---|---|
| 批量大小 | $B$ | 16 | 同时处理多少个请求 | $B$↑ → 硬件利用率↑, KV占用↑ |
| 序列长度 | $S$ | 4096 | 输入文本的token数 | $S$↑ → 计算↑(Prefill), 内存↑(Decode) |
| 隐藏维度 | $d$ | 7168 | Transformer内部表示维度 | $d$↑ → 计算∝$d^2$, 内存∝$d$ |
| 层数 | $L$ | 61 | Transformer块的数量 | $L$↑ → 总参数∝$L$, 推理时间∝$L$ |
| 注意力头数 | $H$ | 128 | 多头注意力的头数 | $H$↑ → TP上限↑ |
并行策略参数
| 参数 | 符号 | 示例值 | 含义 | 关键关系 |
|---|---|---|---|---|
| 张量并行度 | $TP$ | 8 | 将单个tensor切分 | $H \% TP = 0$，每个GPU处理$H/TP$个头 |
| 流水线并行度 | $PP$ | 4 | 将层切分到多个GPU | 每个stage处理$L/PP$层，存在bubble开销 |
| 数据并行度 | $DP$ | 2 | 将batch切分 | 每个GPU处理$B/DP$个请求，各副本独立推理 |
| 专家并行度 | $EP$ | 64 | 将专家切分到多个GPU | 每个GPU存储$E/EP$个专家，需All-to-All通信 |
| 总GPU数 | $N_{GPU}$ | 512 | 集群中GPU总数 | $N_{GPU} = TP \times PP \times DP \times EP$ |
| 通信组大小 | $TP_{comm}$ | 8 | 单次通信涉及的GPU数 | TP需要All-Reduce，EP需要All-to-All |
| 每GPU参数量 | $P_{GPU}$ | ~50GB | 单GPU存储的参数量 | $P_{GPU} \approx P_{total}/(TP \times PP \times EP)$，DP不影响参数量 |
Continuous Batching参数（实际Serving场景）
| 参数 | 符号 | 含义 | 与简化模型的关系 |
|---|---|---|---|
| Prefill请求数 | $B_p$ | 正在做Prefill的请求数 | 动态变化 |
| Decode请求数 | $B_d$ | 正在做Decode的请求数 | 动态变化 |
| 总请求数 | $B = B_p + B_d$ | 当前batch中的总请求数 | 简化模型中$B$为固定值 |
| Prefill token总数 | $T_p = \sum_{i=1}^{B_p} S_i$ | 所有Prefill请求的token总数 | 各请求序列长度$S_i$不同 |
| Decode token数 | $T_d = B_d$ | 所有Decode请求生成的token数 | 每个请求生成1个token |
| 总token数 | $T_{total} = T_p + T_d$ | 实际处理的token总数 | 通信量∝$T_{total}$，而非$B \times S$ |
并行策略详解（简化模型）
注：以下分析为简化模型，假设batch内所有请求处于同一阶段（纯Prefill或纯Decode）。实际serving场景会使用continuous batching，详见后续章节。
TP对Attention的影响
TP的All-Reduce通信量推导
完整计算（无TP）
$$Y = X \times W$$- $X: (N, d_{in})$ - 输入
- $W: (d_{in}, d_{out})$ - 权重矩阵
- $Y: (N, d_{out})$ - 输出
TP切分（权重按行切分）
-
输入切分：每个GPU持有输入的一部分 $(N, d_{in}/TP)$
-
权重按行切分：
1 2 3 4完整权重: W = (d_in, d_out) 切分为TP份： GPU i: W_i = (d_in/TP, d_out) <- 第 i×(d_in/TP) 到 (i+1)×(d_in/TP)-1 行 -
局部计算： 每个GPU计算：$Y_i = X_i \times W_i = (N, d_{in}/TP) \times (d_{in}/TP, d_{out}) \to (N, d_{out})$
注意：每个GPU都产生$(N, d_{out})$的输出，但这些是部分结果！
-
为什么是部分结果？
完整矩阵乘法应该是：
$$Y = \begin{bmatrix} X_0 & X_1 & \cdots & X_{TP-1} \end{bmatrix} \times \begin{bmatrix} W_0 \\ W_1 \\ \vdots \\ W_{TP-1} \end{bmatrix} = \sum_{i=0}^{TP-1} X_i W_i$$每个GPU只计算一项，需要All-Reduce求和得到最终结果。
-
All-Reduce通信量：
- 需要同步的数据大小：$(N \times d_{out})$个元素
- 实际通信量（Ring All-Reduce算法，每GPU）： $$\text{发送} + \text{接收} = 2 \times \frac{TP-1}{TP} \times N \times d_{out} \text{ 个元素}$$
- 当$TP=8$时：实际传输 $1.75 \times N \times d_{out}$ 个元素
通用规律
所有使用TP按行切分权重的层（输出投影、FFN第二层等），All-Reduce通信量取决于输出形状：
$$\text{需要同步} = N \times d_{out}, \quad \text{实际传输（每GPU）} = 2 \times \frac{TP-1}{TP} \times N \times d_{out}$$简化约定
为简洁，后文"通信量"指需要同步的数据大小（即输出大小），实际传输量需乘以系数 $2(TP-1)/TP$。
Prefill vs Decode的本质区别
虽然两者都是矩阵-矩阵乘法（当batch size $B > 1$时），但性质完全不同：
| 维度 | Prefill | Decode |
|---|---|---|
| 输入形状 | $(B, S, d)$，$S$大（如2048） | $(B, 1, d)$，序列长度为1 |
| 矩阵乘法 | $(B, S, d) \times (d, d)$，大矩阵 | $(B, d) \times (d, d)$，小矩阵 |
| 计算量 | $2BSd^2$ | $2Bd^2$ |
| 算术强度 | $\frac{2BSd^2}{2d^2} \approx BS$（高） | $\frac{2Bd^2}{2d^2} \approx B$（低） |
| 瓶颈 | Compute-bound | Memory-bound |
| GPU利用率 | 高（80-90%） | 低（10-30%） |
| 优化方向 | 提高并行度（TP）、计算优化 | 减少访存（量化、KV Cache压缩） |
| 通信影响 | 可被计算掩盖 | 延迟直接暴露 |
标准MHA（Multi-Head Attention）
Prefill阶段
步骤1：QKV投影
计算：
-
完整计算：$Q, K, V = X \times W_Q, X \times W_K, X \times W_V$
- 输入 $X: (B, S, d)$
- 权重 $W_Q, W_K, W_V: (d, d)$
- 输出 $Q, K, V: (B, S, d)$ → reshape成$(B, S, H, d_h)$
-
TP切分：每个GPU处理$H/TP$个头，权重矩阵按列切分
-
每个GPU上的shape：
- 输入 $X$: $(B, S, d)$（广播到所有GPU）
- 权重 $W_Q, W_K, W_V$: $(d, \frac{d}{TP})$（列切分）
- 输出 $Q, K, V$: $(B, S, \frac{d}{TP})$ → reshape成$(B, S, \frac{H}{TP}, d_h)$
-
每个GPU计算量：$3 \times \frac{2BS \times d \times d}{TP} = \frac{6BSd^2}{TP}$（Q、K、V三个投影）
-
KV Cache占用：
- 每GPU每层：$2 \times B \times S \times \frac{d}{TP}$（$\frac{H}{TP}$个头，每头K和V各$B \times S \times d_h$）
- 所有GPU总计：$2 \times B \times S \times d$
通信：无（TP按列切分，各GPU独立计算）
步骤2：注意力计算
计算：Scaled Dot-Product Attention
-
每个GPU上的shape：
- 输入Q, K, V：$(B, S, \frac{H}{TP}, d_h)$（每个GPU处理$\frac{H}{TP}$个头）
- 将batch和序列维度合并：$(B \times S, \frac{H}{TP}, d_h)$ 或看作每个头处理$(B, S, d_h)$
-
计算流程（每个头独立）：
- $\text{Score} = \frac{Q \times K^T}{\sqrt{d_h}}$: $(B, S, d_h) \times (B, d_h, S) \to (B, S, S)$（attention score矩阵）
- $\text{AttentionWeights} = \text{Softmax}(\text{Score})$: $(B, S, S)$（对每行做softmax）
- $\text{Output} = \text{AttentionWeights} \times V$: $(B, S, S) \times (B, S, d_h) \to (B, S, d_h)$
-
输出合并：
- 每个GPU输出：$(B, S, \frac{H}{TP}, d_h)$
- Concat heads后：$(B, S, \frac{d}{TP})$
- 将 $\frac{H}{TP}$ 个头的维度 $d_h$ 拼接：$\frac{H}{TP} \times d_h = \frac{d}{TP}$
-
每GPU计算量：$\frac{2BS^2 \times d}{TP}$（$Q \times K^T$和$\text{AttentionWeights} \times V$各占一半）
通信：无（各GPU独立计算各自的头）
步骤3：输出投影
计算：
-
完整计算：$Y = \text{Concat}(Q, K, V) \times W_O$
- 输入（concat后的attention输出）: $(B, S, d)$
- 权重 $W_O$: $(d, d)$
- 输出 $Y$: $(B, S, d)$
-
TP切分：权重矩阵按行切分
-
每个GPU上的shape：
- 输入: $(B, S, \frac{d}{TP})$（每个GPU的attention输出）
- 权重 $W_O$: $(\frac{d}{TP}, d)$（行切分）
- 局部输出: $(B, S, \frac{d}{TP}) \times (\frac{d}{TP}, d) \to (B, S, d)$（部分结果）
-
每GPU计算量：$\frac{2BS \times d \times d}{TP}$
通信：
- All-Reduce：合并所有GPU的部分结果$(B, S, d)$
- 通信量：$B \times S \times d$ 个元素（每层）
- $S$很大时通信压力大，但可被计算掩盖
- 通信/计算比：$\frac{BSd}{BSd^2/TP} = \frac{TP}{d} \approx 0.1\%$（$d=7168, TP=8$时）
Decode阶段
步骤1：QKV投影
计算：
-
完整计算：$Q, K, V = X_{new} \times W_Q, X_{new} \times W_K, X_{new} \times W_V$
- 输入 $X_{new}$（新token）: $(B, d)$
- 权重 $W_Q, W_K, W_V$: $(d, d)$
- 输出 $Q, K, V$: $(B, d)$ → reshape成$(B, H, d_h)$
- K和V追加到KV Cache
-
TP切分：权重矩阵按列切分
-
每个GPU上的shape：
- 输入 $X_{new}$: $(B, d)$（广播到所有GPU）
- 权重 $W_Q, W_K, W_V$: $(d, \frac{d}{TP})$（列切分）
- 输出 $Q, K, V$: $(B, \frac{d}{TP})$ → reshape成$(B, \frac{H}{TP}, d_h)$
-
计算量：$3 \times \frac{2B \times d^2}{TP} = \frac{6Bd^2}{TP}$（Q、K、V三个投影）
-
KV Cache占用（$S$逐渐增长）：
- 每GPU每层：$2 \times B \times S \times \frac{d}{TP}$
- 所有GPU总计：$2 \times B \times S \times d$
通信：无
步骤2：注意力计算
计算：新token的Q与所有历史KV做注意力
-
每个GPU上的shape：
- $Q$（新token）: $(B, \frac{H}{TP}, d_h)$（序列长度为1）
- $K$（历史，从KV Cache读取）: $(B, S, \frac{H}{TP}, d_h)$（$S$是历史长度，包含新token）
- $V$（历史，从KV Cache读取）: $(B, S, \frac{H}{TP}, d_h)$
-
计算流程（每个头独立，每个GPU处理$\frac{H}{TP}$个头）：
- $\text{Score} = \frac{Q \times K^T}{\sqrt{d_h}}$: $(B, 1, d_h) \times (B, d_h, S) \to (B, 1, S)$
- $\text{Weights} = \text{Softmax}(\text{Score})$: $(B, 1, S)$
- $\text{Output} = \text{Weights} \times V$: $(B, 1, S) \times (B, S, d_h) \to (B, 1, d_h)$
-
输出合并：
- 每个GPU输出：$(B, \frac{H}{TP}, d_h)$
- Concat heads维度：$(B, \frac{d}{TP})$
-
每GPU计算量：$\frac{2B \times S \times d}{TP}$（$S$是历史长度）
访存（主要瓶颈）：需要读取所有KV Cache
- 每GPU访存量：$2 \times B \times S \times L \times \frac{d}{TP}$（L层的K和V）
- 所有GPU总计：$2BSLd$
通信：无
步骤3：输出投影
计算：
-
完整计算：$Y = \text{Concat}(\text{Output heads}) \times W_O$
- 输入（concat后的attention输出）: $(B, d)$
- 权重 $W_O$: $(d, d)$
- 输出 $Y$: $(B, d)$
-
TP切分：权重矩阵按行切分
-
每个GPU上的shape：
- 输入: $(B, \frac{d}{TP})$（每个GPU的attention输出）
- 权重 $W_O$: $(\frac{d}{TP}, d)$（行切分）
- 局部输出: $(B, \frac{d}{TP}) \times (\frac{d}{TP}, d) \to (B, d)$（部分结果）
-
每GPU计算量：$\frac{2B \times d \times d}{TP}$
通信：
- All-Reduce：合并所有GPU的部分结果
- 通信量：$B \times d$ 个元素（相比Prefill少$S$倍）
- 关键问题：通信延迟（latency）成为瓶颈
- Prefill：通信量大但计算更大，通信可被掩盖
- Decode：计算快（访存受限），通信延迟暴露
- 通信/计算比和Prefill一致：$\frac{Bd}{Bd^2/TP} = \frac{TP}{d} \approx 0.1\%$
MLA（Multi-head Latent Attention，DeepSeek V3）
核心思想
通过低秩联合压缩将 K、V 压缩到低维潜在空间存储，同时对 Q 也进行压缩以减少训练激活内存。推理时采用矩阵吸收优化，避免显式解压巨大矩阵，实现极致的 KV Cache 和内存节省。在本分析中，我们简化了这些细节，具体细节可以参照另一篇博客。
DeepSeek V3 关键参数：
- $d = 7168$ (Model Hidden Size)
- $H = 128$ (Attention Heads 数，真实多头结构中)
- $d_v^h = 128$ (Per-head V Dimension，V的单头维度)
- $d_{qk}^h = 192$ (Per-head Q/K Dimension，Q/K的单头维度)
- 其中 $d_{qk}^h = d_C^h + d_R^h = 128 + 64$
- $d_C^h = 128$ (Q, K Non-RoPE head dimension)
- $d_R^h = 64$ (Q, KRoPE head dimension，Decoupled)
- $d^c_{kv} = 512$ (KV 压缩维度，潜在空间中)
- $d^c_q = 1536$ (Query 压缩维度，潜在空间中)
Latent Vector
- 位置：压缩维度中 ($d^c_{kv} = 512$, $d^c_q = 1536$)
- 形式：一维向量，没有显式的多头张量形状
- 存储：存在 KV Cache 中，但由于被压缩成不是一个个头的形式，所以TP的每个GPU都要有一份完整的拷贝。
- 例子：$c_{KV}: (B, S, 512)$ 看不出有128个头的信息
Attention Head的划分与并行
- 位置：Up-Projection时，注意力头会重新出现，此时可以通过TP来切分矩阵。
Prefill 阶段
步骤 1: Q 的 Down-Projection（查询压缩，无头概念）
计算：输入投影到低维潜在空间
$$c^Q_t = h_t \times W^{DQ}$$- 输入 $h_t$: $(B, S, d) = (B, S, 7168)$
- 权重 $W^{DQ}$: $(d, d^c_q) = (7168, 1536)$
- 输出 $c^Q_t$: $(B, S, d^c_q) = (B, S, 1536)$
关键特征：
- 这一步没有任何"头"的概念
- $c^Q_t$ 是一维向量，信息压缩在1536维中，每个GPU都有这个latent vector的副本。
每个GPU计算量: $2BS \times d \times d^c_q$
步骤 2: Q 的 Up-Projection 和多头展开
第一步：线性投影展开维度
$$q_t^C = c^Q_t \times W^{UQ}$$- 输入 $c^Q_t$: $(B, S, 1536)$
- 权重 $W^{UQ}$: $(d^c_q, H \times d_{C}^h) = (1536, 128 \times 128) = (1536, 24576)$
- 输出 $q_t^C$: $(B, S, 24576)$
第二步：Reshape 成多头结构
$$q_t^C = \text{reshape}(q_t^C, (B, S, H, d_C^h)) = (B, S, 128, 128)$$- 现在有 $H=128$ 个头的显式结构
- 每个头的 Non-RoPE 维度：$d_C^h = 128$
第三步：RoPE 部分
$$q^R_t = \text{RoPE}(c^Q_t \times W^{QR})$$- 输入 $c^Q_t$: $(B, S, 1536)$
- 权重 $W^{QR}$: $(d^c_q, H \times d_R^h) = (1536, 128 \times 64) = (1536, 8192)$
- 输出 $q^R_t$: $(B, S, 8192)$ → reshape 为 $(B, S, 128, 64)$
- 应用 RoPE 位置编码
第四步：拼接成完整Q
$$q_{t} = [q_t^C; q_t^R] = (B, S, 128, 192)$$TP 切分（按列切分）：
- 输入 $c^Q_t$: $(B, S, d^c_q)$
- 权重 $W^{UQ}$: $(d^c_q, \frac{H}{TP} \times d_C^h)$，$W^{QR}$: $(d^c_q, \frac{H}{TP} \times d_R^h)$
- 输出 $q_t$: $(B, S, \frac{H \times d_{qk}^h}{TP})$ → reshape 为 $(B, S, \frac{H}{TP}, d_{qk}^h)$
每个GPU计算量 : $\frac{2BS \times d^c_q \times H \times d_{qk}^h}{TP}$
步骤 3: KV 的 Down-Projection（K、V 联合压缩，无头概念）
计算：
$$c^{KV}_t = h_t \times W^{DKV}$$- 输入 $h_t$: $(B, S, d) = (B, S, 7168)$
- 权重 $W^{DKV}$: $(d, d^c_{kv}) = (7168, 512)$
- 输出 $c^{KV}_t$: $(B, S, d^c_{kv}) = (B, S, 512)$
关键特征：
- 这一步也没有任何"头"的概念
- 128个头的K、V信息全部压缩成512维向量
KV Cache 存储：
- 每个TP rank都要重复存储 $c^{KV}_t$: $(B, S, 512)$
- 由于没有头结构，不再可以按头划分到TP
每GPU相比标准MHA：
- MHA KV Cache：$\frac{2 \times B \times S \times d }{TP} = \frac{2 \times B \times S \times 7168}{TP}$
- MLA KV Cache：$B \times S \times d^c_{kv} = B \times S \times 512$
- TP规模越大，KV Cache的重复越多，MLA的优势越低
每GPU计算量：$2 \times B \times S \times d \times d^c_{kv}$
步骤 4: K 的 Up-Projection
第一步：线性投影展开维度
$$k_t^C = c^{KV}_t \times W^{UK}$$- 输入 $c^{KV}_t$: $(B, S, 512)$
- 权重 $W^{UK}$: $(d^c_{kv}, H \times d_C^h) = (512, 128 \times 128) = (512, 16384)$
- 输出 $k_t^C$: $(B, S, 16384)$
第二步：Reshape 成多头结构
$$k_t^C = \text{reshape}(k_t^C, (B, S, H, d_C^h)) = (B, S, 128, 128)$$第三步：RoPE 部分
$$k^R_t = \text{RoPE}(h_t \times W^{KR})$$- 输入 $h_t$: $(B, S, d)= (B, S, 7168)$
- 权重 $W^{KR}$: $(d, d_R^h) = (7168, 64)$
- 输出 $k^R_t$: $(B, S, 64)$ （注意只有一个头）
第四步：拼接成完整K
$$k_t = [k_t^C; k^R_t] = (B, S, 128, 192)$$TP 切分不切分第三步：
- 输入 $c^{KV}_t$: $(B, S, d^c_{kv})$
- 权重 $W^{UK}$: $(d^c_{kv}, \frac{H}{TP} \times d_C^h)$
- 输出 $k_t^C$: $(B, S, \frac{H \times d_C^h}{TP})$
每GPU的计算量: $\frac{2BS \times d^c_{kv} \times H \times d_C^h}{TP} + {2BS \times d \times d_R^h}$
步骤 5: V 的 Up-Projection
$$v_{t} = c^{KV}_t \times W^{UV}$$每个GPU计算：
- 输入 $c^{KV}_t$: $(B, S, 512)$
- 权重 $W^{UV}$: $(d^c_{kv}, \frac{H}{TP} \times d_v^h) = (512, \frac{128}{TP} \times 128)$
- 输出 $v_{t}$: $(B, S, \frac{128}{TP}, 128)$
计算量: $\frac{2BS \times d^c_{kv} \times H \times d_v^h}{TP}$
步骤 6: Attention 计算（128个多头做）
每个 GPU 上的 shape：
- $Q$: $(B, S, \frac{H}{TP}, d_{qk}^h) = (B, S, \frac{128}{TP}, 192)$
- $K$: $(B, S, \frac{H}{TP}, d_{qk}^h) = (B, S, \frac{128}{TP}, 192)$
- $V$: $(B, S, \frac{H}{TP}, d_{v}^h) = (B, S, \frac{128}{TP}, 128)$
计算流程（128个头独立计算）：
-
Score 计算：
$$\text{Score} = \frac{Q \times K^T}{\sqrt{d_{qk}^h}}$$- 形状: $(B, \frac{H}{TP}, S, S)$
-
Attention 权重：
$$\text{Weights} = \text{Softmax}(\text{Score})$$- 形状: $(B, \frac{H}{TP}, S, S)$
-
加权求和：
$$\text{Output} = \text{Weights} \times V$$- 输出: $(B, S, \frac{H \times d_{v}^h}{TP})$
计算量：
- Score: $\frac{H}{TP}\times2BS^2 \times d_{qk}^h$
- Attention×V: $\frac{H}{TP}\times2BS^2 \times d_{v}^h$
- 总计: $\frac{H}{TP} \times 2BS^2 \times (d_{qk}^h + d_{v}^h)$
步骤 6: 输出投影（降维回隐藏维度）
计算：
$$y_t = \text{output} \times W_O$$- 输入（Attention 输出）: $(B, S, H \times d_{v}^h) = (B, S, 16384)$
- 权重 $W_O$: $(H \times d_{v}^h, d) = (16384, 7168)$
- 输出 $y_t$: $(B, S, d) = (B, S, 7168)$
TP 切分（按行切分）：
- 输入: $(B, S, \frac{H \times d_{v}^h}{TP})$
- 权重 $W_O$: $(\frac{H \times d_{v}^h}{TP}, d)$
- 输出: $(B, S, d)$（部分结果）
计算量: $\frac{2BS \times H \times d_{v}^h \times d}{TP}$
通信：
- All-Reduce: $BS \times d = 7168 \times BS$ 元素（与 MHA 相同）
Prefill 阶段总结
| 步骤 | 涉及的头 | 计算量 (每 GPU) | 说明 |
|---|---|---|---|
| Q Down | ❌ 无 | $2BS \times d \times d^c_q$ | Query 压缩，无头概念 |
| Q Up | ✅ 128 | $\frac{2BS \times d^c_q \times H \times (d_C^h + d_R^h)}{TP}$ | Query 解压，多头展开 + RoPE |
| KV Down | ❌ 无 | $2BS \times d \times d^c_{kv}$ | KV 联合压缩，无头概念，KV Cache每个TP都有副本 |
| K Up | ✅ 128 | $\frac{2BS \times d^c_{kv} \times H \times d_C^h}{TP} + {2BS \times d \times d_R^h}$ | K 解压，多头展开 + RoPE |
| V Up | ✅ 128 | $\frac{2BS \times d^c_{kv} \times H \times d_v^h}{TP}$ | V 解压，头聚合形式 |
| Attention | ✅ 128 | $\frac{2BS^2 \times (d_{qk}^h + d_v^h)}{TP}$ | Score 和加权求和 |
| Output | ✅ 128→1 | $\frac{2BS \times H \times d_v^h \times d}{TP}$ | 多头合并，投影回隐藏维度 |
Decode 阶段
步骤 1: Q 的 Down-Projection（查询压缩，无头概念）
与Prefill相同
步骤 2: Q 的 Up-Projection 和多头展开
第一步：线性投影展开维度
$$q_t^C = c^Q_t \times W^{UQ}$$- 输入 $c^Q_t$: $(B, 1, 1536)$
- 权重 $W^{UQ}$: $(d^c_q, H \times d_{C}^h) = (1536, 128 \times 128) = (1536, 24576)$
- 输出 $q_t^C$: $(B, 1, 24576)$
第二步：Reshape 成多头结构
$$q_t^C = \text{reshape}(q_t^C, (B, 1, H, d_C^h)) = (B, 1, 128, 128)$$- 现在有 $H=128$ 个头的显式结构
- 每个头的 Non-RoPE 维度：$d_C^h = 128$
第三步：吸收矩阵$W^{UK}$到$$q_t^C$$
$$q_t^{C\_absorb} = q_t^C \times {W^{UK}}^T$$- 输入$$q_t^C :(B, 1, 128, 128)$$
- 权重 $W^{UK}$: $(d^c_{kv}, H,d_C^h) = (512, 128,128)$
- 输出$$q_t^{C\_absorb}:(B, 1, H, d^c_{kv})=(B, 1, 128, 512)$$
第四步：RoPE 部分
$$q^R_t = \text{RoPE}(c^Q_t \times W^{QR})$$- 输入 $c^Q_t$: $(B, 1, 1536)$
- 权重 $W^{QR}$: $(d^c_q, H \times d_R^h) = (1536, 128 \times 64) = (1536, 8192)$
- 输出 $q^R_t$: $(B, 1, 8192)$ → reshape 为 $(B, 1, 128, 64)$
- 应用 RoPE 位置编码
第五步：拼接成完整Q
$$q_{t} = [q_t^{C\_absorb}; q_t^R] = (B, 1, 128, 576)$$TP 切分（按列切分）：
- 输入 $c^Q_t$: $(B, 1, d^c_q)$
- 权重 $W^{UQ}$: $(d^c_q, \frac{H}{TP} \times d_C^h)$，$W^{UK}$: $(d^c_{kv}, \frac{H}{TP}\times d_C^h)$,$W^{QR}$: $(d^c_q, \frac{H}{TP} \times d_R^h)$
- 输出 $q_t$: $(B, 1, \frac{H \times d_{qk}^h}{TP})$ → reshape 为 $(B, 1, \frac{H}{TP}, d_{qk}^h)$
每个GPU计算量 （忽略RoPE位置编码）: $\frac{2B \times d^c_q \times H \times (d_{qk}^h+d_R^h)+2B \times d^c_{kv} \times H \times d_{qk}^h}{TP}$
步骤 3: KV 的 Down-Projection
第一步与Prefill一致，计算：
$$c^{KV}_t = h_t \times W^{DKV}$$- 输入 $h_t$: $(B, 1, d) = (B, 1, 7168)$
- 权重 $W^{DKV}$: $(d, d^c_{kv}) = (7168, 512)$
- 输出 $c^{KV}_t$: $(B, 1, d^c_{kv}) = (B, 1, 512)$
利用先前压缩latent vector的KV Cache 存储：
第二步：RoPE 部分
$$k^R_t = \text{RoPE}(h_t \times W^{KR})$$- 输入 $h_t$: $(B, 1, d)= (B, 1, 7168)$
- 权重 $W^{KR}$: $(d, d_R^h) = (7168, 64)$
- 输出 $k^R_t$: $(B, 1, 64)$ （注意只有一个头）
第三步：拼接成完整K
$$k_t = [c^{KV}_t; k^R_t] = (B, 1, 576)$$TP 不能切分该步骤
每GPU的计算量: ${2B \times d \times (d^c_{kv} + d_R^h)}$
步骤 4: Attention 计算（MQA）
每个 GPU 上的 shape：
- $Q$: $(B, 1, \frac{H}{TP}, d_{kv}^c+d^h_R) = (B, 1, \frac{128}{TP}, 576)$
- $K$: $(B, S, d_{kv}^c+d^h_R) = (B, S, 576)$
- $V$: $(B, S, d_{kv}^c) = (B, S, 512)$
计算流程（128个头独立计算）：
-
Score 计算：
$$\text{Score} = \frac{Q \times K^T}{\sqrt{d_{qk}^h}}$$- 形状: $(B, \frac{H}{TP}, 1, S)$
-
Attention 权重：
$$\text{Weights} = \text{Softmax}(\text{Score})$$- 形状: $(B, \frac{H}{TP}, 1, S)$
-
加权求和：
$$\text{Output} = \text{Weights} \times V$$- V此时不需要被升维还原成多头形式
- 输出: $(B, \frac{H}{TP}, 1, d_{kv}^c)$
每个GPU计算量：
- Score: $\frac{H}{TP}\times2BS \times (d_{kv}^c+d^h_R) $
- Attention×V: $\frac{H}{TP}\times2BS \times d_{kv}^c$
- 总计: $\frac{H}{TP}\times2BS \times (2d_{kv}^c+d^h_R) $
步骤 5: V 的 Up-Projection应用到MQA结果
$$Out_{up} = \text{Output} \times W^{UV}$$每个GPU计算：
- 输入 上一步MQA的Output: $(B, \frac{H}{TP}, 1, d_{kv}^c)=(B, \frac{H}{TP}, 1, 512)$
- 权重 $W^{UV}$: $(\frac{H}{TP}, d^c_{kv}, d_v^h) = (\frac{128}{TP} ,512, 128)$
- 输出 $Out_{up}$: $(B, 1, \frac{128}{TP}, 128)$
计算量: $\frac{2B \times d^c_{kv} \times H \times d_v^h}{TP}$
步骤 6: 输出投影（降维回隐藏维度）
$$y_t = Out_{up} \times W_O$$- 输入$Out_{up}$: $(B, 1, H \times d_{v}^h) = (B, 1, 16384)$
- 权重 $W_O$: $(H \times d_{v}^h, d) = (16384, 7168)$
- 输出 $y_t$: $(B, 1, d) = (B, 1, 7168)$
TP 切分：
- 输入: $(B, 1, \frac{H \times d_{v}^h}{TP})$
- 权重 $W_O$: $(\frac{H \times d_{v}^h}{TP}, d)$
- 输出: $(B, 1, d)$（部分结果）
计算量: $\frac{2B \times H \times d_{v}^h \times d}{TP}$
通信：
- All-Reduce: $B \times d = 7168 \times B$ 元素
Decode 阶段总结
| 步骤 | 涉及的头 (每 GPU) | 计算量 (每 GPU) (B: Batch, S: History Len) | 关键特征说明 |
|---|---|---|---|
| Step 1: Q Down | ❌ 无 | $2B \times d \times d^c_q$ | 查询压缩 输入 $h_t$ 投影到 latent 空间，无头概念。 |
| Step 2: Q Up | ✅ 128 | $\frac{2B \cdot d^c_q \cdot H \cdot d_C^h}{TP} + \text{AbsorbCost}^*$ $+ \frac{2B \cdot d^c_q \cdot H \cdot d_R^h}{TP}$ | Query 解压 + 吸收 K-Up 展开多头 Q，并吸收 $W^{UK}$ 权重； 生成 RoPE 部分并拼接。Q 处于"Latent KV"空间。 |
| Step 3: KV Down | ❌ 无 | $2B \times d \times (d^c_{kv} + d_R^h)$ | KV 生成 (单头) 仅生成 latent $c^{KV}$ 和 $k^R$； 不进行多头展开，直接存入 KV Cache。 |
| Step 4: Attention | ✅ 128 (Q) vs 1 (KV) | Score: $\frac{H}{TP} \times 2B \cdot S \times (d^c_{kv} + d_R^h)$ Agg: $\frac{H}{TP} \times 2B \cdot S \times d^c_{kv}$ | MQA (Multi-Query Attention) 多头 Q 关注共享的 Latent KV Cache； 计算复杂度随历史长度 $S$ 线性增长。 |
| Step 5: V Up | ✅ 128 | $\frac{2B \cdot d^c_{kv} \cdot H \cdot d_v^h}{TP}$ | V 解压 (后置) 仅对 Attention 的输出结果进行升维； 恢复多头 Value 的表达能力。 |
| Step 6: Output | ✅ 128→1 | $\frac{2B \cdot H \cdot d_v^h \cdot d}{TP}$ | 输出投影 多头合并，投影回隐藏维度 $d$； 需进行 All-Reduce 通信。 |
- AbsorbCost: 指将 $W^{UK}$ 融合进 Q 的计算开销，通常是 $\frac{2B \cdot H \cdot d_C^h \cdot d^c_{kv}}{TP}$，取决于具体算子融合实现。
Prefill vs Decode
| 维度 | Prefill 阶段 | Decode 阶段 | 核心差异点 |
|---|---|---|---|
| Attention 类型 | MHA (Standard) | MQA (Multi-Query) | Decode直接对Latent vector做MQA |
| KV Cache 内容 | $c^{KV}$ (Latent) | $c^{KV}$ (Latent) | 存储内容完全一致，Decode 复用 Prefill 生成的 Cache，无需转换。 |
| KV Up-Proj 时机 | Attention 之前 | Attention 之后 (V) 融合进 Q (K) | Prefill 显式还原所有 KV 头；Decode 避免还原历史 KV，只在结果上还原。 |
| K 的形态 | 多头 $K$ $(B, S, H, d_{qk}^h)$ | 单头 Latent $K$ $(B, S, 1, d^c_{kv}+d_R^h)$ | Prefill 的 K 是物理多头；Decode 的 K 是Latent vector加上RoPE、逻辑广播。 |
| 计算复杂度 | $O(S^2)$ | $O(S)$ | Prefill仍然可以考虑TP，但Decode使用TP会难以掩盖通信开销，此外KV Cache的效率很低。 |
MHA vs MLA 对比总结
- MHA由于头是独立的，因此用TP只需要最后的通信开销。
- MLA的Latent Vector需要在TP的每个rank做副本的保留，因此使用TP的效率会较低。
- MLA的访存瓶颈被缓解，TP通信的影响会更明显。
- MLA推理实践中很常用的是DP- Attention。
- MLA的Latent Vector可以压缩KV Cache的大小，因此对访存密集的Decode会很有优势。
TP对FFN的影响
Prefill阶段和Decode阶段是一致的
区别是Decode的S为1
步骤1：第一层投影（$d \to d_{ff}$）
计算：
-
完整计算：$H = X \times W_1$
- 输入 $X$: $(B, S, d)$
- 权重 $W_1$: $(d, d_{ff})$（第一层FFN权重）
- 输出 $H$: $(B, S, d_{ff})$
-
TP切分：权重矩阵按列切分
-
每个GPU上的shape：
- 输入 $X$: $(B, S, d)$（广播到所有GPU）
- 权重 $W_1$: $(d, \frac{d_{ff}}{TP})$（列切分）
- 输出 $H$: $(B, S, \frac{d_{ff}}{TP})$
-
每GPU计算量：$\frac{2BS \times d \times d_{ff}}{TP}$
通信：无
步骤2：激活函数
计算：各GPU独立计算（SwiGLU/GELU）
- 每个GPU上的shape：输入和输出都是$(B, S, \frac{d_{ff}}{TP})$
- 计算量：$O(BS \times \frac{d_{ff}}{TP})$（相对较小）
通信：无
步骤3：第二层投影（$d_{ff} \to d$）
计算：
-
完整计算：$Y = \text{Activation}(H) \times W_2$
- 输入（激活后）: $(B, S, d_{ff})$
- 权重 $W_2$: $(d_{ff}, d)$（第二层FFN权重）
- 输出 $Y$: $(B, S, d)$
-
TP切分：权重矩阵按行切分
-
每个GPU上的shape：
- 输入: $(B, S, \frac{d_{ff}}{TP})$
- 权重 $W_2$: $(\frac{d_{ff}}{TP}, d)$（行切分）
- 局部输出: $(B, S, \frac{d_{ff}}{TP}) \times (\frac{d_{ff}}{TP}, d) \to (B, S, d)$（部分结果）
-
每GPU计算量：$\frac{2BS \times d_{ff} \times d}{TP}$
通信：
- All-Reduce：合并TP切分的结果
- 通信量：$BS \times d$ 个元素（与Attention相同）
- FFN后立即与Attention结果Add，通信模式一致
通信开销对比
各并行策略的通信量推导总结
约定：本节表格中的“需要同步的数据”指逻辑上的张量大小（即需要被传输的激活值元素个数）。 实际传输量指物理链路上每块 GPU 发送/接收的数据总量，需乘以与算法相关的常数因子。
| 并行策略 | 通信模式 | 需要同步的数据 (逻辑大小) | 实际通信量 (每GPU, 数量级) | 关键推导逻辑 |
|---|---|---|---|---|
| TP | All-Reduce | $N \times d$ | $\approx 2\frac{TP-1}{TP} \times N d$ | 权重按行切分，每卡计算出部分和，需通过 Ring All-Reduce 累加。$N$ 为该算子的 token 总数。 |
| EP | All-to-All（和DP- Attention结合） | $N \times k \times d$ | $\approx 4 \times \frac{N k d}{EP} \times \frac{EP-1}{EP}$ | MoE 层包含 两次 All-to-All (Dispatch + Combine)。每次逻辑数据量为 $N \times k \times d$ (每个 token 选 $k$ 个专家)。最后的系数是因为不需要和自己通信。 |
| PP | P2P Send/Recv | $N \times d$ | $\approx N \times d$ (单向) | Stage 边界的点对点传输。前一个 Stage 计算完 $N$ 个 token 的激活值，发送给下一个 Stage。 |
| DP | 无 | 0 | 0 | 在纯推理阶段，数据并行副本之间无通信。 |
关键变量说明：
- Prefill 阶段：$N \approx B \times S$ (Batch $\times$ Sequence Length)
- Decode 阶段：$N \approx B$ (Batch $\times$ 1)
- 结论：在 Decode 阶段，所有策略的通信量相比 Prefill 都缩小了约 $S$ 倍。
Prefill vs Decode 的 TP / EP / PP 通信量对比
下表显式区分了 Prefill 和 Decode 阶段在不同并行策略下的通信特征。忽略 DP，仅关注单层/单次操作开销。
- 记号：$B$ (Batch Size), $S$ (Seq Len), $d$ (Hidden Dim), $k$ (Top-k experts), $d_{act}$ (Activation Dim).
| 并行策略 | 阶段 | 逻辑通信数据大小 (单层) | 粗略通信量 (每GPU) | 特性分析 |
|---|---|---|---|---|
| TP | Prefill | $BS \times d$ | $\approx 1.75 \times BS d$ (TP=8) | 通信量大，但在高算术强度下容易被计算掩盖。 |
| Decode | $B \times d$ | $\approx 1.75 \times B d$ | 通信量小，但计算量也小，延迟直接暴露，受限于 NVLink 带宽和延迟。 | |
| EP | Prefill | $BS \times k \times d$ | $\approx 4 \frac{BSkd}{EP}\times \frac{EP-1}{EP}$ | 跨节点的 All-to-All 通信。Prefill 阶段数据量大，极易触碰跨机互联带宽瓶颈。 |
| Decode | $B \times k \times d$ | $\approx 4 \frac{Bkd}{EP}\times \frac{EP-1}{EP}$ | 数据量虽小，但 All-to-All 的 启动延迟 (Latency) 和网络拥塞控制成为主要瓶颈。 | |
| PP | Prefill | $BS \times d_{act}$ | $\approx BS d$ | 巨大的激活值传输导致流水线 Bubble 难以填充，且首个 Token 延迟 (TTFT) 显著增加。 |
| Decode | $B \times d_{act}$ | $\approx B d$ | 通信量较小，主要问题是流水线带来的 串行依赖 导致的时延累积。 |
- TP的实现简单，对MHA支持较好。但需要良好的通信性能。每一层都需要All Reduce通信。
- EP用于MoE架构。优点是不需要GPU加载所有的专家，符合MoE减少内存加载的初衷。在DeepSeek中配合DP attention较为常见。每个MoE层都需要All to All通信，也需要关注。
- PP只有在切分层需要通信。计算通信比会更高，适合显存不够但通信性能差的场景。
Static vs Continuous Batching
实际生产系统中，Batch 不再是静态的 $(B, S)$，而是动态的流。
| 维度 | Static Batching (简化模型) | Continuous Batching (实际模型) | 核心差异 |
|---|---|---|---|
| 定义 | 假设 Batch 内所有请求状态一致（全 Prefill 或全 Decode）。 | 允许 Batch 内混合：$B_p$ 个请求在做 Prefill，$B_d$ 个在做 Decode。 | 资源利用率更高，但这增加了调度复杂性。 |
| 总 Token 数 | Prefill: $T = B \times S$ Decode: $T = B$ |
$T_{total} = \sum_{i=1}^{B_p} S_i + B_d$ | 通信量不再由 $B$ 决定，而是由当前 Step 的 有效 Token 总数 $T_{total}$ 决定。 |
| TP 通信 | 每层固定传输 $BSd$ 或 $Bd$ | 每层传输 $\propto T_{total} \times d$ | $T_{total}$ 在不同 Step 间剧烈波动，带宽需求呈现脉冲式。 |
| EP 通信 | 每层固定传输 $BSkd$ 或 $Bkd$ | 每层传输 $\propto T_{total} \times k \times d$ | EP 的 All-to-All 负载随 $T_{total}$ 波动，且专家负载可能不均衡。 |
Continuous Batching 的系统设计影响
引入 Continuous Batching 后，对并行通信系统的设计带来以下挑战：
-
通信带宽的动态性：
- 由于 $T_{total}$ 取决于新进来请求的 Prefill 长度 $S_i$，某个 Step 的通信量可能突然暴增。
- 网络设计必须按 峰值 (Peak) 而非平均值规划，否则大 Prefill 进来时会造成严重的 Head-of-Line Blocking。
-
调度策略权衡：
- Prefill 优先：最大化 $T_{total}$，提高计算密度（算术强度），掩盖 TP/EP 通信延迟。缺点是显存瞬间暴涨，可能挤占 Decode 空间。
- Decode 优先：保证低延迟，但 $T_{total}$ 较小，计算无法掩盖通信，GPU 即使利用率低也必须等待通信完成。
-
内存（KV Cache）管理：
- KV Cache 占用量公式变为：$\text{Mem}_{KV} = (\sum S_{curr}) \times L \times d_{kv} \times \text{Layers}$。
- TP/PP 必须保证所有 Device 上的显存分配一致，而 EP 可能导致不同专家的 KV 分布不均（但在 MLA/MHA 共享 KV 的架构下，KV 通常不做 EP 切分）。
Chunked Prefill 对通信与流水线的影响
为了解决长序列 Prefill 阻塞 Decode 的问题，通常引入 Chunked Prefill：将长序列 $S_{total}$ 拆分为多个 $S_{chunk}$ (e.g., 512 tokens) 分批处理。
A. 对 TP / EP 通信量的影响
- 总量不变，峰值削峰：
- 不分 Chunk：单次通信量 $B \times S_{total} \times d$，巨大的数据包可能导致网络拥塞。
- 分 Chunk：单次通信量 $B \times S_{chunk} \times d$。虽然总传输字节数不变，但单次通信更小，更容易与计算 Kernel 进行 Overlap (重叠)。
B. 对 PP (流水线并行) 的决定性作用
- 消除流水线气泡 (Bubble)：
- 无 Chunk：一个长 Prefill 请求进入 Stage 1，后续所有 Decode 请求必须等待它完全流过整个 Pipeline，造成巨大的时延抖动。
- 有 Chunk：Prefill 的 Chunk 可以与 Decode 请求 交错 (Interleaved) 进入流水线。
- 效果：虽然单个 Prefill 请求的完成时间（TTFT）没有显著减少，但系统的 平均吞吐量 和 Decode 请求的 P99 延迟 得到显著优化。