vllm源码剖析13-vLLM 分布式推理-张量并行详解

文章目录

• • 一 AllReduce 原理
  • • 1.1 Ring-AllReduce 算子原理
    • • 1.1.1 Reduce-Scatter
      • 1.1.2 All-Gather
    • 1.2 Ring AllReduce 的通信成本（重）
  • 二 Transformer 模型的张量并行
  • • 2.1 线性层权重不同切分方式
    • 2.2 MLP 层的张量并行
    • • 2.2.1 拆分原理
      • 2.2.2 拆分方式
      • 2.2.3 MLA 层的通讯量分析(重)
    • 2.3 MHA 层的张量并行
    • • 2.3.1 拆分原理
      • 2.3.2 MHA 层的通讯量分析
      • 2.3.3 attention 中的张量并行实例
    • 2.4 Embedding 层的张量并行
    • • 2.4.1 输入嵌入层
      • 2.4.2 输出嵌入层
  • 三 vLLM 中的张量并行
  • • 3.1 vLLM 中张量并行如何使用
    • 3.2 vLLM 中的分布式资源管理
    • • 模拟分布式并行分分组算法
    • 3.3 vLLM 的并行线性层
    • • ColumnParallelLinear 类源码剖析(列并行)
      • RowParallelLinear 类源码剖析
      • QKVParallelLinear 源码剖析
    • 3.4 vLLM 分布式推理的流程总结
    • • 阶段一：初始化与模型加载
      • 阶段二：分布式前向传播 (Forward Pass)
      • 阶段三：Logits 聚合与 Token 采样

张量并行（TP，Tensor Parallelism）的核心思路，是把模型中的大矩阵计算拆到多张 GPU 上完成。这样可以降低单张 GPU 的显存压力；当计算量足够大、通信开销可以被摊薄时，也可能提升 forward 吞吐。
在 vLLM 的 Megatron 风格 TP 中，常见的切分位置主要有三类：

• Embedding 层：例如 VocabParallelEmbedding，按词表维度切分 embedding 权重。每张卡只负责一部分 token，forward 后通过 All-Reduce 汇总完整 embedding。
• 线性层：例如 MLP 里的 MergedColumnParallelLinear、RowParallelLinear，以及 Attention 里的 QKVParallelLinear 和 o_proj。Column Parallel 通常切输出维，Row Parallel 通常切输入维。
• Attention：Q/K/V 通常按 head 维分到不同 rank 上，各 rank 计算本地 head 的 attention，最后的 o_proj 通常通过 Row Parallel 和 All-Reduce 合并结果。

后面我们就按 VocabParallelEmbedding、Column Parallel、Row Parallel 和 LM Head 这几类层，分别看它们在 forward 中什么时候触发 All-Reduce、All-Gather 或 gather。 下面将结合具体并行层，说明这些集合通信在 forward 中的触发位置，以及它们如何配合 TP 的切分方式完成计算。

一 AllReduce 原理

All-Reduce 的目标，是将所有进程上的数据通过特定操作（如求和、取最大值）聚合后，把结果同步到每一个进程。常见底层实现包括 Ring AllReduce 和 Tree AllReduce。其最终目标，就是让每块 GPU 上的数据都变成汇总/归约后的同一个结果。 在 vLLM 的 TP 中，All-Reduce 典型出现在 Row Parallel 线性层和 Embedding 层：各 rank 先算出自己的部分结果，再通过 All-Reduce 求和，得到完整输出。具体我们会在本节的之后内容中讲到

1.1 Ring-AllReduce 算子原理

Ring-AllReduce 的实现其实分为两个过程 Reduce-Scatter 和 All-Gather。

1.1.1 Reduce-Scatter

以 4 个 GPU 设备为例，可将它们组织成一个逻辑环，使每个 GPU 只与相邻 GPU 交换数据。待归约的数据被切成 4 个 chunk，每轮通信时，各 GPU 同时发送一个 chunk、接收一个 chunk，并将收到的 chunk 与本地对应位置的 chunk 做逐元素相加。 第一次通信和累加完成后，某些位置的 chunk 已包含两个 GPU 的部分结果；这些更新后的 chunk 会在后续轮次中继续传递和累加。经过 N-1 轮（此处 N=4，即 3 轮）后，每个 GPU 会得到一个已完成全局累加的结果分片。 这里每次通信的数据量是一个 chunk。若原始数据量为 K，进程数为 N，则每个 chunk 的大小约为

K

/

N

K/N

K/N。Reduce-Scatter 阶段结束时，每个设备只持有完整结果的一部分；后续 All-Gather 阶段会将这些分片继续传播，使每个设备最终都获得完整结果。

第二次累加完成后的示意图如下，同样，被更新的数据块，会作为下一次传递和累加的起点，继续参与下一轮的通信和计算。

第三次累加完成后的示意图如下： 经过 3 次环形传递和规约后，每块 GPU 上都有一块数据拥有了对应位置完整的累加聚合（下图中红色块）。此时，Reduce-Scatter 通信阶段结束，进入 All-Gather 通信阶段。目标是将红色块继续沿环传播，并填充到其余 GPU 对应的位置上，使所有 GPU 最终都拥有全部数据。

1.1.2 All-Gather

All-Gather 通信操作依然遵循相邻 GPU 对应位置进行通讯的原则，但这一步不再做相加，而是将已归约好的分块拷贝到下一跳对应的位置上。All-Gather 以红色块作为起点，第一轮传递和填充完成后的示意图如下所示：

同样的经过 3 轮更新，使得每块 GPU 上都汇总到了完整的数据，变成如下形式:

1.2 Ring AllReduce 的通信成本（重）

假设有

N

N

N 个设备，原始数据总大小为

K

K

K，在一次 AllReduce 过程中，进行了

N

−

1

N-1

N−1 次 Scatter-Reduce 操作和

N

−

1

N-1

N−1 次 Allgather 操作，又因为每一次操作所需要传递的数据大小为

K

/

N

K/N

K/N，所以整个 AllReduce 过程所传输的数据大小为

2

(

N

−

1

)

∗

K

/

N

2(N-1) * K/N

2(N−1)∗K/N。随着 N 的增大，Ring AllReduce 通信算子的通信量可以近似为 2K。

