PD分离与EP分离技术

一、Prefill与Decode的底层原理剖析

1.1 Prefill阶段的数学本质

（以Transformer架构为例）

计算密集型的核心原因：

# 自注意力计算伪代码
Q = X @ W_Q  # [batch, seq_len, d_model] → [batch, seq_len, d_k]
K = X @ W_K  # O(n^2)复杂度开始显现
V = X @ W_V

attn = (Q @ K.transpose(-2, -1)) / sqrt(d_k)  # [batch, seq_len, seq_len]
attn = softmax(attn)
output = attn @ V  # 立方级计算复杂度

深入解析：

• 注意力机制的平方复杂度：
  Prefill阶段的核心瓶颈在于自注意力计算。对于长度为N的序列，需要计算每个token与其他所有token之间的关联，这导致了O(N²)的计算复杂度。这意味着，当输入序列长度增加一倍时，计算量将增加四倍。
• KV Cache的构建：
  Prefill阶段还需要构建KV Cache，用于存储每个token的key和value向量。这些向量将在Decode阶段被重复使用，以加速注意力计算。KV Cache会占用大量的显存。

# 以Llama-7B模型为例（hidden_size=4096）
kv_cache_per_layer = 2 * batch * seq_len * num_heads * head_dim
# 当batch=1, seq=2048, layers=32时：
总缓存量 = 32 * 2 * 2048 * 4096 * 1 = 536MB

1.2 Decode阶段的访存瓶颈

典型访存模式分析：

for token in output_sequence:  # 自回归循环
    # 读取当前token的embedding（1次显存访问）
    x = embedding[token]

    # 逐层处理（每层需要访问KV Cache）
    for layer in layers:
        # 从HBM读取该层的KV Cache（2次显存访问）
        k_cache = load_from_hbm(layer.k_cache)
        v_cache = load_from_hbm(layer.v_cache)

        # 计算新的k,v并更新缓存（2次显存写入）
        new_k = x @ W_k
        new_v = x @ W_v
        k_cache.append(new_k)
        v_cache.append(new_v)

        # 注意力计算（需要读取整个k_cache和v_cache）
        attn_output = attention(q, k_cache, v_cache)

    # 写入新生成的token（1次显存访问）
    output_sequence.append(next_token)

深入解析：

• 自回归特性导致的访存压力：
  Decode阶段是自回归的，这意味着每个token的生成都依赖于之前生成的token。因此，需要频繁地读取和更新KV Cache，这导致了大量的显存访问。随着生成token的数量增加，KV Cache也会越来越大，导致更多的访存压力。
• KV Cache的重复读取：
  在decode阶段，每一层都会重复读取全部的KV Cache，这是导致访存瓶颈的主要原因。

关键指标对比：

阶段	计算量(FLOPs)	显存访问量	并行度
Prefill	O(N²d)	O(Nd)	高
Decode	O(Nd)	O(N²d)	低

二、PD分离的硬件级优化

2.1 异构计算架构设计

Prefill节点的典型配置：

• NVIDIA H100 GPU（FP8 Tensor Core）
• 内存带宽：3TB/s
• 算力：989 TFLOPS
• 优化目标：最大化SM（流式多处理器）利用率

Decode节点的特殊设计：

• 定制化推理卡（如Groq LPU）
• 内存带宽：5TB/s（HBM3E）
• 片上SRAM：230MB（减少HBM访问）
• 优化目标：降低访存延迟

深入解析：

• 硬件差异化的本质：
  PD分离的核心在于利用不同硬件在计算和访存上的优势。Prefill阶段需要强大的计算能力，而Decode阶段更需要高带宽的内存。
  PD分离部署的核心在于将LLM推理过程解耦为两个在资源需求和执行特点上显著不同的阶段，并针对性地进行优化：
• 逻辑上的流水线 (Pipeline Stages)：
  PD分离将推理过程视为一个两阶段的流水线。一个推理请求首先进入Prefill阶段，然后将Prefill阶段的输出作为Decode阶段的输入。
• 数据流转 (Data Flow)：
  • Prefill阶段：接收用户的prompt（token序列），模型对整个prompt进行并行计算，生成每个token的embedding和attention信息，并构建关键的上下文状态，即KV Cache。Prefill阶段的输出主要是这个KV Cache。
  • Decode阶段：接收Prefill阶段生成的KV Cache，以及当前已生成的token序列（初始为空）。模型根据KV Cache和已生成的token预测下一个token的概率分布，并从中采样生成下一个token。然后，将新生成的token加入到序列中，并更新KV Cache，重复这个过程直到生成完整的回复或达到预设的停止条件。
