DeepEP 辅助工具测试源码分析：Barrier / Engram / Gate

# L19-26: 核心测试逻辑
def loop_barrier(num_tests=1000):
    for i in range(num_tests):
        buffer.barrier()

t = bench_kineto(lambda: loop_barrier(), 'barrier',
                 barrier_comm_profiling=True, barrier=buffer.barrier)
dist_print(f' > EP: {buffer.rank_idx:3}/{buffer.num_ranks:3}, '
           f'barrier time: {t * 1e6:.3f} us')

buffer.barrier()  (elastic.py L497)
  → runtime.barrier(use_comm_stream, with_cpu_sync, sequential)
    → launch_barrier(...)  (barrier.hpp L56)
      → barrier_impl<kIsScaleupNVLink, kNumSMs, kNumThreads,
                      kNumScaleoutRanks, kNumScaleupRanks,
                      kNumTimeoutCycles, kSequential>  (barrier.cuh L17)

// barrier.cuh L24-39: 串行 vs 并行
if constexpr (kSequential) {
    if constexpr (kNumScaleoutRanks > 1)
        comm::gpu_barrier<..., kKernelBarrierTag, false, false, false>(...);
    // scaleup barrier 需要 flush RDMA 请求
    comm::gpu_barrier<..., kKernelBarrierTag, true, true, false>(...);
} else {
    comm::gpu_barrier<..., kKernelBarrierTag, false, false, false>(...);
}

# L17-37: 配置与 buffer 创建
num_gpu_bytes, num_cpu_bytes = deep_ep.ElasticBuffer.get_engram_storage_size_hint(
    args.num_entries, args.hidden, args.num_tokens * args.num_entries_per_token, dtype)
buffer = deep_ep.ElasticBuffer(
    group, num_bytes=num_gpu_bytes + num_cpu_bytes, num_cpu_bytes=num_cpu_bytes,
    explicitly_destroy=True, num_allocated_qps=num_qps, ...)

# L40-53: 写入阶段
local_bf16 = torch.randn((args.num_entries, args.hidden), dtype=torch.bfloat16, device='cuda')
if args.use_fp8:
    local_storage, local_sf = per_token_cast_to_fp8(local_bf16)
buffer.engram_write(local_storage, sf=sf)

# L56-57: 生成随机索引
indices = torch.randint(0, num_ranks * args.num_entries,
                        (args.num_tokens, args.num_entries_per_token), device='cuda', dtype=torch.int)

# L60-71: 基于 all_gather 的参考实现
if not args.skip_check:
    global_storage = torch.empty((num_ranks * args.num_entries, args.hidden), dtype=dtype, device='cuda')
    dist.all_gather_into_tensor(global_storage, local_storage, group)
    ref_data = global_storage[indices.view(-1)].view(args.num_tokens, -1)

    for use_tma_aligned_col_major_sf in (False, True) if args.use_fp8 else (False,):
        data, fetched_sf = buffer.engram_fetch(indices,
                                               use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf)()
        assert torch.equal(ref_data, data), 'data mismatch'

# L79-96: issue + wait 端到端测量
def fetch_and_wait():
    hook = buffer.engram_fetch(indices, use_tma_aligned_col_major_sf=True)
    hook()

issue_t, wait_t = bench_kineto(
    fetch_and_wait,
    kernel_names=('engram_fetch_impl', 'engram_fetch_wait_impl'),
    barrier_comm_profiling=True, barrier=buffer.barrier,
    trace_path=...)
mpps = args.num_tokens * args.num_entries_per_token / (issue_t + wait_t) / 1e6

// L29-35: 每个 block 对应一个 QP，warp 级并行
const auto qp_idx = static_cast<int>(blockIdx.x);
const auto warp_idx = ptx::get_warp_idx();
const auto global_warp_idx = qp_idx * kNumWarps + warp_idx;
const auto gin = handle::NCCLGin(nccl_dev_comm, nccl_window, qp_idx, NCCL_GIN_RESOURCE_SHARING_CTA);

