"""EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py. CPU + gloo, no Neuron device. Run: python test_ep.py --num-processes {1,2,4} """ 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 # DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank) # with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo. DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', '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 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=DEVICE) 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=DEVICE) 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=DEVICE) 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=DEVICE) if num_bias == 1 else None if num_bias == 2: bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) for _ in range(2)) stats = torch.zeros((num_local_experts,), dtype=torch.int32, device=DEVICE) dispatch_args = dict( x=x_in, topk_idx=topk_idx, topk_weights=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) recv_x, recv_topk_idx, recv_weights, handle, _ = launch( buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=stats, **dispatch_args) assert recv_topk_idx is None assert recv_weights is not None assert handle.num_sms == 7 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) 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=DEVICE, 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=DEVICE, 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=DEVICE) input_for_combine[metadata[:, :-2].long().flatten()] = local_y.view(-1, hidden) input_for_combine = input_for_combine[:-1].contiguous() combined_x, combined_weights, _ = launch( buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle, topk_weights=recv_weights, bias=bias, num_sms=7, async_with_compute_stream=True) ref_y = generate_pre_combine_data(rank * cap + torch.arange(num_tokens, device=DEVICE), 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() assert torch.equal(combined_x, ref_combined_x), 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=input_for_combine, handle=handle, bias=bias, num_sms=7, async_with_compute_stream=True) assert torch.equal(combined_x2, ref_combined_x) assert combined_weights2 is None # --- Cached dispatch replay --- cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=handle) del cached_args['topk_idx'] cached_tensors = (handle.psum_num_recv_tokens_per_rank, handle.psum_num_recv_tokens_per_expert, handle.num_unaligned_recv_tokens_per_expert, handle.recv_src_metadata) cached_copies = [t.clone() for t in cached_tensors] padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=DEVICE) padding_mask[valid_slots] = False for send_weights in (True, False): updated_x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) * (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) cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch( buffer, 'dispatch', defer_epilogue, **dict(cached_args, x=updated_in, topk_weights=updated_weights if send_weights else None)) assert cached_handle is handle and cached_topk_idx is None for tensor, saved in zip(cached_tensors, cached_copies): assert torch.equal(tensor, saved) 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 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 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=DEVICE), topk_idx, updated_weights, cap, num_experts) return ref_w def worker(local_rank: int, num_ranks: int, args: argparse.Namespace): os.environ.setdefault('MASTER_ADDR', '127.0.0.1') os.environ.setdefault('MASTER_PORT', str(args.port)) if DEVICE == 'neuron': # Must be set before the Neuron runtime initializes (first device op) in this process. os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank) dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks) group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD rank = group.rank() 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})', 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)