Tensor Parallel：当一张 GPU 放不下一个大模型

前面的文章里，我们一直默认模型可以完整放在一张 GPU 上。

这让系统设计简单很多。Engine 只需要管理请求状态，Scheduler 只需要决定哪些 Sequence 运行，Block Manager 只需要管理这张 GPU 上的 KV cache blocks，ModelRunner 只需要在本地执行一次 forward。

但真实的大模型推理经常不是这样。

一个 7B 模型用 FP16 权重大约需要十几 GB 显存。一个 70B 模型只算权重就可能需要一百多 GB 显存。再加上 KV cache、CUDA buffer、临时激活、通信 buffer，一张 GPU 很可能放不下。

这时我们有两个选择。

第一，把模型换小。

第二，把模型切到多张 GPU 上。

Tensor parallel 要解决的就是第二个问题。

它的核心思想是：把模型中的大矩阵按某个维度切开，让多张 GPU 分别保存一部分权重，并共同完成一次 forward。

它不是把不同请求分给不同 GPU。那叫数据并行或多副本 serving。

Tensor parallel 是一个请求的一次模型计算，会同时使用多张 GPU。

这一篇我们要讨论：为什么 Transformer 适合做 tensor parallel，linear 层怎么切，attention 怎么切，MLP 怎么切，通信发生在哪里，以及它接入 mini vLLM 后，会改变哪些模块。

Tensor parallel 解决的是模型太大，不是请求太多

先区分两个概念。

如果模型可以放进一张 GPU，但请求太多，我们可以复制多份模型。

GPU 0: 一份完整模型，处理一部分请求
GPU 1: 一份完整模型，处理另一部分请求
GPU 2: 一份完整模型，处理另一部分请求

这更像 data parallel serving。

每张 GPU 都有完整模型副本。请求被分发到不同副本上。它提升的是服务吞吐，但不能解决单个模型太大放不下的问题。

Tensor parallel 不一样。

它面对的是模型本身太大。

例如一个模型权重太大，一张 GPU 放不下，只能把权重拆到多张 GPU。

GPU 0: 保存模型的一部分权重
GPU 1: 保存模型的一部分权重
GPU 2: 保存模型的一部分权重
GPU 3: 保存模型的一部分权重

一次 forward 时，多张 GPU 协同计算同一个 batch。

所以 tensor parallel 解决的是：

单个模型无法完整放入单张 GPU
单次 forward 需要多张 GPU 协同完成

这会带来一个直接后果：每层之间需要通信。

如果只是把请求分发给不同 GPU，GPU 之间几乎不需要在每一层通信。

但 tensor parallel 把同一层的权重拆开了，所以 forward 过程中必须把不同 GPU 的部分结果合并或同步。

通信就是 tensor parallel 的核心成本。

Transformer 中最适合切的是大矩阵

Transformer 的主要计算来自 linear 层。

在 decoder-only LLM 中，每层大致包含：

attention projection:
  q_proj
  k_proj
  v_proj
  o_proj

MLP:
  gate_proj
  up_proj
  down_proj

这些本质上都是矩阵乘法。

如果 hidden size 很大，权重矩阵也很大。

Tensor parallel 的基本做法，就是把这些大矩阵切开。

矩阵可以按列切，也可以按行切。

这对应两种常见并行方式：

Column Parallel Linear
Row Parallel Linear

理解这两种切法，是理解 tensor parallel 的关键。

Column Parallel Linear：按输出维度切

先看一个普通 linear。

Y = XW

假设：

X: [batch, in_features]
W: [in_features, out_features]
Y: [batch, out_features]

Column parallel 的做法是按 W 的输出维度切，也就是按列切。

如果有 2 张 GPU：

W = [W0, W1]

其中：

W0: [in_features, out_features / 2]
W1: [in_features, out_features / 2]

GPU 0 计算：

Y0 = XW0

GPU 1 计算：

Y1 = XW1

最后完整输出是：

Y = concat(Y0, Y1)

也就是说，column parallel 的每张 GPU 负责输出的一部分通道。

它的好处是，计算完后如果下一步也能接受切分后的输出，就不一定立刻通信。

