6.2 分布式KV Cache与多GPU推理

6.2 分布式KV Cache与多GPU推理

分布式推理的必然性

当模型规模突破数百亿参数的门槛，单张GPU的显存容量已经无法容纳完整的模型权重和推理所需的KV Cache。以Llama-3-70B为例，其FP16权重大约需要140GB显存，而KV Cache在长序列场景下还需要额外的数十GB。即便是当前最大的A100 80GB或H100 80GB，单卡也无法满足需求。

解决方案是将模型和计算分布到多张GPU甚至多台服务器上，这就是分布式推理。在这一场景中，KV Cache的管理呈现出全新的复杂性：每一层Transformer的KV Cache自然地归属于计算该层的GPU，跨层的注意力计算需要跨设备的数据传输。

张量并行中的KV Cache

张量并行与KV Cache的切分

张量并行（Tensor Parallelism, TP）是分布式推理中最常用的并行策略。在注意力层中，TP将注意力头沿列方向切分到不同的GPU上：

假设模型有 32 个注意力头，TP=4：
GPU 0: 负责头 0-7 的 KV Cache
GPU 1: 负责头 8-15 的 KV Cache
GPU 2: 负责头 16-23 的 KV Cache
GPU 3: 负责头 24-31 的 KV Cache
​

每个GPU只需要存储和计算部分注意力头的KV Cache，显存占用自然降低为原来的1/N（N为TP度数）。

TP中的KV Cache通信模式

在自回归推理的每一步，张量并行需要两次全局通信（All-Reduce）来聚合注意力计算结果：

1. KV计算阶段：各GPU独立计算自己负责的注意力头的K和V投影，存储在本地KV Cache中
2. 注意力分数计算：各GPU从本地KV Cache读取K和V，计算部分注意力分数
3. All-Reduce聚合：跨GPU通信聚合所有注意力头的输出

关键观察：KV Cache数据不需要跨GPU传输——每个GPU只访问自己负责的注意力头的KV Cache，这使得张量并行中的KV Cache管理相对简单。通信开销主要集中在注意力输出的聚合上，而非KV Cache的传递。

TP的KV Cache显存分析

对于TP=N的配置，每张GPU的KV Cache显存占用为：

KV_per_GPU = (num_layers / N_not_relevant) × (num_heads / TP) × (seq_len × head_dim) × 2 × dtype_size
​

注意：在标准的Megatron-style张量并行中，KV Cache按注意力头维度切分，因此每张GPU只需要存储 num_heads/TP 个头的KV数据。但实际上，KV切分方式有两种选择：

• 沿注意力头切分（Head-parallel）：每GPU存部分头的完整KV，最常用
• 沿head_dim切分（Head-dim-parallel）：每GPU存所有头的部分KV维度，通信模式不同

流水线并行中的KV Cache

流水线并行的KV分布

流水线并行（Pipeline Parallelism, PP）将模型的不同层分配到不同的GPU上：

GPU 0: 层 0-15
GPU 1: 层 16-31
GPU 2: 层 32-47
GPU 3: 层 48-63
​

在这种配置下，KV Cache的自然分布是按层划分：每张GPU只存储自己负责的层的KV Cache。

流水线气泡与KV Cache

流水线并行的一个经典问题是"气泡"（bubble）——由于不同GPU处理不同层的计算，数据需要在GPU之间传递，导致部分GPU处于空闲等待状态。在自回归推理中，每生成一个token都需要数据穿越所有流水线阶段：

生成步骤 t:
GPU 0 → GPU 1 → GPU 2 → GPU 3
       ↑ 写入第0-15层KV
                ↑ 写入第16-31层KV
                         ↑ 写入第32-47层KV
                                  ↑ 写入第48-63层KV + 输出token
​

KV Cache在流水线间的传递

虽然KV Cache数据不需要显式跨GPU传递（每层的KV Cache留在对应的GPU上），但中间激活值（hidden states）需要在流水线阶段间传递。这些激活值的大小与KV Cache密切相关——实际上，在推理过程中，层间传递的激活值远小于KV Cache本身。

TP + PP 组合：真实部署场景

在实际的分布式推理部署中，张量并行和流水线并行通常组合使用：

配置示例：TP=4, PP=2, 共8张GPU
GPU 0-3: 处理前半部分层（0-31），每张处理8个注意力头
GPU 4-7: 处理后半部分层（32-63），每张处理8个注意力头
​

在这种配置中，KV Cache的分布更加复杂：

