DeepEP Elastic Test 源码深度分析

test_ep.py (609 行)
├── 导入区 (L1-19)
├── enumerate_ep_modes() (L22-31)     ← 参数空间笛卡尔积生成器
├── launch() (L34-41)                 ← 统一调用 dispatch/combine 的封装器
├── fold_expanded() (L44-55)          ← expand 模式结果折叠验证
├── test_dispatch_combine() (L59-516)  ← 核心测试逻辑
├── test_loop() (L20-561)             ← 分布式入口 + 压力测试
└── __main__ (L564-609)               ← argparse + multiprocessing.spawn

# L522: 通过 envs.init_dist() 初始化 NCCL 进程组
rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks, seed=args.seed)

# L523-534: construct_elastic_buffer 闭包
def construct_elastic_buffer():
    return deep_ep.ElasticBuffer(group,
                                 num_max_tokens_per_rank=args.num_tokens,
                                 hidden=args.hidden,
                                 deterministic=args.deterministic,
                                 allow_hybrid_mode=args.allow_hybrid_mode,
                                 allow_multiple_reduction=args.allow_multiple_reduction,
                                 prefer_overlap_with_compute=bool(args.prefer_overlap_with_compute),
                                 sl_idx=args.sl_idx,
                                 num_allocated_qps=max(args.num_allocated_qps, args.num_qps),
                                 explicitly_destroy=True,
                                 num_gpu_timeout_secs=args.num_gpu_timeout_secs,
                                 num_cpu_timeout_secs=args.num_cpu_timeout_secs)

# L608-609: 使用 torch.multiprocessing.spawn 启动多进程
num_processes = args.num_processes
torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes)

def enumerate_ep_modes():
    for do_handle_copy in (1, 0):                    # 是否拷贝 topk_idx
        for expert_alignment in (128, 1):            # expert 对齐粒度
            for use_fp8_dispatch in (1, 0):          # FP8 vs BF16 dispatch
                for num_bias in (0, 1, 2):           # bias 数量
                    for with_previous_event in (0, 1):  # 是否传入 previous_event
                        for async_with_compute_stream in (0, 1):  # 异步 compute stream
                            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)

def launch(buffer: deep_ep.ElasticBuffer, name: str,
           with_previous_event: int, async_with_compute_stream: int,
           params: dict):
    if with_previous_event:
        params.update(previous_event=buffer.capture())  # 捕获当前 CUDA event
    values = getattr(buffer, name)(**params)  # 动态调用 dispatch 或 combine
    values[-1].current_stream_wait() if async_with_compute_stream else ()
    return values

def fold_expanded(expanded, indices, valid_mask):
    if not isinstance(expanded, torch.Tensor):
        return tuple(fold_expanded(t, indices, valid_mask) for t in expanded)

    gathered = expanded[indices]  # 按 slot 索引收集
    first_valid_idx = valid_mask.to(torch.int).argmax(dim=1)  # 每行第一个有效位置
    folded = gathered[torch.arange(gathered.shape[0], device='cuda'), first_valid_idx]
    result = (gathered == folded.unsqueeze(1)).all(dim=-1)  # 验证所有有效位相同
    result = result | (~valid_mask)
    assert result.all()
    return folded

flowchart TD
    A[构造 expert 选择: scores → topk] --> B[枚举参数组合 enumerate_ep_modes]
    B --> C[生成输入数据 x, bias]
    C --> D[Reference 计算: ref_dispatch + ref_combine]
    D --> E[DeepEP dispatch: 4 种模式]
    E --> F[构造 combine 输入数据]
    F --> G[DeepEP combine: normal + reduced]
    G --> H[性能基准测试 bench_kineto]
    H --> I[正确性验证: bitwise compare]
    I --> J{还有更多组合?}
    J -->|Yes| C
    J -->|No| K[结束]

# L61-66: 获取逻辑域大小和计算参数
num_scaleout_ranks, num_scaleup_ranks = buffer.get_logical_domain_size()
num_max_tokens_per_rank, num_tokens, hidden = args.num_tokens, max(1, args.num_tokens - dist.get_rank()), args.hidden
num_topk, num_experts = args.num_topk, args.num_experts
num_local_experts = num_experts // buffer.num_ranks
num_sms = buffer.get_theoretical_num_sms(num_experts, num_topk) if args.num_sms == 0 else args.num_sms
num_qps = buffer.get_theoretical_num_qps(num_sms) if args.num_qps == 0 else args.num_qps

