DeepEP 测试基础设施与通用工具深度分析

DeepEP/
├── tests/
│   ├── elastic/
│   │   ├── test_ep.py          # 核心: Dispatch/Combine 正确性 + 性能测试
│   │   ├── test_agrs.py        # AllGather/ReduceScatter 测试
│   │   ├── test_barrier.py     # Barrier 同步原语测试
│   │   ├── test_engram.py      # Engram 内存管理测试
│   │   └── test_pp.py          # Pipeline Parallelism 测试
│   ├── legacy/                 # 旧版 (V1) 测试
│   └── utils/                  # 工具测试
└── deep_ep/utils/
    ├── __init__.py             # 导出 EventHandle
    ├── testing.py              # bench / bench_kineto / flush_l2_cache
    ├── gate.py                 # 非均衡 MoE 流量生成 (Unbalanced Score)
    ├── comm.py                 # NCCL Communicator 封装
    ├── envs.py                 # 分布式初始化 / 硬件探测
    ├── event.py                # EventOverlap CUDA 事件封装
    ├── refs.py                 # PyTorch 参考实现 (dispatch / combine)
    ├── math.py                 # FP8 编解码 / 对齐 / 差异计算
    ├── semantic.py             # weak_lru / value_or
    └── find_pkgs.py            # NVIDIA 包路径探测

# envs.py:73-113
def init_dist(local_rank: int, num_local_ranks: int, seed: int = 0) -> Tuple[int, int, dist.ProcessGroup]:
    ip = os.getenv('MASTER_ADDR', '127.0.0.1')
    port = int(os.getenv('MASTER_PORT', '8361'))
    num_nodes = int(os.getenv('WORLD_SIZE', 1))
    node_rank = int(os.getenv('RANK', 0))

    sig = inspect.signature(dist.init_process_group)
    params = {
        'backend': 'nccl',
        'init_method': f'tcp://{ip}:{port}',
        'world_size': num_nodes * num_local_ranks,
        'rank': node_rank * num_local_ranks + local_rank,
    }
    if 'device_id' in sig.parameters:
        params['device_id'] = torch.device(f'cuda:{local_rank}')
    dist.init_process_group(**params)
    torch.set_default_dtype(torch.bfloat16)
    torch.set_default_device('cuda')
    torch.cuda.set_device(local_rank)

    init_seed(seed)
    return dist.get_rank(), dist.get_world_size(), dist.new_group(list(range(num_local_ranks * num_nodes)))

# test_ep.py:609
torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes)

flowchart TD
    A["__main__: argparse 解析参数"] --> B["torch.multiprocessing.spawn<br/>启动 N 个 test_loop 进程"]
    B --> C["init_dist: 设置 MASTER_ADDR/PORT"]
    C --> D["dist.init_process_group<br/>(NCCL backend)"]
    D --> E["torch.set_default_dtype(bfloat16)<br/>torch.set_default_device(cuda)"]
    E --> F["init_seed: global_seed + rank"]
    F --> G["dist.new_group 创建全局 group"]
    G --> H["construct_elastic_buffer<br/>构建 ElasticBuffer"]
    H --> I["test_dispatch_combine<br/>执行正确性 + 性能测试"]
    I --> J{"do_pressure_test?"}
    J -->|Yes| K["循环 seed in range(1e9)<br/>重建 buffer 并重复测试"]
    J -->|No| L["buffer.destroy()"]
    K --> L
    L --> M["dist.destroy_process_group()"]