比如 attention 里的 q_proj、k_proj、v_proj 很适合 column parallel。

因为每个 attention head 可以分配到不同 GPU 上。每张 GPU 只计算自己负责的 heads。

这种切法可以减少权重显存，并让多张 GPU 分担矩阵乘法。

Row Parallel Linear：按输入维度切

Row parallel 是按输入维度切。

仍然看：

Y = XW

现在把 W 按行切。

如果有 2 张 GPU：

W = [W0; W1]

其中：

W0: [in_features / 2, out_features]
W1: [in_features / 2, out_features]

对应地，输入 X 也按最后一维切。

X = [X0, X1]

GPU 0 计算：

Y0 = X0W0

GPU 1 计算：

Y1 = X1W1

完整输出是：

Y = Y0 + Y1

这里不是 concat，而是求和。

所以 row parallel 通常需要一次 all reduce，把不同 GPU 上的部分结果加起来。

这就是 row parallel 的通信点。

它的特点是：输入是切分的，输出需要聚合。

Attention 里的 tensor parallel

现在把这两种切法放回 Transformer attention。

Attention 里通常有四个 projection：

q_proj
k_proj
v_proj
o_proj

q_proj、k_proj、v_proj 通常适合 column parallel。

因为它们的输出可以按 head 切。

假设模型有 32 个 attention heads，tensor parallel size 是 4。

那么每张 GPU 可以负责 8 个 heads。

GPU 0: heads 0 到 7
GPU 1: heads 8 到 15
GPU 2: heads 16 到 23
GPU 3: heads 24 到 31

每张 GPU 只保存自己 heads 对应的 q、k、v projection 权重，也只计算自己 heads 的 attention。

这样 attention 的中间计算可以局部完成。

每张 GPU 有自己的 Q、K、V。

每张 GPU 也可以维护自己 heads 对应的 KV cache。

这很自然，因为 KV cache 本来就有 head 维度。

如果 heads 被切到不同 GPU，KV cache 也可以按 head 切。

最后 attention 得到每张 GPU 上的部分输出。

接下来进入 o_proj。

o_proj 通常适合 row parallel。

因为 attention 输出在 head 维度上被切分了，需要通过 o_proj 投影回 hidden size。

每张 GPU 对自己那部分 heads 的输出做部分 o_proj，得到一个 partial hidden output。

然后多个 GPU 的 partial output 需要相加。

这时需要 all reduce。

所以 attention tensor parallel 的典型模式是：

q/k/v projection: column parallel
attention per GPU: local heads
o projection: row parallel
all reduce 合并输出

这个结构很重要。

它说明 tensor parallel 不是简单地把层平均分给不同 GPU，而是每层内部按矩阵维度切开，并在必要位置通信。

MLP 里的 tensor parallel

Transformer 的 MLP 部分也很适合 tensor parallel。

以常见 gated MLP 为例，它包含：

gate_proj
up_proj
down_proj

其中 gate_proj 和 up_proj 把 hidden size 扩展到 intermediate size。

down_proj 再把 intermediate size 投回 hidden size。

通常 gate_proj 和 up_proj 使用 column parallel。

每张 GPU 负责 intermediate 维度的一部分。

gate_i = X @ gate_W_i
up_i = X @ up_W_i

然后本地做激活和逐元素乘法。

hidden_i = activation(gate_i) * up_i

因为 intermediate 维度已经被切分，每张 GPU 可以独立处理自己那部分。

接着 down_proj 使用 row parallel。

每张 GPU 用自己的 hidden_i 乘以对应的 down_proj 权重，得到 partial output。

最后 all reduce，把各 GPU 的 partial output 相加。

所以 MLP 的典型模式是：

gate/up projection: column parallel
activation: local
down projection: row parallel
all reduce 合并输出

这和 attention 的模式非常相似。

先把中间维度切开，本地计算；最后投回 hidden size 时做 all reduce。

通信是 tensor parallel 的代价

Tensor parallel 可以让多张 GPU 分担权重和计算，但它不是免费的。

每一层 attention 的 o_proj 后可能需要 all reduce。

