第 10 章 V2:EP 之外的 Buffer 与扩展能力
本章回答 V2.5 中还有哪些机制共享底层运行时,以及这些能力为何不应与 EP dispatch/combine 主路径混成一个 API。
10.1 BufferBase 统一的是生命周期
BufferBase 只统一 runtime 引用、memory usage 输出和显式 destroy 入口(deep_ep/buffers/base.py:1-27)。V2.5 README 描述 EPBuffer、EngramBuffer、PPBuffer、BucketBuffer 共用该生命周期,但各类 Buffer 的操作和数据契约保持专用。它们不是一个有多种 operation flag 的万能 buffer。
基类保留 runtime,并把显式销毁检查和 runtime 清理集中在 destroy():
class BufferBase(ABC):
"""Common lifecycle for communication buffers."""
def __init__(self, explicitly_destroy: bool = False):
self.explicitly_destroy = explicitly_destroy
self.runtime = None
@property
def main_context(self):
"""The first communication context, or ``None`` after destruction."""
return self.runtime.main_context if self.runtime is not None else None
def print_memory_usage(self) -> None:
assert self.runtime is not None
self.runtime.print_memory_usage()
@abstractmethod
def destroy(self) -> None:
assert self.explicitly_destroy
if self.runtime is not None:
self.runtime.destroy()
self.runtime = None10.2 与 EP 的关联能力
EPBuffer 还提供冗余专家负载均衡相关操作:lb_prefetch_weights 在 expert compute 前搬运副本权重,lb_reduce_grads 在 backward 聚合冗余副本的梯度(deep_ep/buffers/ep.py:846-920;C++ kernel wrappers 在 csrc/kernels/ep/prefetch_weights.hpp 和 reduce_grads.hpp)。它们复用 EP buffer 的额外 load-balance 区域,构造时需要传入容量或 BufferAllocator 计划。
Python wrapper 会把单个 Tensor 统一包装成序列,并将异步 runtime event 包装为 EventOverlap;梯度接口则将 reduction 参数直接传到底层:
def lb_prefetch_weights(self, redundant_expert_weights, expert_weights,
redundancy_mapping: torch.Tensor, num_sms: int = 0,
previous_event: Optional[EventHandle] = None) -> EventOverlap:
redundant_expert_weights = ([redundant_expert_weights] if isinstance(redundant_expert_weights, torch.Tensor)
else list(redundant_expert_weights))
expert_weights = [expert_weights] if isinstance(expert_weights, torch.Tensor) else list(expert_weights)
num_sms = self.lb_get_theoretical_num_sms() if num_sms == 0 else align(num_sms, 2)
return EventOverlap(self.runtime.lb_prefetch_weights(
redundant_expert_weights, expert_weights, redundancy_mapping, num_sms, previous_event))
def lb_reduce_grads(self, redundant_expert_grads: torch.Tensor,
expert_grads: torch.Tensor, redundancy_mapping: torch.Tensor,
num_sms: int = 0,
previous_event: Optional[EventHandle] = None) -> EventOverlap:
num_sms = self.lb_get_theoretical_num_sms() if num_sms == 0 else align(num_sms, 2)
return EventOverlap(self.runtime.lb_reduce_grads(
redundant_expert_grads, expert_grads, redundancy_mapping, num_sms, previous_event))该功能存在调用顺序契约:先按 redundancy mapping 配置并预取权重,再用副本执行 expert 计算;反向阶段再将副本梯度规约到原专家。源码可确认这些 API 和 C++ kernel 的存在;具体如何与 MoonEP、UltraEP 等上层调度结合,由相应框架集成代码定义。
10.3 实验性通信原语
README 把 Engram、PP 和 bucket collectives 标注为实验能力。EngramBuffer 面向远程记忆存储 fetch/wait;PPBuffer 提供 PP send/recv;BucketBuffer 支持分桶的 all-gather、reduce-scatter、all-reduce 与普通 PyTorch tensor session。它们分别对应不同数据生命周期和消费者,放在单独章节脉络比较比塞入 EP 黄金路径更准确。
flowchart TB
B[BufferBase 生命周期] --> E[EPBuffer: MoE dispatch / combine]
B --> G[EngramBuffer: 远程存储 fetch / wait]
B --> P[PPBuffer: pipeline send / recv]
B --> K[BucketBuffer: 分桶 collective]
E --> L[load-balance 预取权重 / 梯度规约]10.4 边界与风险
- V1 的 normal/low-latency 旧 API 不应被当成当前 EPBuffer 的 alias;V2.5 README 明确它们已移除。
- Engram/PP/Bucket 在当前 README 中标为实验特性;使用时需按对应测试与硬件条件判断成熟度。
- 新增 buffer 类型共享生命周期基类,不代表能够共享同一 workspace size、消息协议或失败语义。
本章收束
- V2.5 扩大了 DeepEP 的系统范围,但 EPBuffer 仍是本书核心主线。
- BufferBase 统一生命周期和诊断入口,各 buffer 保留专属数据面。
- load-balance 原语扩展 EP 的资源用途;Engram、PP、Bucket 是独立通信场景。