# testing.py:12-21
def flush_l2_cache(enabled: bool = True):
    l2_flush_cache_size = 256e6
    if enabled:
        torch.empty(int(l2_flush_cache_size // 4), dtype=torch.int, device='cuda').zero_()

# testing.py:24-60
def bench(fn, num_warmups: int = 50, num_tests: int = 50,
          post_fn: Optional[Callable] = None, flush_l2: bool = True):
    torch.cuda.synchronize()
    for _ in range(num_warmups):
        fn()
    start_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
    end_events = [torch.cuda.Event(enable_timing=True) for _ in range(num_tests)]
    for i in range(num_tests):
        flush_l2_cache(flush_l2)
        start_events[i].record()
        fn()
        end_events[i].record()
        if post_fn is not None:
            post_fn()
    torch.cuda.synchronize()
    times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:]
    return np.average(times), np.min(times), np.max(times)

# testing.py:111-219
def bench_kineto(fn, kernel_names: Union[str, tuple], num_tests: int = 30,
                 suppress_kineto_output: bool = False, trace_path: Optional[str] = None,
                 flush_l2: bool = True, barrier_comm_profiling: bool = False,
                 num_kernels_per_period: int = 1, barrier: Optional[Callable] = None):

flowchart TD
    A["检查 EP_USENVIDIA_TOOLS 环境变量"] -->|已设置| Z["返回 1 (跳过)"]
    A -->|未设置| B["fn() 预热"]
    B --> C["torch.profiler.profile 配置<br/>schedule: wait=0, warmup=1, active=1"]
    C --> D["循环 num_tests 次"]
    D --> E["flush_l2_cache"]
    E --> F{"barrier_comm_profiling?"}
    F -->|Yes| G["torch.cuda._sleep(2e7)<br/>dist.all_reduce 同步"]
    F -->|No| H["fn()"]
    G --> H
    H --> I["profiler.step()"]
    I --> J["解析 key_averages().table()"]
    J --> K["匹配 kernel_names<br/>提取 time × count"]
    K --> L{"num_kernels_per_period > 1?"}
    L -->|Yes| M["Chrome trace JSON 解析<br/>按时间戳排序分段聚合"]
    L -->|No| N["返回平均 duration"]
    M --> N

# testing.py:72-108
class suppress_stdout_stderr:
    def __enter__(self):
        self.outnull_file = open(os.devnull, 'w')
        self.errnull_file = open(os.devnull, 'w')
        self.old_stdout_fileno = os.dup(sys.stdout.fileno())
        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
        os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
        ...

# gate.py:116-137
def get_precise_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio):
    rank_count = generate_rank_count(num_tokens, num_experts, num_ranks, num_topk, ratio)
    threshold = 0.9
    scores = torch.empty((num_tokens, num_experts), dtype=torch.float32, device='cuda')
    scores.uniform_(to=threshold)                  # 低分填充
    topk_idx = generate_topk_idx(rank_count, ...)
    topk_scores = torch.empty((num_tokens, num_topk), ...)
    topk_scores.uniform_(threshold + 1e-6, 1.0)    # 高分填充
    scores[row_idx, topk_idx] = topk_scores
    return scores

# gate.py:167-173
def get_random_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio):
    factor = 1.0
    if ratio != 1.0:
        factor = map_unbalanced_ratio_to_factor(num_tokens, num_experts, num_ranks, num_topk, ratio)
    return get_scores_by_factor(num_tokens, num_experts, num_ranks, factor)

# gate.py:148-164
def map_unbalanced_ratio_to_factor(...):
    factor_l, factor_r = 1.0, 100.0
    for _i in range(num_iterations):  # 20 次迭代
        factor_mid = (factor_l + factor_r) / 2
        scores = get_scores_by_factor(...)
        _, topk_idx = torch.topk(scores, num_topk, ...)
        rank_idx = topk_idx // num_experts_per_rank
        one_hot = torch.nn.functional.one_hot(rank_idx, num_ranks)
        counts = one_hot.any(dim=1).to(torch.float).sum(dim=0)
        if counts[0].item() > counts[1:].mean().item() * ratio:
            factor_r = factor_mid
        else:
            factor_l = factor_mid
    return factor_l

