DeepInfer/tests/neuron/deep_ep/test_engram.py

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