"""EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py. Default CPU + gloo. DEEP_EP_TEST_DEVICE=neuron: buffers on NeuronCores with a neuron process group; inputs and references stay on the host (see dev.py). Run: python test_ep.py --num-processes {1,2,4,8} """ import argparse import os import sys import torch import torch.distributed as dist import torch.multiprocessing as mp sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor")) import deep_ep # noqa: E402 from deep_ep.utils.math import align, per_token_cast_to_fp8 # noqa: E402 from refs import combine as ref_combine # noqa: E402 from refs import dispatch as ref_dispatch # noqa: E402 from refs import generate_pre_combine_data # noqa: E402 from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402 # Inputs, references and checks live on the host; only the buffer calls run on DEVICE. HOST = 'cpu' def enumerate_modes(): for expert_alignment in (128, 1): for use_fp8 in (True, False): for num_bias in (0, 1, 2): for col_major in ((False, True) if use_fp8 else (False,)): yield expert_alignment, num_bias, use_fp8, col_major def launch(buffer, name, defer_epilogue, **params): result = getattr(buffer, name)(defer_epilogue=defer_epilogue, **params) if defer_epilogue: event = result return *event.current_stream_wait(release_handle=True), event result[-1].current_stream_wait(release_handle=True) return result class HostHandle: """Host copies of the tensors of an `EPHandle` (the device handle itself is what the buffer gets back).""" def __init__(self, handle): for name in ('psum_num_recv_tokens_per_rank', 'psum_num_recv_tokens_per_expert', 'num_unaligned_recv_tokens_per_expert', 'recv_src_metadata'): setattr(self, name, to_host(getattr(handle, name))) self.num_recv_tokens = handle.num_recv_tokens self.num_recv_tokens_per_expert_list = list(handle.num_recv_tokens_per_expert_list) self.num_sms = handle.num_sms def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): expert_alignment, num_bias, use_fp8, col_major = mode tag = f'mode={mode} defer={defer_epilogue} cpu_sync={do_cpu_sync}' cap = args.num_tokens num_tokens = max(1, args.num_tokens - rank) hidden, num_topk, num_experts = args.hidden, args.num_topk, args.num_experts num_local_experts = num_experts // num_ranks torch.manual_seed(1234 + rank) scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=HOST) topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) topk_idx = topk_idx.to(torch.int64) if args.masked_ratio > 0: rand_mask = torch.rand_like(topk_idx, dtype=torch.float) topk_idx = topk_idx.masked_fill(rand_mask < args.masked_ratio, -1) topk_weights = topk_weights.masked_fill(topk_idx < 0, 0) topk_idx, topk_weights = topk_idx.contiguous(), topk_weights.contiguous() x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) if use_fp8: x_data, sf = per_token_cast_to_fp8(x) sf = sf.view(torch.int16) num_sf_packs = sf.size(1) if col_major: sf_storage = torch.empty((num_tokens * num_sf_packs,), dtype=torch.int16, device=HOST) input_sf = torch.as_strided(sf_storage, (num_tokens, num_sf_packs), (1, num_tokens)) input_sf.copy_(sf) else: input_sf = sf x_in = (x_data, input_sf) else: x_in = x bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) if num_bias == 1 else None if num_bias == 2: bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) for _ in range(2)) d_stats = to_dev(torch.zeros((num_local_experts,), dtype=torch.int32, device=HOST)) d_topk_idx = to_dev(topk_idx) dispatch_args = dict( x=to_dev(x_in), topk_idx=d_topk_idx, topk_weights=to_dev(topk_weights), num_sms=7, num_max_tokens_per_rank=cap, num_experts=num_experts, expert_alignment=expert_alignment, async_with_compute_stream=True, do_handle_copy=False, do_cpu_sync=do_cpu_sync, do_expand=True, do_zero_padding=True, use_tma_aligned_col_major_sf=col_major) ref_x = (x_data.view(torch.uint8), input_sf) if use_fp8 else x (ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights, ref_recv_src_token_idx, ref_num_recv_tokens_per_rank) = ref_dispatch( ref_x, topk_idx, topk_weights, cap, num_experts) d_recv_x, recv_topk_idx, d_recv_weights, d_handle, _ = launch( buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=d_stats, **dispatch_args) assert recv_topk_idx is None assert d_recv_weights is not None assert d_handle.num_sms == 7 if use_fp8: # strides are checked on the device tensor itself sf_stride_on_device = d_recv_x[1].stride() recv_x, recv_weights, handle, stats = to_host(d_recv_x), to_host(d_recv_weights), HostHandle(d_handle), to_host(d_stats) recv_x_data, recv_sf = recv_x if use_fp8 else (recv_x, None) check_x = recv_x_data.view(torch.uint8) if use_fp8 else recv_x_data num_recv_tokens = handle.psum_num_recv_tokens_per_rank[-1].item() num_expanded_tokens = align(handle.psum_num_recv_tokens_per_expert[-1].item(), expert_alignment) expected_num_recv = num_recv_tokens if do_cpu_sync else num_ranks * cap assert handle.num_recv_tokens == expected_num_recv, (handle.num_recv_tokens, expected_num_recv) assert handle.recv_src_metadata.size(0) == expected_num_recv if use_fp8: expected_stride = (1, recv_x_data.size(0)) if col_major else (num_sf_packs, 1) assert recv_sf.stride() == expected_stride, (recv_sf.stride(), expected_stride) assert sf_stride_on_device == expected_stride, (sf_stride_on_device, expected_stride) ref_recv_x_data, ref_recv_sf = ref_recv_x if use_fp8 else (ref_recv_x, None) assert num_recv_tokens == ref_recv_x_data.size(0), (num_recv_tokens, ref_recv_x_data.size(0)) metadata = handle.recv_src_metadata[:num_recv_tokens] recv_src_token_idx = metadata[:, -2] recv_expanded_slots = metadata[:, :-2] ref_order = torch.argsort(ref_recv_src_token_idx) metadata_order = torch.argsort(recv_src_token_idx) assert torch.equal(recv_src_token_idx[metadata_order], ref_recv_src_token_idx[ref_order]) def reorder(src): out = torch.empty_like(src) out[metadata_order] = src[ref_order] return out ref_x_by_meta = reorder(ref_recv_x_data) ref_sf_by_meta = reorder(ref_recv_sf) if ref_recv_sf is not None else None ref_idx_by_meta = reorder(ref_recv_topk_idx) ref_w_by_meta = reorder(ref_recv_topk_weights) recv_valid_mask = recv_expanded_slots >= 0 assert torch.equal(recv_valid_mask, ref_idx_by_meta >= 0) valid_slots = recv_expanded_slots[recv_valid_mask].long() assert valid_slots.numel() == handle.num_unaligned_recv_tokens_per_expert.sum().item() if valid_slots.numel() > 0: assert valid_slots.min().item() >= 0 and valid_slots.max().item() < num_expanded_tokens assert torch.unique(valid_slots).numel() == valid_slots.numel() expected_x = ref_x_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask] assert torch.equal(check_x[valid_slots], expected_x) if ref_sf_by_meta is not None: expected_sf = ref_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask] assert torch.equal(recv_sf[valid_slots], expected_sf) assert torch.equal(recv_weights[valid_slots], ref_w_by_meta[recv_valid_mask]) # Master top-k slot column topk_range = torch.arange(num_topk, device=HOST, dtype=metadata.dtype) master = torch.where(recv_valid_mask, topk_range.view(1, num_topk), torch.full_like(recv_expanded_slots, -1)).max(dim=1).values assert torch.equal(metadata[:, -1], master) # Per-rank and per-expert counts psum_rank = [0] + handle.psum_num_recv_tokens_per_rank.tolist() for r in range(num_ranks): assert psum_rank[r + 1] - psum_rank[r] == ref_num_recv_tokens_per_rank[r].item() psum_expert = [0] + handle.psum_num_recv_tokens_per_expert.tolist() for e in range(num_local_experts): count = psum_expert[e + 1] - align(psum_expert[e], expert_alignment) ref_count = (ref_recv_topk_idx == e).sum().item() assert count == ref_count, (e, count, ref_count) assert stats[e].item() == ref_count assert handle.num_unaligned_recv_tokens_per_expert[e].item() == ref_count if do_cpu_sync: assert handle.num_recv_tokens_per_expert_list[e] == align(ref_count, expert_alignment) expert_start = align(psum_expert[e], expert_alignment) expert_end = psum_expert[e + 1] expert_slots = recv_expanded_slots[ref_idx_by_meta == e] assert expert_slots.numel() == ref_count if ref_count > 0: assert torch.equal(torch.sort(expert_slots).values, torch.arange(expert_start, expert_end, device=HOST, dtype=expert_slots.dtype)) padding_end = align(expert_end, expert_alignment) assert torch.count_nonzero(recv_x_data[expert_end:padding_end].view(torch.uint8)).item() == 0 if recv_sf is not None: assert torch.count_nonzero(recv_sf[expert_end:padding_end]).item() == 0 # --- Combine --- local_y = generate_pre_combine_data(metadata[:, -2], cap, num_topk, hidden) input_for_combine = torch.empty((recv_x_data.size(0) + 1, hidden), dtype=torch.bfloat16, device=HOST) input_for_combine[metadata[:, :-2].long().flatten()] = local_y.view(-1, hidden) input_for_combine = input_for_combine[:-1].contiguous() d_input_for_combine = to_dev(input_for_combine) combined_x, combined_weights, _ = launch( buffer, 'combine', defer_epilogue, x=d_input_for_combine, handle=d_handle, topk_weights=d_recv_weights, bias=to_dev(bias), num_sms=7, async_with_compute_stream=True) combined_x, combined_weights = to_host(combined_x), to_host(combined_weights) ref_y = generate_pre_combine_data(rank * cap + torch.arange(num_tokens, device=HOST), cap, num_topk, hidden) ref_y[topk_idx == -1] = 0 ref_combined_x = ref_combine(ref_y, topk_idx, num_experts, bias, reduce_in_local=True) max_diff = (combined_x.float() - ref_combined_x.float()).abs().max().item() check_combined(buffer, combined_x, ref_combined_x, ref_y, bias, f'{tag} {max_diff=}') assert torch.equal(combined_weights, topk_weights) if num_bias == 0 and not col_major: combined_x2, combined_weights2, _ = launch( buffer, 'combine', defer_epilogue, x=d_input_for_combine, handle=d_handle, bias=to_dev(bias), num_sms=7, async_with_compute_stream=True) combined_x2 = to_host(combined_x2) check_combined(buffer, combined_x2, ref_combined_x, ref_y, bias, f'{tag} (no weights)') assert combined_weights2 is None # --- Cached dispatch replay --- cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=d_handle) del cached_args['topk_idx'] cached_names = ('psum_num_recv_tokens_per_rank', 'psum_num_recv_tokens_per_expert', 'num_unaligned_recv_tokens_per_expert', 'recv_src_metadata') cached_copies = [getattr(handle, name).clone() for name in cached_names] padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=HOST) padding_mask[valid_slots] = False for send_weights in (True, False): updated_x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=HOST) * (8 if send_weights else 0.125) if use_fp8: updated_data, updated_sf = per_token_cast_to_fp8(updated_x) updated_sf = updated_sf.view(torch.int16) if col_major: updated_sf = updated_sf.t().contiguous().t() updated_in = (updated_data, updated_sf) else: updated_in = updated_x updated_weights = torch.randn_like(topk_weights).masked_fill(topk_idx < 0, 0) ref_updated = (updated_data.view(torch.uint8), updated_sf) if use_fp8 else updated_x ref_cached_x, _, _, _, _ = ref_dispatch(ref_updated, topk_idx, updated_weights, cap, num_experts) d_cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch( buffer, 'dispatch', defer_epilogue, **dict(cached_args, x=to_dev(updated_in), topk_weights=to_dev(updated_weights) if send_weights else None)) assert cached_handle is d_handle and cached_topk_idx is None for name, saved in zip(cached_names, cached_copies): assert torch.equal(to_host(getattr(d_handle, name)), saved), name if use_fp8: cached_sf_stride_on_device = d_cached_x[1].stride() cached_x, cached_w = to_host(d_cached_x), to_host(cached_w) cached_data, cached_sf = cached_x if use_fp8 else (cached_x, None) ref_cached_data, ref_cached_sf = ref_cached_x if use_fp8 else (ref_cached_x, None) cached_check = cached_data.view(torch.uint8) if use_fp8 else cached_data ref_cached_by_meta = reorder(ref_cached_data) assert torch.equal(cached_check[valid_slots], ref_cached_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]) assert torch.count_nonzero(cached_check[:num_expanded_tokens][padding_mask]).item() == 0 if cached_sf is not None: assert cached_sf.stride() == expected_stride assert cached_sf_stride_on_device == expected_stride ref_cached_sf_by_meta = reorder(ref_cached_sf) assert torch.equal(cached_sf[valid_slots], ref_cached_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]) if send_weights: assert torch.equal(cached_w[valid_slots], reorder(updated_weights_ref( updated_weights, topk_idx, cap, num_experts))[recv_valid_mask]) else: assert cached_w is None def check_combined(buffer, combined_x, ref_combined_x, ref_y, bias, tag): """`a2a` reproduces the reference reduction order exactly (bit-exact check). `allgather` sums the per-rank bf16 partials with an fp32 reduce-scatter (collective order) and adds the bias afterwards: `(b0 + b1) + (p_r0 + p_r1 + ...)` instead of `((b0 + b1) + p_m0) + p_m1 + ...` in master top-k slot order. Tokens whose selections all live on one rank stay bit-exact; the rest may differ by fp32 rounding, bounded by one bf16 ulp of the result plus the fp32 error of re-associating the terms. """ if buffer.algo == 'a2a': assert torch.equal(combined_x, ref_combined_x), tag return terms = ref_y.float().abs().sum(1) for b in (bias if isinstance(bias, tuple) else (bias,)): if b is not None: terms = terms + b.float().abs() out, ref = combined_x.float(), ref_combined_x.float() tol = torch.maximum(out.abs(), ref.abs()) * 2 ** -7 + terms * 2 ** -20 diff = (out - ref).abs() assert bool((diff <= tol).all()), f'{tag} max_excess={(diff - tol).max().item()}' CHECK_STATS[0] += int((diff > 0).sum()) CHECK_STATS[1] += diff.numel() CHECK_STATS = [0, 0] # (non-identical elements, compared elements) for the allgather combine def updated_weights_ref(updated_weights, topk_idx, cap, num_experts): _, _, ref_w, _, _ = ref_dispatch( torch.zeros((topk_idx.size(0), 8), dtype=torch.bfloat16, device=HOST), topk_idx, updated_weights, cap, num_experts) return ref_w def worker(local_rank: int, num_ranks: int, args: argparse.Namespace): group = init_dist(local_rank, num_ranks, args.port) rank = group.rank() assert rank == dist.get_rank() # host references use the default (gloo) group buffer = deep_ep.EPBuffer( group, num_max_tokens_per_rank=args.num_tokens, hidden=args.hidden, num_topk=args.num_topk, use_fp8_dispatch=True, allow_hybrid_mode=True, allow_multiple_reduction=True, explicitly_destroy=True) assert buffer.num_bytes > 0 assert deep_ep.get_physical_domain_size(group) == (1, num_ranks) for mode in enumerate_modes(): for defer_epilogue in (True, False): for do_cpu_sync in (True, False): check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync) if rank == 0: print(f' ok: expert_alignment={mode[0]} num_bias={mode[1]} fp8={mode[2]} col_major={mode[3]}', flush=True) buffer.destroy() if rank == 0: print(f'EP tests passed (world_size={num_ranks}, algo={buffer.algo}, ' f'combine non-identical elements on rank 0: {CHECK_STATS[0]}/{CHECK_STATS[1]})', flush=True) dist.destroy_process_group() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--num-processes', type=int, default=2) parser.add_argument('--num-tokens', type=int, default=17) parser.add_argument('--hidden', type=int, default=256) parser.add_argument('--num-topk', type=int, default=4) parser.add_argument('--num-experts', type=int, default=16) parser.add_argument('--masked-ratio', type=float, default=0.2) parser.add_argument('--port', type=int, default=29531) args = parser.parse_args() assert args.num_experts % args.num_processes == 0 if args.num_processes == 1: worker(0, 1, args) else: mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)