mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
302 lines
15 KiB
Python
302 lines
15 KiB
Python
"""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)
|