# L75-76: 生成不平衡的 expert 选择
scores = get_unbalanced_scores(num_tokens, num_experts, buffer.num_ranks, num_topk, 
                               args.unbalanced_ratio, args.precise_unbalanced_ratio)
topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
topk_idx = topk_idx.to(deep_ep.topk_idx_t)

# L78-81: 可选的 masked ratio（模拟 -1 索引）
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)

# L109-111: NCCL reference dispatch
ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, \
    ref_recv_src_token_idx, ref_num_recv_tokens_per_rank = \
    ref_dispatch(x, topk_idx, topk_weights, num_max_tokens_per_rank, num_experts)

# L125-140: Reference combine（两种 recipe）
ref_y = generate_pre_combine_data(...)
ref_reduced_combined_y = ref_combine(ref_y, topk_idx, ..., *reduced_combine_recipe)
ref_combined_y = ref_combine(ref_y, topk_idx, ..., *combine_recipe)

# L144-153: 标准 dispatch
recv_x, recv_topk_idx, recv_topk_weights, handle, dispatch_event = \
    launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, dispatch_args)

# L157-159: Expanding dispatch（do_expand=True）
expanded_recv_x, expanded_recv_topk_idx, expanded_recv_topk_weights, expanded_handle, expanded_dispatch_event = \
    launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, expanded_dispatch_args)

# L163-170: Cached dispatch（复用 handle）
cached_recv_x, cached_recv_topk_idx, cached_recv_topk_weights, cached_handle, cached_dispatch_event = \
    launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, cached_dispatch_args)

# L173-177: Cached expanding dispatch + zero padding
cached_expanded_recv_x, _, cached_expanded_recv_topk_weights, _, _ = \
    launch(buffer, 'dispatch', with_previous_event, async_with_compute_stream, cached_expanded_dispatch_args)

# L186-195: 标准 combine 输入
src_token_global_idx = handle.recv_src_metadata[:num_recv_tokens, 0]
local_y = generate_pre_combine_data(src_token_global_idx, num_max_tokens_per_rank, num_topk, hidden)
local_y[recv_topk_idx[:num_recv_tokens] == -1] = 0
local_reduced_y = ordered_accumulate(local_y)
input_for_combine = torch.empty_like(recv_x_bf16, dtype=torch.bfloat16, device='cuda')
input_for_combine[:num_recv_tokens] = local_reduced_y

# L197-206: Expand combine 输入（按 slot 索引 scatter）
input_for_expand_combine = torch.empty((expanded_recv_x_bf16.shape[0] + 1, hidden), ...)
input_for_expand_combine[expanded_handle.recv_src_metadata[:num_recv_tokens, 2:].flatten()] = local_y_expand.view(-1, hidden)
input_for_expand_combine = input_for_expand_combine[:-1, ...]  # 移除额外行

# L209-217: 标准 combine
combined_x, combined_topk_weights, combine_event = \
    launch(buffer, 'combine', with_previous_event, async_with_compute_stream, combine_args)

# L219-231: Reduced combine（expand 模式）
reduced_combined_x, reduced_combined_topk_weights, reduced_combine_event = \
    launch(buffer, 'combine', with_previous_event, async_with_compute_stream, reduced_combine_args)

dispatch_args = dict(
    x=x,                                    # 输入 token [num_tokens, hidden]
    topk_idx=topk_idx,                      # top-k expert 索引 [num_tokens, num_topk]
    topk_weights=topk_weights,              # top-k 权重 [num_tokens, num_topk]
    num_sms=num_sms,                        # 使用的 SM 数量
    num_qps=num_qps,                        # 使用的 QP 数量
    num_max_tokens_per_rank=num_max_tokens_per_rank,
    num_experts=num_experts,
    expert_alignment=expert_alignment,      # expert 对齐（128 或 1）
    async_with_compute_stream=async_with_compute_stream,
    allocate_on_comm_stream=allocate_on_comm_stream,
    do_handle_copy=do_handle_copy,          # 是否拷贝 topk_idx
    do_cpu_sync=args.do_cpu_sync,           # 是否 CPU 同步
)