每一层 MLP 的 down_proj 后也可能需要 all reduce。

如果模型有 80 层，那么每个 forward 中会有大量通信。

通信开销取决于：

GPU 之间的互联带宽
tensor parallel size
batch size
sequence length
hidden size
实现是否融合通信和计算

如果 GPU 之间是 NVLink，通信会比普通 PCIe 好很多。

如果跨机器做 tensor parallel，通信成本会更高，通常更复杂。

这也是为什么 tensor parallel 的收益不是无限增加的。

TP size 从 1 到 2 可能收益明显。

从 2 到 4 可能还不错。

但继续增加到 8 或 16，通信成本可能吞掉计算收益。

所以 tensor parallel size 是一个需要调优的系统参数。

它既受模型大小约束，也受硬件拓扑影响。

Tensor parallel 对 KV cache 的影响

在单 GPU 版本里，Block Manager 管理整份 KV cache。

每个 block 包含所有 KV heads。

但在 tensor parallel 中，attention heads 被切到不同 GPU。

每张 GPU 只负责一部分 heads。

那么 KV cache 也自然被切开。

例如模型有 32 个 KV heads，TP size 是 4。

每张 GPU 负责 8 个 KV heads。

这意味着每张 GPU 上都有自己的 KV cache pool，只保存本 rank 负责的 heads。

Block table 可以在所有 rank 上保持一致，因为 token 维度上的 block 分配是一样的。

但 physical KV cache 的内容在每张 GPU 上不同。

Rank 0:
  block 3 保存 heads 0 到 7 的 KV

Rank 1:
  block 3 保存 heads 8 到 15 的 KV

Rank 2:
  block 3 保存 heads 16 到 23 的 KV

Rank 3:
  block 3 保存 heads 24 到 31 的 KV

这里 block id 可以保持一致。

Sequence 的 block table 在每个 rank 上相同。

PagedAttention 在每个 rank 上读取同样逻辑 block 对应的 physical block，但只读取本 rank 的 heads。

这样设计会比较清晰。

Block Manager 的元数据可以复制到每个 rank。

每个 rank 各自管理本地 KV cache tensor。

释放 block 时，所有 rank 都释放相同 block id。

这种方式要求各 rank 的 block 分配保持同步。

也就是说，Block Manager 的决策必须在 tensor parallel group 内一致。

不能 Rank 0 给 Sequence A 分配 block 3，而 Rank 1 给它分配 block 5。

否则 block table 就乱了。

所以 tensor parallel 后，Block Manager 的元数据要么由一个 rank 决策后广播，要么每个 rank 执行完全相同的确定性分配逻辑。

教学版可以先采用 rank 0 决策并广播的思路，概念更清楚。

Tensor parallel 对 Scheduler 的影响

Scheduler 的请求调度逻辑总体不变。

它仍然决定哪些 Sequence 做 prefill，哪些 Sequence 做 decode。

但是调度结果需要被所有 tensor parallel ranks 知道。

因为一次 forward 要由所有 ranks 共同完成。

如果 rank 0 调度了 Sequence A 和 B，rank 1 也必须执行同样的 batch。

否则不同 GPU 处理的不是同一批请求，计算结果无法合并。

所以在 tensor parallel 下，Scheduler 不能只存在于某个本地孤岛里。

可以有两种实现方式。

一种是每个 rank 都维护相同 Scheduler 状态，并且所有请求状态更新保持同步。

这种方式减少广播，但实现复杂。

另一种是主 rank 负责 Scheduler，生成调度结果后广播给其他 ranks。

这种方式更容易理解。

mini vLLM 第一版可以采用主 rank 调度。

rank 0:
  接收请求
  维护 Scheduler
  生成 SchedulerOutputs
  广播给其他 ranks

所有 ranks:
  根据同一份 SchedulerOutputs 执行本地 forward

采样也可以先放在 rank 0。

因为 logits 最终需要完整 vocab 分布。是否每个 rank 都有完整 logits，取决于 lm_head 是否也被切分。第一版可以先让 logits 汇总到 rank 0，再由 rank 0 采样。

这样实现不会最高效，但语义清晰。