# gate.py:32-113
def generate_rank_count(num_tokens, num_experts, num_ranks, num_topk, ratio):
    # 1. 随机 topk 选择，计算每个 token 覆盖的 rank 数 a[i]
    random_scores = torch.rand(num_tokens, num_experts, device='cuda')
    topk_weights_, topk_indices = torch.topk(random_scores, num_topk, ...)
    topk_indices //= num_experts_per_rank
    sorted_topk_indices = torch.sort(topk_indices, dim=1)[0]
    topk_indices_diff_mask = sorted_topk_indices[:, 1:] != sorted_topk_indices[:, :-1]
    a = topk_indices_diff_mask.sum(dim=1) + 1

    # 2. 计算特殊 rank 必须/可选出现的 token 数
    upper_bound_per_token = int(num_normal_ranks / ratio) + 1
    a = torch.clamp(a, None, upper_bound_per_token)
    must_mask = (a == num_ranks)
    ...

    # 3. 生成 rank 排列并插入特殊 rank
    rank_perm = (torch.randperm(num_normal_ranks, device='cuda') + 1).repeat(...)
    rank_perm_with_special_rank = torch.zeros(num_tokens * num_topk, ...)
    # ... 复杂的索引计算 ...

    # 4. scatter_add 生成最终 count
    rank_count = torch.zeros((num_tokens, num_ranks), dtype=torch.int32, device='cuda')
    rank_count.scatter_add_(dim=1, index=result, src=torch.ones_like(result, dtype=torch.int32))
    return rank_count

flowchart TD
    A["get_unbalanced_scores<br/>(ratio, precise)"] -->|precise=True| B["get_precise_unbalanced_scores"]
    A -->|precise=False| C["get_random_unbalanced_scores"]
    B --> D["generate_rank_count<br/>生成 (num_tokens, num_ranks) count 矩阵"]
    D --> E["generate_topk_idx<br/>count → 具体 Expert 索引"]
    E --> F["高低分分离填充 scores"]
    C --> G["map_unbalanced_ratio_to_factor<br/>二分搜索 factor"]
    G --> H["get_scores_by_factor<br/>rank 0 用 factor 加权"]
    F --> I["返回 scores (num_tokens, num_experts)"]
    H --> I

# comm.py:11-37
class NCCLCommHandle:
    def __init__(self, nccl_comm: int, managed: bool):
        self.nccl_comm = nccl_comm
        self.managed = managed
        self.destroy = _C.destroy_nccl_comm

    def __del__(self):
        if self.managed:
            self.destroy(self.nccl_comm)

# comm.py:42-75
_storage = dict()

def get_nccl_comm_handle(group: dist.ProcessGroup, force_new_comm: bool = False) -> NCCLCommHandle:
    global _storage
    if not force_new_comm and group in _storage:
        return _storage[group]

    backend = group._get_backend(torch.device('cuda'))
    if not force_new_comm and hasattr(backend, '_comm_ptr') and int(os.getenv('EP_REUSE_NCCL_COMM', '1')):
        _storage[group] = NCCLCommHandle(backend._comm_ptr(), False)
        return _storage[group]

    # 旧版 PyTorch: 手动 all_gather unique_id 后创建新 comm
    nccl_unique_ids = [None, ] * group.size()
    dist.all_gather_object(nccl_unique_ids, _C.get_local_nccl_unique_id(), group)
    root_unique_id = nccl_unique_ids[0]
    key = time.time_ns() if force_new_comm else group
    _storage[key] = NCCLCommHandle(
        _C.create_nccl_comm(root_unique_id, group.size(), group.rank()), True)
    return _storage[key]

# envs.py:24-35
def init_seed(global_seed: int) -> None:
    global _local_seed, _global_seed
    _local_seed = global_seed + dist.get_rank()
    _global_seed = global_seed
    torch.manual_seed(_local_seed)
    random.seed(_local_seed)