combine_args = dict(
    x=input_for_combine,                    # combine 输入
    topk_weights=recv_topk_weights,         # dispatch 返回的 topk_weights
    bias=bias,                              # 可选 bias
    handle=handle,                          # dispatch 返回的 handle
    num_sms=num_sms,
    num_qps=num_qps,
    async_with_compute_stream=async_with_compute_stream,
    allocate_on_comm_stream=allocate_on_comm_stream,
)

# 仅需 x 和 handle，跳过 layout 重计算
cached_dispatch_args = dict(
    x=x,
    num_sms=num_sms, num_qps=num_qps,
    async_with_compute_stream=async_with_compute_stream,
    allocate_on_comm_stream=allocate_on_comm_stream,
    handle=handle,  # 复用之前的 handle
)

# L364: do_handle_copy=1 时 topk_idx 应该被拷贝（data_ptr 不同）
assert (topk_idx.data_ptr() != handle.topk_idx.data_ptr()) == do_handle_copy
# L365: cached handle 也应拷贝
assert (topk_idx.data_ptr() != cached_handle.topk_idx.data_ptr()) == do_handle_copy
# L366: 两次 dispatch 返回的 handle 共享同一 topk_idx
assert handle.topk_idx.data_ptr() == cached_handle.topk_idx.data_ptr()

if args.deterministic:
    # 再次 dispatch，结果应 bitwise 相同
    recv_x_twice, recv_topk_idx_twice, recv_topk_weights_twice, handle_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])
    assert torch.equal(recv_topk_weights, recv_topk_weights_twice[:num_recv_tokens])

# L448-453: 验证 cumulative_local_expert_recv_stats
for i in range(num_local_experts if args.do_cpu_sync else 0):
    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]

# L454-462: 验证 psum 计数
for i in range(num_local_experts):
    ref_count = (ref_recv_topk_idx == i).sum().item()
    count = psum_num_recv_tokens_per_expert_list[i + 1] - psum_num_recv_tokens_per_expert_list[i]
    expanded_count = (expanded_psum_num_recv_tokens_per_expert_list[i + 1] -
                      align(expanded_psum_num_recv_tokens_per_expert_list[i], expert_alignment))
    assert align(ref_count, expert_alignment) == count
    assert ref_count == expanded_count

# L476-500: 逐 rank 验证数据 bitwise 一致
for i in range(buffer.num_ranks):
    rank_start_idx = sum(ref_num_recv_tokens_per_rank[:i])
    rank_end_idx = rank_start_idx + ref_num_recv_tokens_per_rank[i]
    sorted_metadata = torch.sort(check_handle.recv_src_metadata[:, 0])
    sorted_indices = sorted_metadata.indices[rank_start_idx:rank_end_idx]

    # Data should be bitwise identical
    for ref_t, t, do_mask in check_list:
        ref_t = ref_t[rank_start_idx:rank_end_idx]
        t = t[sorted_indices]
        if do_mask:
            ref_t = ref_t.masked_fill(ref_mask, 0)
            t = t.masked_fill(ref_mask, 0)
        assert torch.equal(ref_t, t)

# L503: 标准 combine 结果 bitwise 一致
assert torch.equal(combined_x, ref_combined_y), \
    f'Diff: {calc_diff(combined_x, ref_combined_y)}'

# L505: Reduced combine 结果 bitwise 一致
assert torch.equal(reduced_combined_x, ref_reduced_combined_y), \
    f'Diff: {calc_diff(reduced_combined_x, ref_reduced_combined_y)}'

# L507: topk_weights 应还原
assert torch.equal(combined_topk_weights, topk_weights)

# 验证 cached expand 模式的 zero padding 正确性
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()

# L253-255: dispatch 字节数计算
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

# L256-263: dispatch 性能测试
t, copy_t = bench_kineto(lambda: buffer.dispatch(**dispatch_args),
                         kernel_names=('dispatch_impl', 'dispatch_copy_epilogue_impl'),
                         barrier_comm_profiling=True, barrier=buffer.barrier, 
                         trace_path=get_trace_path('dispatch'))