lm_head 怎么处理

最后一层 lm_head 把 hidden states 投影到 vocab size。

logits = hidden @ lm_head

如果 vocab 很大，lm_head 也很大。

它也可以做 tensor parallel。

常见做法是按 vocab 维度切，也就是 column parallel。

每张 GPU 负责一部分 vocab logits。

Rank 0: logits for vocab 0 到 V/4
Rank 1: logits for vocab V/4 到 V/2
...

这样每张 GPU 只计算部分 logits。

但 sampling 需要在完整 vocab 上做 argmax、top k、top p。

这就需要额外通信。

greedy 情况下，可以每个 rank 先找本地最大 logit，然后 all gather 或 all reduce 找全局最大。

top k / top p 更复杂，因为需要在分布层面处理完整 vocab。

为了简化，mini vLLM 第一版可以不切 lm_head，或者把完整 logits gather 到 rank 0 再采样。

这不一定高效，但有助于先跑通 tensor parallel 主链路。

后续再优化分布式 sampling。

这里要再次强调：教学实现不要一开始追求所有细节都生产级。

先保证读者能理解模型层如何切分、通信在哪里发生、系统状态如何同步。

Tensor parallel 对 Engine 架构的影响

单 GPU Engine 里，ModelRunner 是一个本地对象。

Tensor parallel 后，ModelRunner 变成了分布式对象。

每个 rank 都有自己的 ModelRunner，保存本 rank 的权重分片和本地 KV cache。

一次 Engine.step 需要所有 ranks 一起执行。

可以理解成：

rank 0 Engine:
  schedule
  broadcast batch metadata

all ranks ModelRunner:
  run local forward shard

communication:
  all reduce / all gather

rank 0:
  collect logits
  sample
  update Sequence states
  broadcast state updates if needed

这会让 Engine 的边界更复杂。

尤其是状态更新。

如果只有 rank 0 维护完整 Sequence 状态，那么其他 ranks 只需要知道下一轮执行所需的 metadata，比如 input tokens、position ids、block table、slot mapping。

如果每个 rank 都维护 Sequence 状态，就需要同步每轮采样结果、finished 状态和 block 释放。

第一版建议采用 rank 0 作为控制平面。

也就是说：

rank 0 负责请求状态、调度、采样、输出
所有 ranks 负责模型分片计算

这种控制平面和执行平面的分离，更适合教学版本。

它也符合我们前面一直强调的模块边界。

Tensor parallel 和 Pipeline parallel 的区别

讲 tensor parallel 时，经常会有人问：为什么不把不同层放到不同 GPU？

这其实是另一种并行方式，叫 pipeline parallel。

Tensor parallel 是把同一层的大矩阵切到多张 GPU。

Pipeline parallel 是把不同层放到不同 GPU。

例如：

GPU 0: layer 0 到 19
GPU 1: layer 20 到 39
GPU 2: layer 40 到 59
GPU 3: layer 60 到 79

Pipeline parallel 可以减少每张 GPU 保存的层数，但它会带来流水线调度问题。

对于训练，pipeline parallel 很常见。

对于在线 decode 推理，pipeline parallel 会遇到气泡和延迟问题，尤其是 batch 不够大时。

Tensor parallel 更适合让一次 forward 在多张 GPU 上并行完成，但通信频繁。

实际生产系统可能同时使用 tensor parallel 和 pipeline parallel。

但 mini vLLM 第一版不需要同时讲两个。

我们先专注 tensor parallel，因为它是推理中非常常见的模型切分方式，也和 attention heads、MLP 矩阵切分关系最直接。

第一版 mini vLLM 如何支持 tensor parallel

对 mini vLLM 来说，第一版 tensor parallel 可以做几个简化。

第一，只支持单机多卡。

先不考虑跨机器通信。

第二，只支持固定 tensor parallel size。

模型启动时指定 TP size，运行中不动态变化。

第三，attention heads 必须能被 TP size 整除。

这样每个 rank 分到相同数量的 heads。

第四，MLP intermediate size 也最好能被 TP size 整除。

这样 column parallel 和 row parallel 更容易实现。

