DeepEP Elastic AGRS 与 PP 测试源码分析

# L12-17
def all_gather_ref(shape: tuple, rank_idx: int, num_ranks: int, round_idx: int = 0):
    ref_list = []
    for i in range(num_ranks):
        torch.manual_seed(42 + round_idx * 43 + i)  # 确定性种子
        ref_list.append(torch.randn(shape, dtype=torch.bfloat16, device='cuda'))
    return ref_list[rank_idx], torch.stack(ref_list, dim=0)

# L20-54
def generate_stress_ops(num_ops, num_max_inflight_agrs, shape, rank_idx, num_ranks):
    tensors, refs = zip(*(all_gather_ref(...) for i in range(num_ops)))
    unprocessed = random.sample(range(num_ops), num_ops)
    inflight, ops = [], [('create_session', (-1,))]
    limit = num_max_inflight_agrs
    while unprocessed or inflight:
        # 随机选择：ag / fetch / destroy
        op = random.choice(choices)
        if op == 'ag':
            b = tuple(unprocessed[-random.randint(1, max_g):])  # 随机批量
            ...
        elif op == 'fetch':
            ops.append(('fetch', inflight.pop(random.randrange(len(inflight)))))
        else:
            ops.extend([('destroy_session', (-1,)), ('create_session', (-1,))])