# L259-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')

flowchart LR
    A[test_ep.py] --> B[deep_ep.ElasticBuffer]
    A --> C[deep_ep.utils.math]
    A --> D[deep_ep.utils.gate]
    A --> E[deep_ep.utils.envs]
    A --> F[deep_ep.utils.refs]
    A --> G[deep_ep.utils.testing]

    C --> C1[per_token_cast_to_fp8]
    C --> C2[per_token_cast_back]
    C --> C3[count_bytes]
    C --> C4[calc_diff]
    C --> C5[align]

    D --> D1[get_unbalanced_scores]

    E --> E1[init_dist]
    E --> E2[init_seed]
    E --> E3[dist_print]

    F --> F1[ref_dispatch]
    F --> F2[ref_combine]
    F --> F3[generate_pre_combine_data]
    F --> F4[ordered_accumulate]

    G --> G1[bench_kineto]

# 确定性公式，无需通信即可生成相同数据
y[j,k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k+1) + sin(seed))

# L114-124: combine recipe 选择
if args.allow_multiple_reduction:
    if args.allow_hybrid_mode and num_scaleout_ranks > 1:
        reduced_combine_recipe = (True, True)   # Hybrid: local + scaleup reduction
        combine_recipe = (True, True)
    else:
        reduced_combine_recipe = (True, False)  # Non-hybrid: local only
        combine_recipe = (True, False)
else:
    reduced_combine_recipe = (False, False)     # No reduction
    combine_recipe = (True, False)

# L40: async_with_compute_stream 控制同步行为
values[-1].current_stream_wait() if async_with_compute_stream else ()

# L548-557: 压力测试循环
for seed in range(int(1e9) if args.do_pressure_test else 0):
    if not args.reuse_elastic_buffer:
        buffer.destroy()
        buffer = construct_elastic_buffer()
    init_seed(seed)
    test_dispatch_combine(buffer, args)

sequenceDiagram
    participant Main as __main__
    participant Loop as test_loop
    participant Init as init_dist
    participant Buf as ElasticBuffer
    participant Test as test_dispatch_combine
    participant Ref as ref_dispatch/combine
    participant DEP as DeepEP kernels
    participant Bench as bench_kineto

    Main->>Loop: torch.multiprocessing.spawn
    Loop->>Init: init_dist(local_rank, ...)
    Init-->>Loop: (rank, world_size, group)
    Loop->>Buf: construct_elastic_buffer()
    Loop->>Test: test_dispatch_combine(buffer, args)

    loop 192~288 参数组合
        Test->>Test: get_unbalanced_scores()
        Test->>Ref: ref_dispatch(x, topk_idx, ...)
        Ref-->>Test: (ref_recv_x, ref_recv_topk_idx, ...)
        Test->>Ref: ref_combine(ref_y, ...)
        Ref-->>Test: ref_combined_y

        Test->>DEP: buffer.dispatch(x, topk_idx, ...)
        DEP-->>Test: (recv_x, recv_topk_idx, ..., handle)
        Test->>DEP: buffer.dispatch(do_expand=True)
        DEP-->>Test: expanded_recv_x, expanded_handle
        Test->>DEP: buffer.dispatch(handle=handle)  // cached
        DEP-->>Test: cached_recv_x, cached_handle
        Test->>DEP: buffer.dispatch(do_zero_padding=True)
        DEP-->>Test: cached_expanded_recv_x

        Test->>DEP: buffer.combine(input_for_combine, handle=handle)
        DEP-->>Test: combined_x, combined_topk_weights
        Test->>DEP: buffer.combine(input_for_expand_combine, handle=expanded_handle)
        DEP-->>Test: reduced_combined_x

        Test->>Bench: bench_kineto(dispatch)
        Bench-->>Test: (t, copy_t)
        Test->>Bench: bench_kineto(combine)
        Bench-->>Test: (t, copy_t)

        Test->>Test: assert torch.equal(combined_x, ref_combined_y)
        Test->>Test: assert bitwise checks...
    end

    Loop->>Buf: buffer.destroy()
    Loop->>Init: dist.destroy_process_group()