第五，Scheduler 和请求状态只在 rank 0 维护。

其他 rank 接收 batch metadata 并执行本地 forward。

第六，采样先在 rank 0 完成。

logits 可以先 gather 到 rank 0，后续再优化。

第七，Block Manager 元数据由 rank 0 决策并广播，所有 rank 保持同样 block table。

这些限制能让系统先跑起来。

一旦基本链路正确，再继续优化通信、分布式采样、KV cache 同步和 backend kernel。

如何验证 tensor parallel 正确

Tensor parallel 的正确性验证非常重要。

最基本的测试是和单 GPU 结果对齐。

给定同一个模型权重、同一个输入，在 TP size 为 1 和 TP size 为 N 时，输出 logits 应该接近。

如果使用 greedy decoding，生成 token 序列也应该一致。

可以分层测试。

先测试一个 column parallel linear。

把完整权重切成几份，分别计算，再 concat，结果应该等于完整 linear。

再测试 row parallel linear。

把输入和权重切开，各自计算 partial output，再 all reduce，结果应该等于完整 linear。

然后测试 attention layer。

再测试 MLP layer。

最后测试完整 Transformer。

这种分层测试很重要。

因为如果完整模型输出不一致，很难知道是 q_proj 切错了，还是 o_proj all reduce 错了，还是 MLP down_proj 切错了，还是 lm_head logits 合并错了。

对于分布式系统，越要从小模块开始验证。

性能上应该关注什么

Tensor parallel 的 benchmark 不能只看能不能跑大模型，还要看扩展效率。

假设 TP size 从 1 增加到 2，理想情况下速度接近 2 倍，显存接近减半。

但真实情况不会这么完美，因为有通信成本。

可以记录：

每层 forward 时间
all reduce 时间
all gather 时间
decode tokens/s
TTFT
TPOT
每张 GPU 显存占用
通信占总时间比例

如果 TP size 增加后，显存下降了，但 TPOT 没有改善，可能说明通信成为瓶颈。

如果 batch 很小，tensor parallel 的计算并行收益可能不足以抵消通信开销。

如果 batch 较大，收益可能更明显。

所以 tensor parallel 的性能要结合 batch size 和硬件拓扑看。

它首先是为了解决模型放不下的问题，其次才是加速。

这一点要说清楚。

否则读者可能误以为 TP 一定让推理更快。

实际上，如果模型本来单卡放得下，TP size 增加不一定总是收益。

小结

这一篇我们讨论了 tensor parallel。

它解决的是单张 GPU 放不下大模型的问题。不同于把请求分发到多个完整模型副本，tensor parallel 是把同一个模型层内的大矩阵切到多张 GPU 上，让一次 forward 由多张 GPU 协同完成。

核心切法有两种。

Column Parallel Linear 按输出维度切，每张 GPU 计算一部分输出通道。

Row Parallel Linear 按输入维度切，每张 GPU 计算 partial output，然后通过 all reduce 合并。

在 Transformer 中，q/k/v projection 和 MLP 的 gate/up projection 通常适合 column parallel。attention 的 o_proj 和 MLP 的 down_proj 通常适合 row parallel，并在输出处进行 all reduce。

Tensor parallel 会影响 KV cache。attention heads 被切到不同 ranks 后，每张 GPU 只保存自己负责 heads 的 KV cache。Sequence 的 block table 可以在所有 ranks 上保持一致，但每个 rank 的 physical KV cache 内容不同。

它也会影响 Engine 架构。第一版 mini vLLM 可以采用 rank 0 作为控制平面，负责 Scheduler、Sequence 状态、采样和输出；所有 ranks 作为执行平面，保存权重分片并执行本地 forward。

Tensor parallel 的代价是通信。TP size 越大，权重和计算分摊越多，但 all reduce、all gather 等通信也越多。它不是无条件加速，首先是让大模型能放进多卡，其次才可能带来吞吐提升。

下一篇文章，我们会讨论量化推理。

Tensor parallel 用更多 GPU 解决模型太大的问题，而量化则从另一个方向出发：能不能让模型权重和 KV cache 用更低精度表示，从而减少显存占用并提升推理效率。