在张量并行加速时，Ring AllReduce 的吞吐通常受到环上最慢链路和实现开销的共同影响。每次传输的数据块只有

K

/

N

K/N

K/N，所以当

N

N

N 增大时，单轮数据块变小。对于大张量，Ring AllReduce 容易形成流水线并行，通信效率通常较高；但对于小张量，固定调度开销和协议开销的占比会明显上升，实际有效带宽利用率会下降。

另外，Ring AllReduce 需要完成

2

×

(

N

−

1

)

2 \\times (N-1)

2×(N−1) 次通信（两个阶段各

N

−

1

N-1

N−1 次）。例如 8 GPU 需要 14 轮通信，即 Reduce-Scatter 和 All-Gather 各 7 轮。小张量本身的计算时间通常很短，但多轮通信带来的延迟累积会成为瓶颈。因此，Ring AllReduce 一般不适合小张量频繁同步、且对延迟要求很高的场景。

vLLM 在 CUDA all-reduce 路径中并不简单地固定使用一种传统 Ring-AllReduce。源码中会根据当前并行组、环境开关、world size、硬件拓扑、P2P 能力、张量大小、dtype、连续性等多种条件，尝试不同的通信实现，选择更合适的 all-reduce 路径，以降低延迟或提高带宽利用率。 以下伪代码用来抽象说明低延迟 AllReduce 的核心思想：各 GPU 通过预先建立的通信资源访问其他 rank 的数据，并完成逐元素求和，使每个 GPU 最终得到与 AllReduce(sum) 等价的全局规约结果。

def low_latency_all_reduce(x):
"""
逻辑效果等价于 AllReduce(sum)。
每个 GPU 最终都会得到所有 rank 输入张量的逐元素求和结果。
"""
# 所有 GPU 事先完成必要的通信资源初始化，
# 例如注册 buffer、创建 workspace，或建立 symmetric memory 句柄。
peer_buffers = get_all_peer_buffers()

# 同步和就绪检查省略。
out = zeros_like(x)

for buf in peer_buffers:
# 抽象表示：读取各个 rank 的输入数据，
# 并在 GPU 上完成逐元素累加。
out += load_from(buf)

return out

从逻辑效果上看，每个 GPU 都会执行一次等价的规约过程，并得到与 AllReduce(sum) 相同的输出。和 Ring AllReduce 相比，这类低延迟路径不强调多轮分块传递，而是利用预先建立的通信资源、共享缓冲区等资源，尽量减少通信轮次和 kernel 调度次数，从而降低小张量频繁同步时的总延迟

二 Transformer 模型的张量并行

decoder-only 架构的 LLM 中标准的 transformer 层如图 2 所示，其由一个自注意力（self-attention）模块和一个两层的多层感知机 (MLP)组成，可在这两个模块中分别引入模型并行（也叫张量并行）技术。 基于 transformer 网络 pytorch 代码的基础上，通常只需添加几个同步操作代码（synchronization primitives），就可实现一个简单的模型并行方案。下文我将会描述 Megatron-LM 的张量并行算法原理，以及在 transformer 模型中的应用。

2.1 线性层权重不同切分方式

张量并行的底层逻辑，可以先从矩阵乘法的拆分计算来理解。下面分别介绍线性层中的列并行和行并行，以及它们对应的矩阵形式。 其中

X

X

X

是输入，

是输入，

是输入，A

是权重矩阵。按这个数学形式，

是权重矩阵。按这个数学形式，

是权重矩阵。按这个数学形式，A$的第一维对应输入维，第二维对应输出维。后面说的“列并行”和“行并行”，都基于这个矩阵形式来理解

线性层是 Transformer 中最主要的 GEMM 来源之一，既出现在 MLP 中，也出现在 Attention 的 Q/K/V projection 和 output projection 中。开启张量并行后，vLLM 会让不同 TP rank 加载各自负责的权重分片，并在 forward 时分别执行本地 GEMM。 按不同方式切分权重后（TP=2）的线性层推理 forward 操作的可视化对比图如下图所示：

• 权重 $A$
• 权重 $A$

2.2 MLP 层的张量并行

MLP 的张量并行实现相对 Self-Attention 简单。

2.2.1 拆分原理

MLP 模块的第一个操作是通用矩阵乘法 (GEMM)，随后是一个 GeLU 非线性激活函数，计算公式如下所示。

Y

=

G

e

L

U

(

X

A

)

Y = GeLU(XA)

Y=GeLU(XA) 后续最新的 Llama 及 qwen 系列的 llm 用的激活函数是

S

i

L

U

SiLU

SiLU。 为了实现 GEMM 的并行计算。

将第一个线性层的权重按照列并行方式切分后，第二个线性层的权重矩阵 B 自然沿着行方向拆分，使其能够直接处理来自 GeLU 层的输出而无需任何通信，如图 3a 所示。

2.2.2 拆分方式

import numpy as np

def gelu(x):
"""GeLU 激活函数 (使用 tanh 近似)"""
return 0.5 * x * (1 + np.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * np.power(x, 3))))

# 1. 初始化参数
np.random.seed(42)
B = 4 # Batch size
H = 8 # Hidden dimension (输入维度)
D = 4 # Output dimension (输出维度)

# 创建输入矩阵 X (B, H) 和 权重矩阵 A (H, D)
X = np.random.randn(B, H)
A = np.random.randn(H, D)

# ==========================================
# 方法 1: 单机基准计算 (不拆分)
# ==========================================
# 直接计算 Y = GeLU(X @ A)
target_output = gelu(np.dot(X, A))
print(f"基准输出形状: {target_output.shape}")

# ==========================================
# 方法 2: 模拟行拆分并行 (Row Parallelism)
# ==========================================
print("\\n— 开始模拟行拆分并行 —")

# 假设我们有 2 个设备 (GPU)，将 A 按行切分，将 X 按列切分
# split_size = H // 2