# envs.py:58-70
def dist_print(s: str = '', once_in_node: bool = False) -> None:
    global _local_rank
    assert _local_rank is not None
    if not once_in_node or _local_rank == 0:
        print(s, flush=True)
    dist.barrier()

# envs.py:116-142
def get_physical_domain_size(group) -> Tuple[int, int]:
    return _C.get_physical_domain_size(get_nccl_comm_handle(group).get())

def get_logical_domain_size(group, allow_hybrid_mode=True) -> Tuple[int, int]:
    return _C.get_logical_domain_size(get_nccl_comm_handle(group).get(), allow_hybrid_mode)

# envs.py:192-219
@functools.lru_cache()
def get_nvlink_gbs(factor: float = 0.9) -> float:
    result = subprocess.run(['nvidia-smi', 'nvlink', '-s'], ...)
    pattern = r'GPU \d+:.*?(?=^GPU \d+:|^$)'
    match = re.search(pattern, output, re.MULTILINE | re.DOTALL)
    gpu_block = match.group(0)
    link_pattern = r'Link \d+:\s*([\d\.]+) GB/s'
    link_matches = re.findall(link_pattern, gpu_block)
    return sum(float(bw) for bw in link_matches) * factor

# envs.py:245-268
@functools.lru_cache()
def get_rdma_gbs(nic_name: str = _DEFAULT_NIC_NAME) -> float:
    result = subprocess.run(['ibstat'], ...)
    pattern = rf"CA '{nic_name}'.*?Port \d+:\s*.*?Rate:\s*(\d+)"
    match = re.search(pattern, output, re.DOTALL)
    rate = int(match.group(1))
    return rate / 8  # Gb/s → GB/s

# envs.py:222-242
@functools.lru_cache()
def check_fast_rdma_atomic_support(nic_name: str = _DEFAULT_NIC_NAME) -> bool:
    result = subprocess.run(['ibstat'], ...)
    pattern = rf"CA '{nic_name}'.*?CA type:\s*(\S+)"
    match = re.search(pattern, output, re.DOTALL)
    return match.group(1) == 'MT4131'

# event.py:8-96
class EventOverlap:
    def __init__(self, event: Optional[EventHandle] = None,
                 extra_tensors: Optional[Tuple[torch.Tensor]] = None):
        self.event = event
        self.extra_tensors = extra_tensors
        self._release_handle_by_call = False
        self.hook_after_wait: Optional[Callable] = None

# event.py:39-54
def current_stream_wait(self, release_handle: bool = False) -> None:
    assert self.event is not None
    self.event.current_stream_wait()
    if self.hook_after_wait is not None:
        self.hook_after_wait()
        self.hook_after_wait = None
    if release_handle:
        self.event = None

# event.py:74-96
def __enter__(self) -> Any:
    return self

def __exit__(self, exc_type, exc_val, exc_tb) -> None:
    if self.event is not None:
        self.current_stream_wait(release_handle=self._release_handle_by_call)
    self._release_handle_by_call = False

event_overlap = event_after_all_to_all_kernels()
with event_overlap:
    do_something_on_current_stream()  # 与通信并行执行
# 退出 with 时自动 current_stream_wait

sequenceDiagram
    participant Compute as Compute Stream
    participant Comm as Communication Stream
    participant Event as EventOverlap

    Compute->>Comm: dispatch / combine 启动
    Comm->>Event: event.capture()
    Compute->>Compute: do_something_on_current_stream()
    Compute->>Event: __exit__ → current_stream_wait()
    Event->>Comm: wait for event
    Comm->>Compute: 通信完成，继续执行

