"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py. Default CPU + gloo; DEEP_EP_TEST_DEVICE=neuron runs the buffer on NeuronCores with a neuron process group while the tables, indices and reference lookups stay on the host (see dev.py). """ 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 from dev import init_dist, to_dev, to_host # noqa: E402 HOST = 'cpu' # inputs + references; the buffer runs on dev.DEVICE 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=HOST) 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=HOST) 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). # Host reference on the default (gloo) group. dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8)) sfs.append(sf) else: local_storage = local_bf16 local_storages.append(local_storage.contiguous()) buffer.write([to_dev(t) for t in local_storages], sfs=None if sfs is None else [to_dev(t) for t in 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=HOST) 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) # host, default (gloo) 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=HOST) 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(to_dev(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) if fetched_sf is not None: sf_stride_on_device = fetched_sf.stride() data, fetched_sf = to_host(data), to_host(fetched_sf) 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 sf_stride_on_device == expected_stride, (sf_stride_on_device, 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): group = init_dist(local_rank, num_ranks, args.port) assert group.rank() == dist.get_rank() 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)