# [切分权重矩阵 A] -> A1, A2
# A 的形状是 (H, D)，按行切分 (axis=0)
A1 = A[:H//2, :]
A2 = A[H//2:, :]
print(f"设备1 权重 A1 形状: {A1.shape}")
print(f"设备2 权重 A2 形状: {A2.shape}")

# [切分输入矩阵 X] -> X1, X2
# X 的形状是 (B, H)，为了配合 A 的行切分，X 需要按列切分 (axis=1)
X1 = X[:, :H//2]
X2 = X[:, H//2:]
print(f"设备1 输入 X1 形状: {X1.shape}")
print(f"设备2 输入 X2 形状: {X2.shape}")

# [并行计算]
# 每个设备独立计算自己的部分积: Xi * Ai
# 结果形状均为 (B, D)
Y1_partial = np.dot(X1, A1)
Y2_partial = np.dot(X2, A2)

# [同步点 / All-Reduce]
# 在 GeLU 之前，必须将各设备的部分结果相加
# 对应公式: XA = X1A1 + X2A2
Y_combined = Y1_partial + Y2_partial

# [应用非线性激活函数]
# 聚合后才能执行 GeLU
parallel_output = gelu(Y_combined)

# ==========================================
# 验证结果
# ==========================================
# 检查并行计算结果与基准结果是否一致
is_close = np.allclose(target_output, parallel_output)
print(f"\\n结果验证: {'成功' if is_close else '失败'}")
print(f"两者误差 (Max Diff): {np.max(np.abs(target_output – parallel_output))}")

# 演示如果直接在局部做 GeLU 再相加是错误的 (数学原理验证)
wrong_output = gelu(Y1_partial) + gelu(Y2_partial)
print(f"\\n错误做法 (先 GeLU 后聚合) 误差: {np.max(np.abs(target_output – wrong_output))}")
print("结论: GeLU(X1A1 + X2A2) != GeLU(X1A1) + GeLU(X2A2)")

• f 的前向推理（forward）计算：对应列并行线性层前后的 identity 操作。输入 $X$
• g 的前向推理（forward）计算：对应行并行线性层后的归约操作。每块 GPU 完成本地 GEMM 后，分别得到局部输出 $Z_1$

MLP 的张量并行过程的形状变换公式拆解如下：

• $b$
• $s$
• $h$
• $i$
• GPU0： $[b,s,h]\times [h,2i/2]\to [b,s,2i/2]$
• GPU1： $[b,s,h]\times [h,2i/2]\to [b,s,2i/2]$

"""
f operator 实现：
– 前向传递：直接返回输入（恒等运算）
– 反向传递：对梯度执行 All-Reduce
对应的 g operator 行为对称：
– 前向传递：执行 All-Reduce
– 反向传递：直接返回梯度（恒等运算）

Implementation of f operator. g is similar to f with
identity in the backward and all-reduce in the forward functions.
"""class f(torch.autograd.Function):
@staticmethoddef forward(ctx, x):
return x # 前向传递无通信 @staticmethoddef backward(ctx, grad_output):
all_reduce(grad_output) # 反向传递触发 All-Reducereturn grad_output

class g(torch.autograd.Function):
@staticmethoddef forward(ctx, x):
all_reduce(x) # 前向传递触发 All-Reduce @staticmethoddef backward(ctx, grad_output):
return grad_output # 反向传递无通信

总结：当列并行与行并行级联使用时，在常见的 ColumnParallelLinear(gather_output=False) 接 RowParallelLinear(input_is_parallel=True) 的组合下，前级输出本来就是按最后一维切分后的局部结果，正好可以作为后级行并行所需的局部输入。 需要注意的是，先列后行的并行级联方式只是表示中间无需通信，并不意味着整个模块没有集合通信。RowParallelLinear 默认 reduce_results=True，通常会在自身输出侧执行 AllReduce，将各 GPU 的部分结果求和成完整输出。如果它的前面或后面还要与其他并行方式连接，也可能根据张量布局继续使用 AllGather、AllReduce、ReduceScatter 等集合通信操作。 先列后行的并行级联的 MLP 前向传播的可视化连接如下图所示: vLLM 框架中 MLP 模块针对线性层的张量并行，它的实际代码实现也是先列后行线性层：

2.2.3 MLA 层的通讯量分析(重)

总结：MLP 层在 forward(前向推理时) 做一次 All-Reduce 操作，在 backward(前向推理时) 做一次 All-Reduce 操作。而 All-Reduce 的过程分为两个阶段，Reduce-Scatter 和 All-Gather，每个阶段的通讯量是相等的。假设输入张量大小为 [b, s, h]，数据类型为 fp16，则每次 All-Reduce 操作通讯量为 2bsh。 模型训练和推理阶段，MLP 层的张量并行通信量如下所示：

• 模型训练时，包含前向传播和反相传播两个过程，即两次 All-Reduce 操作，所以 MLP 层的总通讯量为：$4bsh$
• 模型推理时，只有前向传播过程，即一次 All-Reduce 操作，所以 MLP 层的总通讯量为：$2bsh$

2.3 MHA 层的张量并行

2.3.1 拆分原理

多头注意力模块的结构如下图所示，可以看出，在设计上，MHA 层对于每个头（head），就有都有独立的 q/k/v 三个线性变换层以及对应的 self-attention 计算结构，然后将每个 head 输出的结果做拼接 concat，最后将拼接得到结果做线性变换得到最终的注意力层输出张量。

下图展示了当 num_attention_heads = 2 时 attention 层的 Q/K/V 线性变换的并行计算方法。对每一块权重，我们都沿着列方向（k_dim）维度切割一刀。此时每个 head 上的

W

Q

W^Q

WQ、

W

K

W^K

WK、

W

V

W^V

WV的维度都变成 (d_model, k_dim//2)。每个 head 上能独立做矩阵计算，最后将计算结果 concat起来即可。整个流程如下图所示：

可以发现，从多头注意力结构看，其计算机制真的是天然适合模型的张量并行计算，即每个头上都可以在每个设备上独立计算，即可以把每个头（也可以是 n 个头）的参数放到一块 GPU 上，最后将子结果 concat 后得到最终的张量。具体来说，多头注意力结构的张量并行计算过程拆解如下，实际模型中，会存在一个或多个 head 占用一块 GPU 的情况，且我们尽量保证 heads 总数能被 GPU 个数整除。

对输出线性层 o_proj 做行并行切分 各 rank 完成本地 heads 的 attention 计算后，会得到本地 attention 输出。接下来，这些输出会进入输出线性层，即 o_proj。在 vLLM 中，o_proj 通常使用 RowParallelLinear，即按输入维做行并行切分。每个 rank 先基于自己的本地 heads 输出计算一个局部和，然后在 forward 末尾通过 All-Reduce 将各 rank 的局部和相加，恢复完整的 hidden states。 因此，Attention 的 TP 结构和 MLP 中先列并行、再行并行的模式非常相似：Q/K/V projection 类似列并行，中间的本地 head attention 不需要立即汇总完整张量，最后的 o_proj 类似行并行，并通过 All-Reduce 合并输出。 Attention 模块的张量并行加速过程如图 3(b) 所示：

vLLM 框架中 Attention 模块针对线性层的张量并行，实际代码实现也是先列后行线性层：

2.3.2 MHA 层的通讯量分析

很明显上述的设计对 MLP 和自注意力层均采用了将两组 GEMM 运算融合的策略，从而消除了一个同步步骤，并获得了更好的扩展性。基于此技术方案，在一个标准 transformer 层中，前向传播只需执行两次 all-reduce 操作，反向传播也仅需两次 all-reduce（详见图 4）。 和 MLP 模块类似，模型训练和推理阶段，MHA 层的张量并行通信量如下所示：

• 模型训练时，包含前向传播和反相传播两个过程，即两次 All-Reduce 操作，所以 MLP 层的总通讯量为：4bsh。
• 模型推理时，只有前向传播过程，即一次 All-Reduce 操作，所以 MLP 层的总通讯量为：2bsh。

简单理解，这里的两次All-Reduce分别来自于self-attention层和MLP层

2.3.3 attention 中的张量并行实例

在多头注意力中，Q、K、V 的输出维度可以看成由多个 attention heads 拼接而成。因此，当我们对 wq、wk、wv 的第二维进行切分时，本质上就是把不同的 heads 分配到不同的 tensor parallel rank 上计算。

在这个例子中，tp_size = 2，所以我们将 wq、wk、wv 按列切成两份。第一张卡（tp_rank = 0）负责前一半 heads；第二张卡（tp_rank = 1）负责后一半 heads。与 Q/K/V 投影不同，输出投影 wo 采用行切分。因为 attention 输出已经按 head 维度分布在不同的 rank 上，所以 wo 需要按输入维度（即第一维）进行切分：

wq_sub1 = wq_weight[:, :hidden_dim // 2]
wk_sub1 = wk_weight[:, :hidden_dim // 2]
wv_sub1 = wv_weight[:, :hidden_dim // 2]
wo_sub1 = wo_weight[:hidden_dim // 2, :]

随后，将输入复制到每个 tensor parallel rank 上，并分别与本地切分后的 Q/K/V 权重进行矩阵乘法。下面是第一张卡上的计算结果，分别记作 q1、k1、v1：

q1 = np.matmul(inputs, wq_sub1) # [bsz, seq_len, hidden_dim//2]
k1 = np.matmul(inputs, wk_sub1)
v1 = np.matmul(inputs, wv_sub1)

第二张卡的计算方式完全相同，只是使用的是 Q/K/V 权重的后一半切片。这样，每张卡都会得到自己负责的那部分 heads 对应的 q、k、v。接下来，每张卡使用本地的 q、k、v 独立计算 self-attention。由于标准多头注意力中不同 heads 之间是相互独立的，因此每个 rank 可以只计算自己负责的 heads，而不需要在 attention 计算阶段与其他 rank 通信。 得到局部 attention 输出之后，每张卡再将自己的 attention 输出与 wo 的对应行分块相乘。由于完整的输出投影可以拆成多个局部矩阵乘法结果之和，因此最后将各个 rank 的局部结果相加，就可以得到完整的 attention 输出：

output_tp_parallel = attn_output1 @ wo_sub1 + attn_output2 @ wo_sub2

就是说，整个 self-attention 的张量并行过程可以概括为： Q/K/V 投影：列并行

• wq、wk、wv 按输出维度切分
• 等价于按 attention heads 切分
• 每个 rank 计算一部分 heads 的 q、k、v

Attention 计算：

• 每个 rank 独立计算本地 heads 的 attention
• 标准 MHA 中不同 heads 之间不需要通信

输出投影 wo：行并行

• wo 按输入维度切分
• 每个 rank 计算局部输出
• 最后通过求和（即实际分布式实现中的 all-reduce）得到完整结果 因此，这个例子体现的正是 self-attention 中常见的“先列并行、后行并行”的张量并行模式：Q/K/V 投影阶段按列切分权重（即切分 heads）；输出投影阶段按行切分 wo，再对各个 rank 的局部输出进行求和。 以下是完整的代码，详见 code/course12/attention_tp.py，过程基本符合Megatron-LM 的张量并行算法中所述。

import numpy as np

if __name__ == '__main__':
# 参数设置
bsz = 4
seq_len = 16
hidden_dim = 128
num_heads = 8
head_dim = hidden_dim // num_heads

np.random.seed(42) # 设置随机种子，保证结果可复现

# 随机初始化权重矩阵
wq_weight = np.random.randn(hidden_dim, hidden_dim)
wk_weight = np.random.randn(hidden_dim, hidden_dim)
wv_weight = np.random.randn(hidden_dim, hidden_dim)
wo_weight = np.random.randn(hidden_dim, hidden_dim)

inputs = np.random.randn(bsz, seq_len, hidden_dim)

print("===== 标准版本的多头注意力计算 =====")
# 1. 标准版本 – 不分割计算
q = np.matmul(inputs, wq_weight) # [bsz, seq_len, hidden_dim]
k = np.matmul(inputs, wk_weight) # [bsz, seq_len, hidden_dim]
v = np.matmul(inputs, wv_weight) # [bsz, seq_len, hidden_dim]

# 重塑为多头形式
q = q.reshape(bsz, seq_len, num_heads, head_dim) # [bsz, seq_len, num_heads, head_dim]
k = k.reshape(bsz, seq_len, num_heads, head_dim) # [bsz, seq_len, num_heads, head_dim]
v = v.reshape(bsz, seq_len, num_heads, head_dim) # [bsz, seq_len, num_heads, head_dim]

# 调整维度顺序
q = np.transpose(q, (0, 2, 1, 3)) # [bsz, num_heads, seq_len, head_dim]
k = np.transpose(k, (0, 2, 1, 3)) # [bsz, num_heads, seq_len, head_dim]
v = np.transpose(v, (0, 2, 1, 3)) # [bsz, num_heads, seq_len, head_dim]

# 注意力分数计算
scores = np.matmul(q, np.transpose(k, (0, 1, 3, 2))) / np.sqrt(head_dim) # [bsz, num_heads, seq_len, seq_len]

# 应用softmax
attn_probs = np.exp(scores – np.max(scores, axis=–1, keepdims=True))
attn_probs = attn_probs / np.sum(attn_probs, axis=–1, keepdims=True) # [bsz, num_heads, seq_len, seq_len]

# 注意力输出计算
attn_output = np.matmul(attn_probs, v) # [bsz, num_heads, seq_len, head_dim]

# 恢复原始维度
attn_output = np.transpose(attn_output, (0, 2, 1, 3)) # [bsz, seq_len, num_heads, head_dim]
attn_output = attn_output.reshape(bsz, seq_len, hidden_dim) # [bsz, seq_len, hidden_dim]

# 最终输出投影
output = np.matmul(attn_output, wo_weight) # [bsz, seq_len, hidden_dim]

print("===== 张量并行版本的多头注意力计算 =====")
# 2. 张量并行版本 – 按头切分
# 每个并行组处理一半的头
heads_per_gpu = num_heads // 2

# GPU 1处理前half_heads个头
# 按列切分QKV权重 – 每个GPU负责一半头的权重
wq_sub1 = wq_weight[:, :hidden_dim // 2] # 前half_heads个头的权重
wk_sub1 = wk_weight[:, :hidden_dim // 2]
wv_sub1 = wv_weight[:, :hidden_dim // 2]
wo_sub1 = wo_weight[:hidden_dim // 2, :]
# GPU 1上的计算
q1 = np.matmul(inputs, wq_sub1) # [bsz, seq_len, hidden_dim//2]
k1 = np.matmul(inputs, wk_sub1)
v1 = np.matmul(inputs, wv_sub1)

# 重塑为多头形式
q1 = q1.reshape(bsz, seq_len, heads_per_gpu, head_dim)
k1 = k1.reshape(bsz, seq_len, heads_per_gpu, head_dim)
v1 = v1.reshape(bsz, seq_len, heads_per_gpu, head_dim)

# 调整维度顺序
q1 = np.transpose(q1, (0, 2, 1, 3)) # [bsz, heads_per_gpu, seq_len, head_dim]
k1 = np.transpose(k1, (0, 2, 1, 3))
v1 = np.transpose(v1, (0, 2, 1, 3))

# 计算注意力分数
scores1 = np.matmul(q1, np.transpose(k1, (0, 1, 3, 2))) / np.sqrt(head_dim)

# 应用softmax
attn_probs1 = np.exp(scores1 – np.max(scores1, axis=–1, keepdims=True))
attn_probs1 = attn_probs1 / np.sum(attn_probs1, axis=–1, keepdims=True)

# 注意力输出计算
attn_output1 = np.matmul(attn_probs1, v1) # [bsz, heads_per_gpu, seq_len, head_dim]

# 恢复原始维度
attn_output1 = np.transpose(attn_output1, (0, 2, 1, 3)) # [bsz, seq_len, heads_per_gpu, head_dim]
attn_output1 = attn_output1.reshape(bsz, seq_len, hidden_dim // 2) # [bsz, seq_len, hidden_dim//2]

# GPU 2处理后half_heads个头
wq_sub2 = wq_weight[:, hidden_dim // 2:] # 后half_heads个头的权重
wk_sub2 = wk_weight[:, hidden_dim // 2:]
wv_sub2 = wv_weight[:, hidden_dim // 2:]
wo_sub2 = wo_weight[hidden_dim // 2:, :]

# GPU 2上的计算
q2 = np.matmul(inputs, wq_sub2)
k2 = np.matmul(inputs, wk_sub2)
v2 = np.matmul(inputs, wv_sub2)

# 重塑为多头形式
q2 = q2.reshape(bsz, seq_len, heads_per_gpu, head_dim)
k2 = k2.reshape(bsz, seq_len, heads_per_gpu, head_dim)
v2 = v2.reshape(bsz, seq_len, heads_per_gpu, head_dim)

# 调整维度顺序
q2 = np.transpose(q2, (0, 2, 1, 3)) # [bsz, heads_per_gpu, seq_len, head_dim]
k2 = np.transpose(k2, (0, 2, 1, 3))
v2 = np.transpose(v2, (0, 2, 1, 3))

# 计算注意力分数
scores2 = np.matmul(q2, np.transpose(k2, (0, 1, 3, 2))) / np.sqrt(head_dim)

# 应用softmax
attn_probs2 = np.exp(scores2 – np.max(scores2, axis=–1, keepdims=True))
attn_probs2 = attn_probs2 / np.sum(attn_probs2, axis=–1, keepdims=True)

# 注意力输出计算
attn_output2 = np.matmul(attn_probs2, v2) # [bsz, heads_per_gpu, seq_len, head_dim]

# 恢复原始维度
attn_output2 = np.transpose(attn_output2, (0, 2, 1, 3)) # [bsz, seq_len, heads_per_gpu, head_dim]
attn_output2 = attn_output2.reshape(bsz, seq_len, hidden_dim // 2) # [bsz, seq_len, hidden_dim//2]

# 合并结果 (相当于在head维度上concatenate)
output_tp_parallel = attn_output1 @ wo_sub1 + attn_output2 @ wo_sub2

print(np.mean(np.abs(output – output_tp_parallel)))

2.4 Embedding 层的张量并行

Transformer 语言模型的输出侧通常会通过 LM head 将 hidden states 投影到词表维度。这里需要区分两个概念：LM head 的权重矩阵通常是 [vocab_size, hidden_size]，而计算得到的 logits 通常是 [batch_size 或 token 数, vocab_size]。由于当前 Transformer 语言模型的词汇表通常至少有数万个 token，因此对 LM head 这类大词表投影进行张量并行，通常可以减少单张 GPU 上的计算量和显存占用。 在 vLLM 中，输入 embedding 和输出 LM head 都采用按词表维度切分的方式。需要注意的是，输入 embedding 的通信方式和输出 LM head 不一样：

2.4.1 输入嵌入层

Embedding 层开启 TP 时，会将输入嵌入层的权重矩阵

E

E

E（尺寸为 [vocab_size, hidden_size]）按词汇维度拆分。由于词汇维度对应矩阵的第 0 维，所以可以理解为“按行拆分”：每个 GPU 只保存一段 token id 范围对应的 embedding 行。例如 TP=2 时，可以把词表分成两个连续范围，rank 0 负责前半段 token，rank 1 负责后半段 token。 因为每个分块只包含嵌入表的一部分，所以每张卡只能直接查到属于自己词表范围内的 token。对于不属于当前 rank 的 token，vLLM 会先通过 mask 将其排除，并把对应位置的 embedding 输出置为 0。随后，各个 rank 对局部 embedding 结果执行一次 all-reduce。由于同一个 token 只会在负责它的 rank 上产生非零 embedding，all-reduce 求和后就能得到完整的输入嵌入结果。 简单理解这个过程：假设现在有两个 GPU，每个 GPU 负责一半词表。输入请求的 seq=2，input_ids=[31, 33]。假设 token 31 属于 GPU 0 的词表范围，token 33 属于 GPU 1 的词表范围。那么在 GPU 0 上，token 31 能查到真实 embedding，token 33 不属于本卡词表范围，会被置为 0：

GPU 0:
[
[0.11, 0.22, 0.33, …, 0.12],
[0.00, 0.00, 0.00, …, 0.00]
]

在 GPU 1 上则相反，token 31 会被置为 0，token 33 能查到真实 embedding：

GPU 1:
[
[0.00, 0.00, 0.00, …, 0.00],
[0.31, 0.36, 0.13, …, 0.62]
]

最后对两个 GPU 的结果做 all-reduce 求和，就得到完整的输入嵌入：

[
[0.11, 0.22, 0.33, …, 0.12],
[0.31, 0.36, 0.13, …, 0.62]
]

也就是说，输入 embedding 的 TP 流程可以概括为：先判断每个 token 是否属于当前 rank 的词表范围；属于本 rank 的 token 正常查表；不属于本 rank 的 token 先 mask 掉并把输出置 0；最后通过 all-reduce 把各个 rank 的局部结果相加，恢复完整的输入 embedding。

import numpy as np

# ==========================================
# 0. 初始化
# ==========================================
np.random.seed(42)
Vocab = 10
Hidden = 4
B, S = 2, 2

# 完整的 Embedding 表 (模拟 Ground Truth)
# 形状: (10, 4)
E_full = np.random.randn(Vocab, Hidden)

# 输入 Token IDs (Batch=2, Seq=2)
# 包含落在两个 GPU 范围内的 ID
# 1, 3 -> GPU 1 (0-4)
# 6, 8 -> GPU 2 (5-9)
Input_IDs = np.array([
[1, 6],
[3, 8]
])

print("输入 IDs:\\n", Input_IDs)

# ==========================================
# 1. 单机基准 (Lookup)
# ==========================================
# Numpy 的高级索引模拟 Embedding Lookup
Output_Ref = E_full[Input_IDs]
print(f"\\n基准输出形状: {Output_Ref.shape}") # (2, 2, 4)

# ==========================================
# 2. 并行模拟 (Parallel Embedding)
# ==========================================
print("\\n— 开始并行模拟 —")

# [切分权重] 按词汇维度切分 (Row Parallel in terms of Matrix,
# 但通常称为 Vocab Parallel)
# Split Size = 5
V_per_gpu = Vocab // 2

# GPU 1: 负责 ID 0-4
E_gpu1 = E_full[:V_per_gpu, :] # (5, 4)
range_start_1 = 0
range_end_1 = V_per_gpu

# GPU 2: 负责 ID 5-9
E_gpu2 = E_full[V_per_gpu:, :] # (5, 4)
range_start_2 = V_per_gpu
range_end_2 = Vocab

def parallel_embedding_forward(input_ids, local_weight, start_idx, end_idx):
"""
每个 GPU 独立执行的 Forward 函数
"""
# 1. 创建掩码: 找出哪些 ID 属于当前 GPU
mask = (input_ids >= start_idx) & (input_ids < end_idx)

# 2. 将全局 ID 映射为本地 ID (Offset)
# 例如 GPU 2 负责 5-9，ID=6 对应的本地索引是 1
local_ids = input_ids – start_idx

# 3. 为了避免索引越界，将不属于自己的 ID 置为 0
safe_ids = np.where(mask, local_ids, 0)

# 4. 查表 (Lookup)
local_output = local_weight[safe_ids]

mask_expanded = mask[:, :, np.newaxis]

# 5. 只有属于自己的 ID 保留 Lookup 结果，其他的变成 0.0
final_local_output = local_output * mask_expanded

return final_local_output

# — GPU 1 计算 —
Out_gpu1 = parallel_embedding_forward(Input_IDs, E_gpu1, range_start_1, range_end_1)
print("\\nGPU 1 输出 (部分为0):")
print(Out_gpu1[0]) # 看第一行: [Vector(ID=1), Vector(0.0)]

# — GPU 2 计算 —
Out_gpu2 = parallel_embedding_forward(Input_IDs, E_gpu2, range_start_2, range_end_2)
print("\\nGPU 2 输出 (部分为0):")
print(Out_gpu2[0]) # 看第一行: [Vector(0.0), Vector(ID=6)]

# ==========================================
# 3. 同步聚合 (All-Reduce)
# ==========================================
Output_Fused = Out_gpu1 + Out_gpu2

print(f"\\n验证结果: {np.allclose(Output_Ref, Output_Fused)}")
print("逻辑: Embedding(ID) = Embedding_GPU1(ID) + Embedding_GPU2(ID)")
print(" 其中一个必然是 0向量，另一个是真实向量")

2.4.2 输出嵌入层

对于输出 LM head，先考虑一种朴素的完整 logits 汇总方案。假设 hidden states 为

X

X

X，输出权重矩阵按词表维度切分为

E

1

,

E

2

E_1, E_2

E1​,E2​。若

E

i

E_i

Ei​ 的形状沿用 embedding 权重的表示，即 [vocab_size/tp_size, hidden_size]，那么每个 rank 上的本地 logits 应写成：

Y

i

=

X

E

i

T

Y_i = X E_i^T

Yi​=XEiT​ 两个 rank 并行计算后，可以得到词表分片上的 logits：

[

Y

1

,

Y

2

]

=

[

X

E

1

T

,

X

E

2

T

]

[Y_1, Y_2] = [X E_1^T, X E_2^T]

[Y1​,Y2​]=[XE1T​,XE2T​]，接下来再通过 gather 或 all-gather 将这些本地 logits 按词表维度拼接起来，得到完整 logits：

Y

=

gather/all-gather

(

[

Y

1

,

Y

2

]

)

Y = \\text{gather/all-gather}([Y_1, Y_2])

Y=gather/all-gather([Y1​,Y2​])

三 vLLM 中的张量并行

3.1 vLLM 中张量并行如何使用

使用 vLLM 启动模型服务时，如果模型太大，无法放入单个 GPU，但可以放入单个节点中的多个 GPU，就可以使用张量并行。只考虑 TP、没有叠加 PP、DP、DCP 等其他并行维度时，tensor_parallel_size 通常就等于你希望用于单个模型副本的 GPU 数量。例如，单节点有 4 个 GPU，可以将张量并行大小设置为 4。 对于多 GPU 的离线推理，可以在 LLM 类中设置 tensor_parallel_size 为所需的 GPU 数量。例如，要在 4 个 GPU 上运行推理：：

from vllm import LLM
llm = LLM("Qwen/Qwen3-32B", tensor_parallel_size=4) # 或者改为模型权重的本地路径
output = llm.generate("San Francisco is a")

对于多 GPU 服务，也就是在线推理，可以在启动服务器时包含 –tensor-parallel-size。例如，在 4 个 GPU 上运行 API server：

# 前提是安装成功了 vllm，在可通过下述命令启动多 GPU 服务
vllm serve Qwen/Qwen3-32B \\
–tensor-parallel-size 4

在 vLLM 的多进程执行路径中，前端进程通过 ZMQ socket 与后台 EngineCore 通信。开启 DP 时，vLLM 会按 DP 模式管理多个 EngineCore；MoE 的 DP 场景会使用 DPEngineCoreProc，非 MoE 的 DP rank 则更接近多个相互独立的 EngineCore。 如果同时开启 DP 和 TP，可以把每个 DP rank 理解为一个模型副本，副本内部再按 tensor_parallel_size 切分到多个 GPU。只开启 DP+TP 时，1 个 DP 副本内部的 GPU 数量等于 TP 数；如果还叠加 PP（pipeline parallel）或 prefill context parallel，则还要继续乘上这些并行维度。

3.2 vLLM 中的分布式资源管理

vLLM 中与分布式资源管理相关的核心逻辑主要位于 vllm/distributed 目录。其中，distributed/parallel_state.py 负责分布式并行状态管理。它封装了 PyTorch ProcessGroup 等底层通信机制，为推理过程提供统一的并行组管理和通信接口。 工作流如下：

在 initialize_model_parallel() 中，vLLM 会根据当前 world_size、rank，以及 TP、PP、DP、prefill/decode context parallel、EP 等配置，构造全局 rank 布局，布局顺序是：ExternalDP x DP x PP x TP 随后，initialize_model_parallel() 通过 reshape、transpose 等操作生成不同并行维度的 group_ranks。例如，TP 组会把同一个 PP stage 内的 rank 放到一起；PP 组会把同一个 TP rank、跨不同 PP stage 的 rank 放到一起。 init_model_parallel_group() 更偏底层：它接收已经计算好的 group_ranks，并创建对应的 GroupCoordinator。GroupCoordinator 是对 PyTorch ProcessGroup 的封装，会为这些 rank 创建 device 通信组，并记录当前进程在组内的 rank_in_group、world_size 等信息。后续的 all_reduce、all_gather、send/recv 等操作，都会通过对应的 GroupCoordinator 在正确的 rank 范围内执行。 vLLM 使用模块级全局变量保存并行组状态。例如，_TP 记录张量并行组，_PP 记录流水线并行组，_DP 记录数据并行组；当前版本还包含 _DCP、_PCP、_EP、_EPLB 等组状态。其中，_EP 主要用于 MoE 场景，dense model 通常不会创建 EP group。后续代码可以通过 get_tp_group()、get_pp_group()、get_dp_group() 等函数获取对应通信组。 例如，假设有 8 张 GPU，并设置：

• tensor_parallel_size = 2
• pipeline_parallel_size = 4
• data_parallel_size = 1

此时world_size = 8，avLLM 会先将 global rank 组织成一个二维结构：

PP stage 0: GPU0 GPU1
PP stage 1: GPU2 GPU3
PP stage 2: GPU4 GPU5
PP stage 3: GPU6 GPU7
TP0 TP1

然后生成 TP 和 PP 两类通信组： TP groups：

• [GPU0, GPU1]
• [GPU2, GPU3]
• [GPU4, GPU5]
• [GPU6, GPU7]

PP groups：

• [GPU0, GPU2, GPU4, GPU6]
• [GPU1, GPU3, GPU5, GPU7]

其中，TP 组负责同一个 PP stage 内的张量并行通信。例如，[GPU2, GPU3] 共同负责 stage 1 的模型层，RowParallelLinear 的 AllReduce 只会在 GPU2 和 GPU3 之间发生。 PP 组负责同一个 TP rank 跨不同 PP stage 的通信。例如，[GPU1, GPU3, GPU5, GPU7] 表示 TP rank 1 这一路的流水线通道。前向执行时，stage 0 的 GPU1 会将对应激活发送给 stage 1 的 GPU3，随后再沿 GPU3 -> GPU5 -> GPU7 向后传递

_TP = init_model_parallel_group(group_ranks,
get_world_group().local_rank,
backend,
use_message_queue_broadcaster=True,
group_name="tp")

四种并行组作用如下：

• TP（张量并行）：每组负责模型参数的分区计算，通常用于分割权重矩阵。
• PP（流水线并行）：每组负责模型的不同层或阶段，适合深层模型分段计算。
• DP（数据并行）：每组处理不同的数据批次，实现批量训练或推理;每个 DP group 的成员在模型上参数完全相同，只负责数据并行。
• EP（专家并行）：用于 MoE（Mixture of Experts）等结构，每组分配部分专家网络。专家并行 = 数据并行 × 张量并行的组合，每个 group 内 rank 会共享一组专家层的负载。

实例理解，TP + PP 并行组的概念。在 8 卡 GPU 上配置 流水线并行度（PP）为 4 和 张量并行度（TP）为 2 时，并行分组逻辑如下：

• 全局并行设置：
  • world_size（总 GPU 数）= 8
  • pipeline_model_parallel_size（流水线并行组数）= 4
  • tensor_model_parallel_size（张量并行组大小）= 2
• 并行分组详解：
• 张量并行（TP）分组：模型参数在组内的 2 张 GPU 间进行切分。系统会形成 4 个 TP 组，每个组处理模型的一部分。
• 流水线并行（PP）分组：模型层被分配到 4 个连续的流水线阶段。系统会形成 2 个 PP 组，每个组内的 4 张 GPU 分别负责一个阶段，共同完成一个完整的批次处理。

# 共 4 个张量并行（TP）组，每组在 2 张卡上切分模型参数
TP Group 0: [GPU0, GPU1]
TP Group 1: [GPU2, GPU3]
TP Group 2: [GPU4, GPU5]
TP Group 3: [GPU6, GPU7] # 共 2 个流水线并行（PP）组，每组由 4 张卡构成一个完整流水线
PP Group 0: [GPU0, GPU2, GPU4, GPU6] # 处理微批次 A
PP Group 1: [GPU1, GPU3, GPU5, GPU7] # 处理微批次 B

结合图表与分组信息，可以这样理解 TP + PP 并行组内的交互：

模拟分布式并行分分组算法

可以通过下述示例代码（CPU模拟，不需要多机）来模拟分布式并行分分组算法，得到不同 gpu 设备和不同推理配置下的并行分组信息。rank网格简化为 [ExternalDP, DP, PP, TP]。

import torch

def compute_groups(world_size, tp, pp, dp):
assert world_size % (tp * pp * dp) == 0
E = world_size // (tp * pp * dp) # ExternalDP 大小（通常为 1）
all_ranks = torch.arange(world_size).reshape(E, dp, pp, tp)

# TP 组
tp_groups = [x.tolist() for x in all_ranks.view(–1, tp).unbind(0)]

# PP 组
pp_groups = [x.tolist() for x in all_ranks.transpose(2, 3).reshape(–1, pp).unbind(0)]

# DP 组（模型内部 DP）
dp_groups = [x.tolist() for x in all_ranks.transpose(1, 3).reshape(–1, dp).unbind(0)]

# EP 组（同 PP stage 下合并 DP×TP）
ep_groups = [x.tolist() for x in all_ranks.transpose(1, 2).reshape(–1, dp * tp).unbind(0)]

return tp_groups, pp_groups, dp_groups, ep_groups

print("TP groups:", tp_groups)
print("PP groups:", pp_groups)
print("DP groups:", dp_groups)
print("EP groups:", ep_groups)

if __name__ == "__main__":
"""
world_size = 16
tp_size = 4
pp_size = 1
dp_size = 4
"""
compute_groups(16, 4, 1, 4) # 卡数 16

输出结果如下所示:

TP groups: [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]
PP groups: [[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], [12], [13], [14], [15]]
DP groups: [[0, 4, 8, 12], [1, 5, 9, 13], [2, 6, 10, 14], [3, 7, 11, 15]]
EP groups: [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]]

这里把这些 rank 按照不同类型拆分到 4 维上。这里的 ext_dp_size 指的是本系统之外的一层数据并行维度，在本节课中不做深究。

• DP 维度对应传统意义上的数据并行组（组内 rank 拥有完整模型、处理不同子 batch）
• PP 维度对应流水线并行的层切分（每个 stage 只持有一段层）
• TP 维度对应张量并行的参数切分（同一层内部的大矩阵按维度拆到多个 rank 上计算，并 All-Reduce 合并） 在此基础上，引入 EP 时，通常是把同一个 PP stage 下若干个 (DP, TP) 的笛卡尔积子网格看作一个 EP group。也就是说，同一个 PP stage 内的不同 DP、TP rank 共同承担同一批 token 的 MoE 路由与专家计算，同时各自只持有部分专家，即让一个 EP group 里的所有 GPU 一起负责一整批 token 的 MoE 计算。 以上 rank 网格中我们忽略了 PCP（prefill_context_parallel_size）。这里的 prefill_context_parallel_size 指的是：在 prefill 阶段，将长上下文按 token 维度切分到多个 rank 上并行处理，因此实际的 rank 网格里除了 DP、PP、TP 之外，还会额外增加一个 PCP 维度。

3.3 vLLM 的并行线性层

vllm/model_executor/models/qwen3.py 代码中的 Qwen3Attention 模块的核心网络层组成如下所示:

Qwen3Attention (继承自 nn.Module)
├── qkv_proj: QKVParallelLinear (用于合并计算 Q, K, V)
│ └── (继承自) ColumnParallelLinear (列并行层)
│ └── (继承自) LinearBase (并行层抽象基类)
│
├── o_proj: RowParallelLinear (用于 Attention 输出)
│ └── (继承自) LinearBase (并行层抽象基类)
│
├── rotary_emb: RotaryEmbedding (旋转位置编码)
│
├── attn: Attention (底层 Attention 计算核心)
│ └── (使用) Context (用于获取 prefill/decode 状态及相关参数)
│
├── q_norm: RMSNorm (对 Query 向量进行归一化)
│
└── k_norm: RMSNorm (对 Key 向量进行归一化)

其中 qkv 线性层使用 TP 是按列切分权重，o 线性层是按行切分权重。

ColumnParallelLinear 类源码剖析(列并行)

ColumnParallelLinear 和 RowParallelLinear 类的基类是 LinearBase，它为所有并行及量化的线性层提供了基础框架和通用属性，其核心支持了普通与量化两种线性层实现。 ColumnParallelLinear 类是一种实现列并行（Column Parallelism）的线性层。其中权重 A 按第二维（输出维度）列切分，被划分为多个子矩阵：

A

=

[

A

1

,

A

2

,

…

,

A

p

]

A = [A_1, A_2, \\dots, A_p]

A=[A1​,A2​,…,Ap​]，每个 GPU 只存储并计算其中一块

A

i

A_i

Ai​，对应输出

Y

i

=

X

A

i

Y_i = X A_i

Yi​=XAi​。如果需要完整的输出，多个 GPU 上的

Y

i

Y_i

Yi​<span class=