# refs.py:10-123
def dispatch(x, topk_idx, topk_weights, num_max_tokens_per_rank, num_experts):
    # 1. 按目标 rank 分组发送数据
    for dst_rank_idx in range(num_ranks):
        expert_start_idx = dst_rank_idx * num_experts_per_rank
        expert_end_idx = expert_start_idx + num_experts_per_rank
        mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
        indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
        # ... 选择 x, sf, topk_idx, topk_weights ...

    # 2. all_to_all_single 交换大小和数据
    dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
    dist.all_to_all_single(recv_x, send_x, ...)
    # ... 其他张量同理 ...

    # 3. 本地 Expert 索引转换 + mask
    recv_topk_idx = recv_topk_idx - expert_start_idx
    recv_topk_idx.masked_fill_(~mask, -1)

# refs.py:177-243
def combine(y, topk_idx, num_scaleout_ranks, num_scaleup_ranks, num_experts,
            bias, reduce_in_local, reduce_in_scaleup):

# refs.py:201-234
def grouped_reduce(data_to_reduce, group_id):
    group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
    # 重排 data 使同组连续
    data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices]...
    # 逐 topk 维度扫描，检测 segment break
    for i in range(num_topk):
        is_segment_break = ... group_id[:, i] != group_id[:, i + 1]
        cur_accum_buf += data_to_reduce[:, i, :].float()
        data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[...]
        cur_accum_buf[segment_break_token_indices] = 0.0

