"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py (CPU + gloo).""" 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 ceil_div, per_token_cast_to_fp8 # 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 run(rank: int, num_ranks: int, group, args, use_fp8: bool): dtype = torch.float8_e4m3fn if use_fp8 else torch.bfloat16 num_entries_per_layer = [int(v) for v in args.num_entries_per_layer.split(',')] num_layers = len(num_entries_per_layer) num_sf_packs = 2 * ceil_div(args.hidden, 128) if use_fp8 else 0 num_bytes, num_rdma_bytes = deep_ep.EngramBuffer.get_storage_size_hint( group, num_entries_per_layer, args.hidden, args.num_tokens, args.num_entries_per_token, dtype, num_sf_packs) num_qps, qp_depth, restrict = deep_ep.EngramBuffer.get_theoretical_config( group, num_layers, args.num_tokens, args.num_entries_per_token) assert (num_qps, qp_depth, restrict) == (0, 0, False) buffer = deep_ep.EngramBuffer( group, num_gpu_bytes=num_bytes, num_rdma_storage_size=num_rdma_bytes, use_cpu_rdma_storage=True, num_allocated_qps=num_qps, qp_depth=qp_depth, restrict_rd_atomic=restrict, allow_hybrid_mode=True, explicitly_destroy=True) buffer.set_config(num_entries_per_layer, args.hidden, args.num_tokens, args.num_entries_per_token, dtype, num_sf_packs) torch.manual_seed(7 + rank) local_storages, sfs = [], ([] if use_fp8 else None) for num_entries in num_entries_per_layer: local_bf16 = torch.randn((num_entries, args.hidden), dtype=torch.bfloat16, device=DEVICE) if use_fp8: local_storage, local_sf = per_token_cast_to_fp8(local_bf16) local_sf = local_sf.view(torch.int16).contiguous() sf = torch.empty((num_ranks * num_entries, num_sf_packs), dtype=local_sf.dtype, device=DEVICE) if num_ranks == 1: sf.copy_(local_sf) else: # gloo has no int16 collective; gather the bytes instead (DeepInfer's # NgramEmbedding.gather_sf gathers int16 directly on the accelerator). dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8), group=group) sfs.append(sf) else: local_storage = local_bf16 local_storages.append(local_storage.contiguous()) buffer.write(local_storages, sfs=sfs) # Reference global tables global_bits = [] for local_storage in local_storages: bits = torch.empty((num_ranks * local_storage.shape[0], args.hidden * dtype.itemsize), dtype=torch.uint8, device=DEVICE) local_bits = local_storage.view(torch.uint8).view(local_storage.shape[0], -1).contiguous() if num_ranks == 1: bits.copy_(local_bits) else: dist.all_gather_into_tensor(bits, local_bits, group=group) global_bits.append(bits) def generate_indices(): return torch.stack([ torch.randint(0, num_ranks * num_entries, (args.num_tokens, args.num_entries_per_token), dtype=torch.int, device=DEVICE) for num_entries in num_entries_per_layer]) for col_major in ((False, True) if use_fp8 else (False,)): for layer_order in (range(num_layers), reversed(range(num_layers))): indices = generate_indices() ref_data = [global_bits[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1) for l in range(num_layers)] ref_sfs = [sfs[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1) for l in range(num_layers)] if use_fp8 else None hooks = buffer.fetch(indices, use_tma_aligned_col_major_sf=col_major) assert len(hooks) == num_layers for l in layer_order: fetched = hooks[l]() data, fetched_sf = fetched if use_fp8 else (fetched, None) assert data.shape == (args.num_tokens, args.num_entries_per_token * args.hidden) assert torch.equal(ref_data[l], data.reshape(args.num_tokens, -1).contiguous().view(torch.uint8)), \ f'data mismatch (layer {l}, fp8={use_fp8}, col_major={col_major})' if use_fp8: expected_stride = (1, args.num_tokens) if col_major else (fetched_sf.size(1), 1) assert fetched_sf.stride() == expected_stride, (fetched_sf.stride(), expected_stride) assert torch.equal(ref_sfs[l], fetched_sf), f'sf mismatch (layer {l})' buffer.destroy() def worker(local_rank: int, num_ranks: int, args): 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 for use_fp8 in (False, True): run(group.rank(), num_ranks, group, args, use_fp8) if group.rank() == 0: print(f' ok: fp8={use_fp8}', flush=True) if group.rank() == 0: print(f'Engram 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-entries-per-layer', type=str, default='37,64') parser.add_argument('--hidden', type=int, default=128) parser.add_argument('--num-tokens', type=int, default=11) parser.add_argument('--num-entries-per-token', type=int, default=3) parser.add_argument('--port', type=int, default=29541) args = parser.parse_args() if args.num_processes == 1: worker(0, 1, args) else: mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)