// L55-74: 每个 warp 协作 fetch 一个 token 的多个 entry
for (int i = global_warp_idx; i < num_tokens * kNumEntriesPerToken; i += kNumQPs * kNumWarps) {
    const auto global_idx = __ldg(indices + i);
    const auto owner_rank_idx = global_idx / kNumEntriesPerRank;
    const auto local_entry_idx = global_idx % kNumEntriesPerRank;
    const auto peer_idx = owner_rank_idx / kNumRanksPerRDMAPeer;
    gin.get<team_t, ncclCoopThread, ncclGin_SegmentMixed>(...);
}

# L7-53: 穷举参数空间测试
def test_unbalanced_scores():
    for num_tokens in [1, 4096]:
        for num_experts_per_rank in [1, 4, 8, 16]:
            for num_ranks in [2, 4, 8, 16, 64, 72]:
                num_experts = num_experts_per_rank * num_ranks
                for num_topk in [1, 2, 4, 6, 8, 9]:
                    for ratio in [1.0, 2.0, 4.0]:
                        for precise in [1, 0]:
                            # L18-25: 边界条件检查
                            lower_bound_per_token = max(1, ceil_div(num_topk, num_experts_per_rank))
                            upper_bound_per_token = min(min(num_topk, num_ranks), int((num_ranks - 1) / ratio) + 1)
                            if lower_bound_per_token > upper_bound_per_token:
                                continue

                            # L27-36: 逐 rank 生成并统计
                            for rank_idx in range(num_ranks):
                                scores = get_unbalanced_scores(num_tokens, num_experts, num_ranks, num_topk, ratio, precise)
                                _topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, ...)
                                topk_idx = topk_idx // num_experts_per_rank
                                ...  # 统计每个 rank 被选中次数

                            # L39-52: 验证 ratio 与 inequality
                            practical_ratio = total_rank_count[0].item() / max(total_rank_count[1:].min().item(), 1)
                            inequality = total_rank_count[1:].max().item() / max(total_rank_count[1:].min().item(), 1)
                            if precise:
                                assert abs(practical_ratio - ratio) / ratio < 0.1 and inequality < 1.02

scores = get_unbalanced_scores(num_tokens, num_experts, buffer.num_ranks, num_topk,
                               args.unbalanced_ratio, args.precise_unbalanced_ratio)

# 标准入口模式
@torch.inference_mode()
def test_loop(local_rank, num_ranks, args):
    rank, num_ranks, group = init_dist(local_rank, num_ranks)
    buffer = deep_ep.ElasticBuffer(group, ...)
    # ... 测试逻辑
    buffer.destroy()
    dist.destroy_process_group()

if __name__ == '__main__':
    torch.multiprocessing.spawn(test_loop, args=(num_processes, args), nprocs=num_processes)

t = bench_kineto(fn, kernel_names=...,
                 barrier_comm_profiling=True, barrier=buffer.barrier,
                 trace_path=...)

                    ┌─────────────────────┐
                    │   test_ep.py        │
                    │   (主 EP 测试)       │
                    └─────────┬───────────┘
                              │
          ┌───────────────────┼───────────────────┐
          │                   │                   │
          ▼                   ▼                   ▼
   ┌─────────────┐    ┌──────────────┐    ┌──────────────┐
   │ barrier     │    │ engram       │    │ gate         │
   │ (同步原语)   │    │ (RDMA 原语)   │    │ (负载生成)    │
   └──────┬──────┘    └──────┬───────┘    └──────┬───────┘
          │                  │                   │
          ▼                  ▼                   ▼
   ┌─────────────┐    ┌──────────────┐    ┌──────────────┐
   │ barrier_impl│    │engram_fetch_ │    │generate_rank_│
   │ (NCCL Gin   │    │impl / wait   │    │count /       │
   │  GPU barrier)│    │(RDMA get)    │    │topk_idx      │
   └─────────────┘    └──────────────┘    └──────────────┘

hook = buffer.engram_fetch(indices)  # 异步发起
# ... 可插入计算 ...
hook()  # 阻塞等待，返回 (data, sf)