# refs.py:126-153
def generate_pre_combine_data(src_token_global_idx, num_max_tokens_per_rank, num_topk, hidden):
    token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
                   torch.arange(num_topk, ...).unsqueeze(0))
    max_seed = num_ranks * num_max_tokens_per_rank * num_topk
    result = torch.sin(
        (((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
        torch.arange(1, hidden + 1, ...) +
        math.sin(float(get_global_seed()))
    )
    return result.to(torch.bfloat16)

# math.py:30-56
@torch.compile(dynamic=True)
def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    # 1. align hidden 到 128
    x_padded = torch.nn.functional.pad(x, (0, aligned_n - n), mode='constant', value=0)
    # 2. 按 128 分组计算 amax
    x_padded_view = x_padded.view(m, -1, 128)
    x_amax = x_padded_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
    # 3. scale = 448 / amax (FP8 E4M3 max = 448)
    return (x_padded_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view(m, aligned_n)[:, :n].contiguous(), \
           (x_amax / 448.0).view(m, -1)

# math.py:5-9
def calc_diff(x: torch.Tensor, y: torch.Tensor) -> float:
    x, y = x.double() + 1, y.double() + 1
    denominator = (x * x + y * y).sum()
    sim = 2 * (x * y).sum() / denominator
    return (1 - sim).item()

# test_ep.py:94-98
x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
x = per_token_cast_to_fp8(x) if use_fp8_dispatch else x
bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') if num_bias == 1 else None

# test_ep.py:78-81
if args.masked_ratio > 0:
    rand_mask = torch.rand_like(topk_idx, dtype=torch.float)
    topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1)
    topk_weights.masked_fill_(topk_idx < 0, 0)

# test_ep.py:22-31
def enumerate_ep_modes():
    for do_handle_copy in (1, 0):
        for expert_alignment in (128, 1):
            for use_fp8_dispatch in (1, 0):
                for num_bias in (0, 1, 2):
                    for with_previous_event in (0, 1):
                        for async_with_compute_stream in (0, 1):
                            for allocate_on_comm_stream in ((1, ) if with_previous_event else (0, 1)):
                                yield (do_handle_copy, expert_alignment, use_fp8_dispatch, num_bias,
                                       with_previous_event, async_with_compute_stream, allocate_on_comm_stream)

# test_ep.py:503-508
assert torch.equal(combined_x, ref_combined_y), \
    f'Diff: {calc_diff(combined_x, ref_combined_y)}'
assert torch.equal(reduced_combined_x, ref_reduced_combined_y), \
    f'Diff: {calc_diff(reduced_combined_x, ref_reduced_combined_y)}'
assert torch.equal(combined_topk_weights, topk_weights), \
    f'{calc_diff(combined_topk_weights, topk_weights)}'

# test_ep.py:500
assert torch.equal(ref_t, t), f'{ref_t=}, {t=}'

# test_ep.py:448-462
for i in range(num_local_experts):
    ref_count = (ref_recv_topk_idx == i).sum().item()
    aligned_ref_count = align(ref_count, expert_alignment)
    assert ref_count == cumulative_local_expert_recv_stats[i].item()
    assert aligned_ref_count == handle.num_recv_tokens_per_expert_list[i]

# test_ep.py:424-428
for expert_idx in range(num_local_experts):
    start = expanded_handle.psum_num_recv_tokens_per_expert[expert_idx].item()
    end = align(start, expert_alignment)
    assert (cached_expanded_recv_x_bf16[start:end] == 0).all()
    assert (cached_expanded_recv_topk_weights[start:end] == 0).all()

# test_ep.py:388-397
if args.deterministic:
    recv_x_twice, ... = launch(buffer, 'dispatch', ...)
    assert torch.equal(recv_x_bf16, recv_x_twice_bf16)
    assert torch.equal(recv_topk_idx, recv_topk_idx_twice[:num_recv_tokens])

# test_ep.py:253-255
num_bytes_per_dispatch_token = safe_div(count_bytes(recv_x, recv_topk_idx, recv_topk_weights), recv_topk_idx.size(0))
num_scaleup_bytes = num_bytes_per_dispatch_token * num_scaleup_recv_tokens
num_scaleout_bytes = num_bytes_per_dispatch_token * num_scaleout_send_tokens

# test_ep.py:259-263
dist_print(f'   * EP: {buffer.rank_idx:3}/{buffer.num_ranks} | '
        f'dispatch: '
        f'{num_scaleout_bytes / t / 1e9:.0f} GB/s (SO), '
        f'{num_scaleup_bytes / t / 1e9:.0f} GB/s (SU), {t * 1e6:.3f} us, {num_scaleup_bytes:.0f} bytes | '
        f'copy: {2 * num_recv_tokens * num_bytes_per_dispatch_token / copy_t / 1e9:.0f} GB/s, {copy_t * 1e6:.3f} us')

# test_ep.py:240-245
dst_scaleout_rank_idx = topk_idx // (num_experts // num_scaleout_ranks)
num_scaleout_send_tokens = 0
for i in range(num_scaleout_ranks if num_scaleout_ranks > 1 else 0):
    if args.ignore_local_traffic and i == dist.get_rank() // num_scaleup_ranks:
        continue
    num_scaleout_send_tokens += (dst_scaleout_rank_idx == i).any(dim=1).sum().item()

# semantic.py:9-27
def weak_lru(maxsize: Optional[int] = 128, typed: bool = False):
    def wrapper(func):
        @functools.lru_cache(maxsize, typed)
        def _func(_self, *args, **kwargs):
            return func(_self(), *args, **kwargs)
        @functools.wraps(func)
        def inner(self, *args, **kwargs):
            return _func(weakref.ref(self), *args, **kwargs)
        return inner
    return wrapper

# find_pkgs.py:8-54
def find_pkg_root(name: str, lib_name: Optional[str] = None, optional: bool = False):
    # 1. 检查环境变量 EP_{NAME}_ROOT_DIR 和 {NAME}_DIR
    # 2. 遍历 importlib.metadata.distributions()
    # 3. 匹配 nvidia-{name} 包名
    # 4. 查找 lib_name 文件确定 root

# 基本正确性测试 (8 卡)
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --num-tokens 4096 --hidden 7168 --num-experts 256 --num-topk 6

# 跳过正确性验证，仅做性能测试
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --skip-check

# 压力测试 (循环多种子)
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --do-pressure-test

# 非均衡流量测试
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --unbalanced-ratio 2.0

# Deterministic 模式
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --deterministic

# 导出 profiling trace
torchrun --nproc_per_node=8 tests/elastic/test_ep.py --dump-profile-traces ./traces