• 资源隔离与定制 (Resource Isolation and Customization)：
  • Prefill节点：通常部署在具有强大计算能力的硬件上（例如，配备大量高性能GPU），因为 Prefill阶段需要大量的并行计算。
  • Decode节点：通常部署在具有较大内存容量和高内存带宽的硬件上，因为Decode阶段需要频繁地读写和更新KV Cache。
• 状态管理 (State Management)：
  Prefill阶段生成的KV Cache是Decode阶段的关键状态。在分离部署中，需要有效地将KV Cache从Prefill节点传递到Decode节点。对于并发推理，可能需要管理多个请求的KV Cache。
• 任务调度与协调 (Task Scheduling and Coordination)：
  需要一个系统来协调和调度Prefill和Decode阶段的任务。例如，当一个请求完成Prefill后，需要将其KV Cache路由到可用的Decode节点进行后续的生成。
• 通信机制 (Communication Mechanism)：
  Prefill节点和Decode节点之间需要高效的通信机制来传输KV Cache和其他控制信息。这在分布式部署中尤其重要。

三、EP分离与门控网络的深度协同

3.1 MoE架构的运行时解析

专家并行执行流程：

# 假设有8个专家分布在2个GPU上
def moe_layer(x):
    # 门控网络计算（GPU0）
    gates = gating_network(x)  # [batch, num_experts]

    # 专家路由（跨设备通信）
    selected_experts = topk(gates, k=2)
    x_shards = alltoall(x, selected_experts)  # 跨设备传输

    # 专家计算（分布式执行）
    expert_outputs = []
    for expert_id in selected_experts:
        device = expert_id // 4  # 每个GPU存放4个专家
        with device(device):
            out = experts[expert_id](x_shards[expert_id])
            expert_outputs.append(out)

    # 结果聚合（跨设备规约）
    combined = allreduce(expert_outputs)
    return combined * gates  # 门控加权

深入解析：

• EP作用于Decode阶段的专家网络：
  EP（专家并行）主要应用于Decode阶段的专家网络。在MoE模型中，Decode阶段的FFN层被拆分为多个“专家”模块。EP通过将这些专家模块分配到不同的硬件设备上，实现并行计算，加速推理过程。
• 门控网络的本质：
  门控网络的作用是根据输入数据的特征，动态地选择激活哪些专家。门控网络的输出，本质上就是对于多个专家的一个概率分布。
• 专家路由的挑战：
  EP分离的核心挑战在于如何高效地进行专家路由，即将输入数据分配给不同的专家。All to all的通信，是专家路由中的主要通信方式。
• 专家并行的优势：
  通过将不同的专家分配到不同的硬件设备上，可以实现并行计算，加速推理过程。

什么是Expert Parallelism (EP)?

在拥有Mixture of Experts (MoE) 架构的大型语言模型中，模型包含多个“专家”子网络。Expert Parallelism 是一种将这些不同的专家子网络分布到多个计算设备（例如GPU）上的并行计算技术。

• 每个计算设备负责存储和运行模型中一部分的专家。
• 当输入数据进入模型时，一个“门控网络”（gating network）会决定将哪些输入路由到哪些专家进行处理。
• 只有被选中的专家会对特定的输入进行计算，最终它们的输出会被合并以产生最终结果。
• 这种方法允许模型拥有非常庞大的参数量（因为有很多专家），但对于每个单独的输入，只有一小部分参数是激活的，从而在一定程度上控制了计算成本。

EP与PD分离的关系：

