DeepInfer/tests/neuron/deep_ep/test_engram.py

"""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)