• 沿注意力头维度被TP切分
• 沿层维度被PP切分
• 每张GPU存储约 num_heads/TP × num_layers/PP 个KV Cache块

跨节点分布式推理中的KV Cache挑战

NVLink vs. PCIe vs. 网络

不同互联方式的带宽差异巨大，直接影响KV Cache相关操作的性能：

KV Cache的跨节点迁移

在弹性推理（Elastic Inference）场景中，请求可能需要在推理过程中从一个节点迁移到另一个节点（例如，原节点负载过高时）。这就需要将已计算的KV Cache从源节点迁移到目标节点。

KV Cache迁移的效率取决于：

• 数据量：KV Cache可能达到数GB（长序列 × 大模型）
• 网络带宽：跨节点传输可能需要数百毫秒
• 迁移策略：全量迁移 vs. 增量迁移（只传输新计算的token对应的Cache）

Sequence Parallelism：新的并行范式

序列并行（Sequence Parallelism, SP）是专门针对KV Cache分布式管理而提出的新范式。其核心思想是沿序列维度切分KV Cache：

序列长度 = 8192, SP=4:
GPU 0: KV Cache for Token 0-2047
GPU 1: KV Cache for Token 2048-4095
GPU 2: KV Cache for Token 4096-6143
GPU 3: KV Cache for Token 6144-8191
​

这种切分方式使得每张GPU的KV Cache显存占用降低为1/SP，特别适合超长序列场景。

Ring Attention：去中心化的分布式注意力

Ring Attention是一种优雅的去中心化分布式注意力算法。其核心思想是让KV Cache在GPU之间环形传递，每个GPU在接收KV块的同时计算对应的注意力分数：

步骤1: GPU 0 计算 Q×K[0], GPU 1 计算 Q×K[1], ...
步骤2: KV块循环移位: GPU 0 获取 K[1], GPU 1 获取 K[2], ...
步骤3: 继续计算并累积注意力分数
...
步骤N: 所有KV块已被所有GPU处理，得到完整的注意力结果
​

Ring Attention的优势：

• 线性扩展：KV Cache可以线性分布到任意数量的GPU上
• 通信高效：每次只传递一个KV块，带宽需求可控
• 灵活适配：适合任意序列长度和GPU数量的组合

实战：vLLM中的分布式KV Cache

vLLM的分布式支持

vLLM通过集成Ray框架实现了分布式推理支持。在分布式配置中：

# vLLM 分布式推理配置示例
# TP=2, 在单机的2张GPU上运行
llm = LLM(model="meta-llama/Llama-2-70b-hf", 
          tensor_parallel_size=2)

# TP=4 + PP=2, 在2台机器×2张GPU上运行
llm = LLM(model="meta-llama/Llama-2-70b-hf",
          tensor_parallel_size=4,
          pipeline_parallel_size=2)
​

分布式PagedAttention

vLLM将PagedAttention扩展到了分布式场景。关键设计：

1. 分布式块管理器：每张GPU维护自己的KV Cache块表，由中心调度器协调全局分配
2. 按层+按头的块分配：每个KV Cache页块对应特定层和特定头的范围
3. 跨GPU显存感知：调度器知道每张GPU的可用显存，做出全局最优的块分配决策

分布式前缀缓存的实现

在分布式环境中，前缀缓存的实现需要额外考虑：

• 缓存位置透明：缓存块可能分布在任意GPU上，请求复用时需要知道缓存块的位置
• 跨GPU引用：一个请求可能需要引用分布在多张GPU上的缓存块
• 原子性保证：缓存块的分配和引用更新需要保证原子性，避免竞争条件

性能调优要点

1. TP优先于PP：张量并行的通信开销（All-Reduce）通常小于流水线并行的气泡开销，优先使用TP
2. 控制PP深度：流水线深度越大，气泡比例越高；通常PP不超过4
3. GPU亲和性：同一TP组的GPU应尽量在同一节点内，利用NVLink高带宽
4. 序列并行用于超长上下文：当序列长度超过32K时，考虑启用SP
5. KV Cache量化：FP8/INT8的KV Cache量化可以减少显存占用和传输量，代价是轻微的质量损失

小结

分布式KV Cache管理是大语言模型推理从实验室走向大规模生产的关键技术。从张量并行的头级切分到流水线并行的层级分布，再到Ring Attention的去中心化传递，不同的并行策略对KV Cache的管理提出了不同的要求和优化空间。

随着模型规模和序列长度的持续增长，分布式KV Cache技术将在推理效率中扮演越来越重要的角色。理解这些底层机制，对于构建高效、可扩展的大模型推理服务至关重要。