# L57-84
def do_all_gather(buffer, is_inplace, is_batched, tensors, start_event=None):
    if is_inplace:
        ag_tensors = buffer.agrs_get_inplace_tensor(tuple(t.shape for t in tensors), torch.bfloat16)
        for x, y in zip(ag_tensors, tensors, strict=True):
            x.copy_(y)  # 拷贝到弹性缓冲区
    else:
        ag_tensors = tensors

    if start_event is not None:
        torch.zeros(int(256e6 // 4), dtype=torch.int, device='cuda')  # flush L2 cache
        start_event.record()

    if is_batched:
        *out_tensors, handle = buffer.all_gather(ag_tensors)
        return out_tensors, [handle]
    else:
        # 逐个调用，每个返回独立 handle
        out_tensors, handles = [], []
        for t in ag_tensors:
            out_tensor, handle = buffer.all_gather(t)
            ...

# L93-110: 配置阶段
shape = (32, 64, 2048)  # 固定形状
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_num_max_session_bytes(
    group, [shape for _ in range(num_max_inflight_agrs)], torch.bfloat16)
num_max_session_bytes = deep_ep.ElasticBuffer.get_agrs_buffer_size_hint(
    group, num_max_session_bytes)
buffer = deep_ep.ElasticBuffer(group, explicitly_destroy=True, num_bytes=num_max_session_bytes)
buffer.agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)

# L136-138
for i in range(num_ops):
    assert results[i] is not None and torch.equal(results[i], refs[i]), \
        f'Rank {rank_idx}: stress mismatch at seed={seed}, op={i}'

# L156-187
for num_bytes in (2 ** p for p in range(20, 27)):  # 1MB 到 64MB
    shape = (num_bytes // 2, )  # bfloat16 = 2 bytes
    for is_inplace in (False, True):
        for is_batched in (False, True):
            # 50 次迭代，丢弃首次（warmup）
            times = np.array([s.elapsed_time(e) / 1e3 ...])[1:]
            bandwidth_info = f'{num_bytes * num_ranks * num_max_inflight_agrs / avg_t / 1e9:.3f} GB/s'

# L13-37
def generate_stress_ops(rank_idx, num_ranks, num_sends, shape):
    send_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d}
    recv_times = {(s, d): [] for s in range(num_ranks) for d in range(num_ranks) if s != d}

    for _ in range(num_sends):
        src_rank_idx = random.randint(0, num_ranks - 1)
        dst_rank_idx = (src_rank_idx + (1 if random.randint(0, 1) else -1)) % num_ranks  # 仅相邻 rank
        st = random.randint(0, 10 ** 8)
        rt = st + random.randint(1, 3 * 10 ** 6)  # recv 在 send 之后
        send_times[(src_rank_idx, dst_rank_idx)].append(st)
        recv_times[(src_rank_idx, dst_rank_idx)].append(rt)

# L44-59: 配置阶段
shape = (args.num_tokens, args.hidden)  # 默认 (4096, 7168)
num_max_tensor_bytes = math.prod(shape) * 2  # bfloat16
buffer = deep_ep.ElasticBuffer(
    group, explicitly_destroy=True, allow_hybrid_mode=False,  # 禁用 hybrid 模式
    num_bytes=deep_ep.ElasticBuffer.get_pp_buffer_size_hint(
        num_max_tensor_bytes, num_max_inflight_tensors))
buffer.pp_set_config(num_max_tensor_bytes, num_max_inflight_tensors)

# L69-80
prev = 0
for j, (op, timestamp, peer, _, tensor) in enumerate(ops):
    if op == 'send':
        buffer.pp_send(tensor, peer)
    else:
        result = torch.empty_like(tensor)
        buffer.pp_recv(result, peer)
        assert torch.equal(result, tensor), \  # 逐张量精确匹配
            f'Rank {rank_idx}: mismatch at op {j}'
    if timestamp > prev:
        torch.cuda._sleep(int((timestamp - prev) / 10 ** 8 * args.num_sleep_cycles))
    prev = timestamp

# L86-116
num_approx_rdma_cycles = int(num_max_tensor_bytes * 2 / get_rdma_gbs() * 1.5)

for hide_rdma_latency in (True, False):
    for num_concurrent in (1, 2, 3):
        def loop(_hide_rdma_latency=hide_rdma_latency):
            torch.zeros((131072, 32768), dtype=torch.int, device='cuda')  # L2 flush
            for t in send_tensors:
                buffer.pp_send(t, (rank_idx + 1) % num_ranks)
            if _hide_rdma_latency:
                torch.cuda._sleep(num_approx_rdma_cycles * num_concurrent)
            for t in recv_tensors:
                buffer.pp_recv(t, (rank_idx - 1) % num_ranks)

        send_t, recv_t = bench_kineto(
            loop, kernel_names=('send_impl', 'recv_impl'),
            barrier_comm_profiling=True, barrier=buffer.barrier,  # 使用 buffer 的 barrier
            trace_path=get_trace_path(...))

ElasticBuffer(group, explicitly_destroy=True, num_bytes=...)
  ├── .agrs_set_config(num_max_session_bytes, num_max_inflight_agrs)
  ├── .create_agrs_session() / .destroy_agrs_session()
  ├── .agrs_new_session()  [context manager]
  ├── .agrs_get_inplace_tensor(shapes, dtype)  [inplace 模式]
  └── .all_gather(tensors) → (gathered_tensors, handle)

ElasticBuffer(group, explicitly_destroy=True, allow_hybrid_mode=False, num_bytes=...)
  ├── .pp_set_config(num_max_tensor_bytes, num_max_inflight_tensors)
  ├── .pp_send(tensor, dst_rank_idx)
  ├── .pp_recv(tensor, src_rank_idx)
  └── .barrier  [用于 profiling 同步]

# L136-138
assert results[i] is not None and torch.equal(results[i], refs[i])

# L76-77
assert torch.equal(result, tensor)

# AGRS: L71
torch.zeros(int(256e6 // 4), dtype=torch.int, device='cuda')  # 256MB

# PP: L98
torch.zeros((131072, 32768), dtype=torch.int, device='cuda')  # 131072*32768*4 = 16GB (!)

# AGRS: L165-177
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):
    with buffer.agrs_new_session():
        _, wait_handles = do_all_gather(..., start_event=start_events[i])
        for h in wait_handles:
            h()
    end_events[i].record()
times = np.array([s.elapsed_time(e) / 1e3 for s, e in zip(start_events, end_events)])[1:]  # 丢弃 warmup

# PP: L78-80
if timestamp > prev:
    torch.cuda._sleep(int((timestamp - prev) / 10 ** 8 * args.num_sleep_cycles))
prev = timestamp

DeepEP Core (Production)
├── Dispatch / Combine (MoE expert parallelism)
├── Intra-node (NVLink)
└── Inter-node (RDMA)

Elastic Buffer (Experimental)
├── AGRS (All-Gather / Reduce-Scatter)
├── PP (Pipeline Parallelism send/recv)
└── Hybrid Mode (NCCL fallback)