• Expert Parallelism 和 PD 分离是两种不同的、但可以互补的优化策略，它们关注推理过程的不同方面：
  • 关注点不同：
    • PD分离 (Prefill/Decode Separation) 主要关注推理流程的两个主要阶段，并针对它们不同的资源需求进行优化部署，将计算密集型的Prefill阶段和访存密集型的Decode阶段分开处理。
    • EP (Expert Parallelism) 主要关注模型内部的并行化，特别是针对MoE架构，通过将模型中的不同专家分布到多个设备上，实现模型参数的并行处理和存储。
  • 解决的问题不同：
    • PD分离 主要解决的是Prefill和Decode阶段的资源竞争问题，提高整体的吞吐量和降低延迟。
    • EP 主要解决的是MoE模型参数量过大，单个设备可能无法存储或高效计算的问题，通过分布式计算来扩展模型容量。
  • 联系与协同：
    • 在一个使用MoE架构的LLM推理系统中，可以同时采用EP和PD分离策略。
    • 例如，在Decode阶段，模型需要访问KV Cache并进行后续的token生成。如果模型采用了EP，那么Decode阶段的不同专家可能分布在不同的设备上，需要通过网络进行通信。PD分离可以将Decode阶段整体部署在特定的节点上，而这些节点内部又可以通过EP来并行计算不同的专家。
    • 在Prefill阶段，如果模型是MoE的，那么Prefill过程也需要涉及到专家的计算，同样可以利用EP将这些计算分布到多个设备上。
  • 总结来说：
    • PD分离 是从推理流程的宏观阶段来优化资源分配和并行执行。
    • EP 是从模型内部的微观结构（特别是MoE模型中的专家）来实现参数和计算的并行化。
    • 它们可以结合使用，共同提升基于MoE架构的LLM的推理性能和可扩展性。例如，一个推理集群可以采用PD分离，将Prefill和Decode部署在不同的节点组上，而在执行Decode阶段时，如果模型是MoE的，那么每个Decode节点组内部又可以通过Expert Parallelism来并行处理不同的模型专家。

3.2 门控网络的硬件优化

动态路由的三种实现方式：

1. 软件路由（Baseline）
   • 在CPU上执行门控计算，产生3-5ms额外延迟
2. 硬件加速（NVIDIA MaaS方案）
   • 使用GPU上的TensorRT插件实现路由，延迟降至0.2ms
3. 专用路由单元（Groq方案）
   • 在LPU中集成路由协处理器，延迟<0.05ms

路由算法演进：

# 传统softmax门控
gates = softmax(x @ W_g)

# 最新研究成果（DeepSeek-V2）
sparse_gates = entmax(x @ W_g)  # 产生更稀疏的专家选择
load_balance_loss = CV(gates)  # 控制专家负载均衡

深入解析：

• 硬件加速的必要性：
  门控网络的计算和专家路由的通信，会产生额外的延迟。因此，需要使用硬件加速来降低这些延迟。
• 稀疏门控的优势：
  稀疏门控可以减少激活的专家数量，从而降低计算和通信开销。负载均衡损失，能够有效地控制专家的使用频率。

门控网络（Gating Network）是什么，一个层么？详细介绍一下找到对应专家的大致过程

• 门控网络的概念：
  门控网络是MoE模型的核心组件，它负责根据输入token的特征，动态选择激活哪些专家。
  门控网络通常是一个小型神经网络，它可以是一个或多个全连接层（FFN）。
• 找到对应专家的大致过程：
  1. 输入token的特征向量被输入到门控网络中。
  2. 门控网络输出一个或多个权重向量，每个权重向量对应一个专家。
  3. 权重向量中的每个元素表示对应专家的激活程度。
  4. 模型根据权重向量，选择激活一个或多个权重最高的专家。
  5. 被激活的专家处理输入token，并将它们的输出进行加权融合，得到最终的输出。

能详细介绍一下FFN层么？本质上为什么会在这层激活专家呢？如何找到对应的专家？

• FFN层：
  • FFN层（Feed-Forward Network）是Transformer模型中的一个重要组件，它由两个或多个全连接层组成，用于对输入特征进行非线性变换。
  • FFN层的作用是增强模型的表达能力，使其能够学习更复杂的特征表示。
• 在FFN层激活专家：
  • 在MoE模型中，FFN层被拆分为多个“专家”模块，每个专家都是一个独立的FFN层。
  • 这样做的好处是，可以让模型根据不同的输入数据，选择激活不同的专家，从而提高模型的性能。
  • 本质上来说，FFN层是模型中参数量最大的层，也是模型中最具有特征提取能力的层，将FFN层进行专家化，能够有效提高模型的性能。
• 如何找到对应的专家：
  • 通过门控网络，门控网络会根据输入token的特征，输出每个专家的权重。
  • 根据权重的大小，来选择对应的专家。

专家权重是动态计算的，而不是预先计算好的

• 动态计算：
  • 门控网络的作用是根据当前的输入token的特征，动态地计算每个专家的权重。
  • 这意味着，对于不同的输入token，门控网络会产生不同的权重分布。
  • 这种动态性是MoE模型的核心优势之一，它允许模型根据输入数据的不同，灵活地选择最合适的专家进行处理。
• 计算过程：
  • 当一个输入token被送入MoE模型时，它的特征向量会被传递给门控网络。
  • 门控网络通过一系列的计算（通常是全连接层和激活函数），输出一个表示每个专家权重的向量。
  • 这个权重向量中的每个元素，代表了对应专家应该被激活的程度。
  • 因此，每次输入一个新的token，门控网络都会重新进行计算，生成新的专家权重。
• 原因：
  • 模型需要根据不同的输入，选择不同的专家。如果权重是预先计算好的，那么模型就失去了动态调整的能力，无法适应不同的输入模式。
• 总结：
  • 门控网络会根据当前输入token的特征，动态计算每个专家的权重。
  • 这种动态的计算，是MoE模型的核心，也是MoE模型能够提高模型表现的原因。

门控网络有些类似分类网络

• 门控网络与分类的相似性：
  • 目标：两者都涉及将输入数据分配到不同的“类别”。在门控网络中，“类别”就是不同的专家。两者都使用类似的数学工具，例如Softmax函数，将输出转换为概率分布。
  • 机制：门控网络通常使用全连接层，这也是分类网络中常用的结构。门控网络的输出可以被视为输入数据属于每个专家的“概率”。
• 门控网络与分类的差异性：
  • 应用场景：分类任务通常用于将数据分配到预定义的类别，例如图像分类或文本分类。门控网络用于MoE模型中，用于动态选择激活哪些专家来处理输入数据。
  • 目标的弹性：分类任务通常是“硬性”的，即每个输入数据必须被分配到一个明确的类别。门控网络则更加“软性”，它允许模型根据输入数据的特征，对不同的专家进行加权，甚至同时激活多个专家。
• 本质：
  • 分类网络，着重于将输入数据进行明确的分类。
  • 门控网络，着重于根据输入数据，来选择最适合处理该数据的专家。
• 总结：
  • 门控网络可以被视为一种特殊的分类任务，其目标是根据输入数据的特征，选择最适合处理该数据的专家。
  • 然而，门控网络与传统分类任务之间存在一些关键差异，例如应用场景和目标的弹性。
  • 因此，虽然门控网络在某些方面与分类任务相似，但它们并非完全相同。它们是为了实现MoE模型中的动态专家选择而设计的。

四、PD与EP的联合优化实践

4.1 资源分配策略：混合部署架构

                  [Load Balancer]
                 /                \
      Prefill Cluster           Decode Cluster
      (H100 x4)           (LPU x8 + EP组)
                                /       \
                      文本专家组    多模态专家组
                      (LPU0-3)      (LPU4-7)

深入解析：

• 资源差异化分配：
  • PD与EP的联合优化，核心在于根据Prefill和Decode阶段的计算和访存需求，以及专家网络的特性，对硬件资源进行差异化分配。
  • Prefill集群使用擅长计算的GPU，而Decode集群则使用擅长访存的推理卡，并结合EP技术，将专家网络分配到多个推理卡上。
• 负载均衡：
  • Load Balancer负责将用户请求合理分配到Prefill和Decode集群，确保各个硬件资源的充分利用。
• 多模态支持：
  • 通过将不同模态的专家分配到不同的专家组，可以实现多模态的支持。

4.2 通信优化关键技术：KV Cache的压缩传输

# Prefill节点生成KV Cache后：
compressed_kv = FP8_quantize(kv_cache)  # 4:1压缩比
encrypted_data = AES_gcm_encrypt(compressed_kv)

# Decode节点接收时：
decrypted = AES_gcm_decrypt(encrypted_data)
kv_cache = FP8_dequantize(decrypted)

深入解析：

• 压缩传输的必要性：
  • Prefill和Decode集群之间需要传输大量的KV Cache。通过压缩KV Cache，可以减少传输的数据量，从而降低通信延迟。
• RDMA网络：
  • 结合RDMA网络，可以实现低延迟的跨集群通信。

4.3 关键突破：专家负载预测、门控网络协同、跨集群调度

• 专家负载预测：
  通过预测专家的负载情况，可以实现更精细的资源分配和调度。
• 门控网络协同：
  将门控网络的设计与硬件特性相结合，可以进一步提高路由效率。
• 跨集群调度：
  通过优化跨集群的流水线调度，可以减少整体的推理延迟。

五、技术演进与开发者建议

5.1 前沿研究方向

1. 动态PD调度
   根据请求特征自动分配Prefill/Decode资源
2. 3D混合并行
   在EP基础上结合Tensor/Pipeline并行
3. 存算一体架构
   • 使用HBM直接处理KV Cache

六、总结

大模型推理优化的三个维度：

1. 时间维度：解耦Prefill/Decode实现流水线并行
2. 空间维度：通过EP实现计算资源的空间划分
3. 硬件维度：针对不同计算模式定制加速单元

原文地址: https://www.cnblogs.com/menkeyi/p/18778750