From c1424354729d3da0ca0f2ffb7898e8afa9d69677 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 22 Oct 2026 04:01:28 +0800 Subject: [PATCH] Neuron deep_ep: MoE dispatch/combine via all-gather + reduce-scatter (issue #12) Uneven all_to_all on the neuron process group is unreliable at scale (EP=64: "Invalid NEFF" in prefill; test_ep at 64 ranks hung after 9/18 modes). EPBuffer now moves tokens without all_to_all: - dispatch: one all_gather of x / sf / weights / topk_idx packed as uint8 rows (row count = global max via a 0-dim MAX all_reduce, or the static num_max_tokens_per_rank with DEEPINFER_NEURON_EP_AG_ROWS=cap); each rank selects the rows routed to its local experts in the same (source rank, source token) order as before, so layout / metadata / cached dispatch are unchanged. - combine: per-token fp32 partials over local experts (bf16-rounded as before) + weights, one fp32 SUM reduce_scatter (all_reduce fallback with DEEPINFER_NEURON_EP_REDUCE=ar), chunked under DEEPINFER_NEURON_EP_CHUNK_BYTES (64 MiB). - The a2a path stays behind DEEPINFER_NEURON_EP_ALGO=a2a. test_ep passes on the neuron group at 2/4/8/16/32/64 ranks (combine bit-identical to the reference in practice); new test_collectives.py found that int32 all_gather with full-range values is wrong on a 16-rank neuron group for >= 1 MiB outputs (uint8 / fp32 fine) -- payloads move as bytes. Tiny model EP=8 on NeuronCores: 25/27 argmax vs the reference, same as EP=1. Co-Authored-By: Claude Opus 5.5 --- .../neuron_ops/vendor/deep_ep/__init__.py | 3 +- .../neuron_ops/vendor/deep_ep/buffers/ep.py | 483 ++++++++++++++---- deepinfer/neuron_ops/vendor/deep_ep/comm.py | 44 ++ tests/neuron/deep_ep/test_collectives.py | 86 ++++ tests/neuron/deep_ep/test_ep.py | 33 +- 5 files changed, 551 insertions(+), 98 deletions(-) create mode 100644 tests/neuron/deep_ep/test_collectives.py diff --git a/deepinfer/neuron_ops/vendor/deep_ep/__init__.py b/deepinfer/neuron_ops/vendor/deep_ep/__init__.py index d431fb5..94729ae 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/__init__.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/__init__.py @@ -7,7 +7,8 @@ is implemented with `torch.distributed` collectives on the process group DeepInf asynchronous event objects become synchronous no-ops with the same interface. Implemented: - EPBuffer / EPHandle expanded-layout MoE dispatch & combine (all-to-all) + EPBuffer / EPHandle expanded-layout MoE dispatch & combine (all-gather + reduce-scatter by default, + all-to-all with DEEPINFER_NEURON_EP_ALGO=a2a) EngramBuffer sharded n-gram embedding lookup (all-to-all of indices + rows) BucketBuffer / registered-storage sessions + batched all-gather (CP only) BucketSession diff --git a/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py index a8862ae..a6da26d 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/buffers/ep.py @@ -2,7 +2,16 @@ Expert-parallel dispatch/combine in the Ascend **expanded** layout (`do_expand=True`, one slot per (token, top-k selection) grouped by local expert), implemented with plain -`torch.distributed` all-to-all on the process group DeepInfer passes in. +`torch.distributed` collectives on the process group DeepInfer passes in. + +Transport, `DEEPINFER_NEURON_EP_ALGO` (issue #12): + * ``allgather`` (default): dispatch = ONE all-gather of every rank's tokens (x / sf / weights / routing as raw + bytes, padded to a common row count per rank, `DEEPINFER_NEURON_EP_AG_ROWS`); each rank then picks the rows + with a selection on its local experts. combine = each rank's bf16 partial for every gathered token (zero + where it has none) -> ONE fp32 SUM reduce-scatter (`DEEPINFER_NEURON_EP_REDUCE=rs|ar`). No all-to-all: + uneven `all_to_all_single` on the neuron backend fails / hangs at 16+ ranks. Same outputs as ``a2a`` + except the combine reduction order (see `combine`). + * ``a2a``: uneven all-to-all with host-synced split sizes (`comm.all_to_all_single`). Semantics follow `vendor/DeepEP-Ascend`: * dispatch: `deep_ep/utils/refs.py:dispatch` (token selection/ordering, `-1` masking of @@ -118,19 +127,65 @@ class EPHandle: class _RoutingPlan: - """Everything derived from `topk_idx` once, reused by cached dispatch and combine.""" + """Everything derived from `topk_idx` once, reused by cached dispatch and combine. - __slots__ = ('num_tokens', 'num_topk', 'num_local_experts', 'expert_alignment', + `algo == 'a2a'`: `send_*` / `recv_counts` describe the all-to-all exchange. + `algo == 'allgather'`: every rank gathers `gather_rows` rows from each rank (`[num_ranks * gather_rows]` + gathered rows, row `r * gather_rows + t` = token `t` of rank `r`); `recv_rows` are the gathered rows with at + least one selection on this rank, in gathered order (= the a2a receive order). + """ + + __slots__ = ('algo', 'num_tokens', 'num_topk', 'num_local_experts', 'expert_alignment', 'send_counts', 'recv_counts', 'num_send_tokens', 'num_recv_tokens', 'send_token_idx', 'send_row_rank', 'master_per_rank', + 'gather_rows', 'recv_rows', 'chunk_cache', 'local_valid', 'slot', 'slot_rows', 'valid_flat', 'num_expanded_tokens', 'counts_per_expert') def __init__(self, **kwargs): + for name in self.__slots__: + setattr(self, name, None) for key, value in kwargs.items(): setattr(self, key, value) +_EP_ALGOS = ('allgather', 'a2a') + + +def _ep_algo_from_env() -> str: + """`DEEPINFER_NEURON_EP_ALGO`: ``allgather`` (default) or ``a2a`` (the original all-to-all path).""" + import os + algo = os.environ.get('DEEPINFER_NEURON_EP_ALGO', 'allgather') + assert algo in _EP_ALGOS, f'DEEPINFER_NEURON_EP_ALGO must be one of {_EP_ALGOS}, got {algo!r}' + return algo + + +def _ag_rows_policy_from_env() -> str: + """`DEEPINFER_NEURON_EP_AG_ROWS`: rows each rank contributes to the all-gather. + + ``max`` (default): the global max token count of this call (one 0-dim MAX all-reduce, host sync) — keeps + prefill buffers small. ``cap``: always `num_max_tokens_per_rank` (static shapes, the decode-graph layout). + """ + import os + policy = os.environ.get('DEEPINFER_NEURON_EP_AG_ROWS', 'max') + assert policy in ('max', 'cap'), f'DEEPINFER_NEURON_EP_AG_ROWS must be max|cap, got {policy!r}' + return policy + + +# `DEEPINFER_NEURON_EP_DEBUG=1`: print per-dispatch routing sizes (host sync; diagnostics only). +_EP_DEBUG = __import__('os').environ.get('DEEPINFER_NEURON_EP_DEBUG', '0') == '1' + + +def _ag_chunk_bytes_from_env() -> int: + """`DEEPINFER_NEURON_EP_CHUNK_BYTES` (default 64 MiB): upper bound of one all-gather output / one + reduce-scatter input. Larger exchanges are split into row chunks (one collective each), so the transient + `num_ranks * rows * row_bytes` buffers stay bounded for long prefills (fp32 combine rows are 4 * hidden bytes).""" + import os + value = int(os.environ.get('DEEPINFER_NEURON_EP_CHUNK_BYTES', str(64 << 20))) + assert value > 0 + return value + + class EPBuffer(BufferBase): """Expert-parallel dispatch/combine over `torch.distributed` collectives. @@ -171,6 +226,10 @@ class EPBuffer(BufferBase): self.num_max_tokens_per_rank = num_max_tokens_per_rank self.num_allocated_qps = 0 self.num_lb_bytes = 0 + # Transport (issue #12): `allgather` (default) or the original `a2a`; see `_ep_algo_from_env`. + self.algo = _ep_algo_from_env() + self.ag_rows_policy = _ag_rows_policy_from_env() + self.ag_chunk_bytes = _ag_chunk_bytes_from_env() super().__init__(explicitly_destroy) if num_bytes is None: num_bytes = self.get_buffer_size_hint( @@ -228,9 +287,61 @@ class EPBuffer(BufferBase): # --- Routing ------------------------------------------------------------------- + def _expand_layout(self, recv_topk_idx: torch.Tensor, num_experts: int, + num_max_tokens_per_rank: int, expert_alignment: int, do_cpu_sync: bool): + """Expanded slots for the received (token, top-k) pairs, grouped by local expert. + + Shared by both algorithms; `recv_topk_idx` (`[num_recv_tokens, num_topk]`, global expert ids) is in + receive order (source rank, then source token). + """ + device = recv_topk_idx.device + num_ranks, rank_idx = self.num_ranks, self.rank_idx + num_recv_tokens, num_topk = recv_topk_idx.shape + num_local_experts = num_experts // num_ranks + + expert_base = rank_idx * num_local_experts + local_expert = recv_topk_idx - expert_base + local_valid = (recv_topk_idx >= expert_base) & (recv_topk_idx < expert_base + num_local_experts) + keys = torch.where(local_valid, local_expert, + torch.full_like(local_expert, num_local_experts)).reshape(-1) + counts_all = torch.zeros(num_local_experts + 1, dtype=torch.int64, device=device) + counts_all.index_add_(0, keys, torch.ones_like(keys)) + counts_per_expert = counts_all[:num_local_experts] + aligned_counts = align_tensor(counts_per_expert, expert_alignment) + expert_start = exclusive_cumsum(aligned_counts) + + # Stable sort by expert id => slots ordered by (expert, received row, top-k index). + perm = torch.argsort(keys, stable=True) + sorted_keys = keys[perm] + first_in_sorted = exclusive_cumsum(counts_all) + padded_start = torch.cat([expert_start, expert_start.new_zeros(1)]) + slot_sorted = padded_start[sorted_keys] + \ + (torch.arange(keys.numel(), device=device, dtype=torch.int64) - first_in_sorted[sorted_keys]) + slot_flat = torch.empty_like(slot_sorted) + slot_flat[perm] = slot_sorted + slot = torch.where(local_valid.reshape(-1), slot_flat, + torch.full_like(slot_flat, -1)).view(num_recv_tokens, num_topk) + + if do_cpu_sync: + # NEURON: host sync — exact expanded size. + num_expanded_tokens = int(aligned_counts.sum().item()) + else: + # Worst case, mirroring `csrc/buffers/ep.hpp:317-319`. + num_expanded_tokens = num_ranks * num_max_tokens_per_rank * min(num_topk, num_local_experts) + num_expanded_tokens += (expert_alignment - 1) * num_local_experts + num_expanded_tokens = align(num_expanded_tokens, expert_alignment) + + valid_flat = local_valid.reshape(-1).nonzero(as_tuple=True)[0] # NEURON: host sync + slot_rows = torch.div(valid_flat, num_topk, rounding_mode='floor') + return dict(num_local_experts=num_local_experts, expert_alignment=expert_alignment, + local_valid=local_valid, slot=slot, slot_rows=slot_rows, valid_flat=valid_flat, + num_expanded_tokens=num_expanded_tokens, counts_per_expert=counts_per_expert), \ + aligned_counts, expert_start + def _build_plan(self, topk_idx: torch.Tensor, num_experts: int, num_max_tokens_per_rank: int, expert_alignment: int, do_cpu_sync: bool): + """All-to-all routing (`DEEPINFER_NEURON_EP_ALGO=a2a`).""" device = topk_idx.device num_ranks, rank_idx = self.num_ranks, self.rank_idx num_tokens, num_topk = topk_idx.shape @@ -277,61 +388,194 @@ class EPBuffer(BufferBase): recv_src_idx = torch.empty((num_recv_tokens,), dtype=torch.int64, device=device) comm.all_to_all_single(recv_src_idx, send_src_idx, recv_counts, send_counts, self.group) - # Local expanded layout: group the received (token, top-k) pairs by local expert. - expert_base = rank_idx * num_local_experts - local_expert = recv_topk_idx - expert_base - local_valid = (recv_topk_idx >= expert_base) & (recv_topk_idx < expert_base + num_local_experts) - keys = torch.where(local_valid, local_expert, - torch.full_like(local_expert, num_local_experts)).reshape(-1) - counts_all = torch.zeros(num_local_experts + 1, dtype=torch.int64, device=device) - counts_all.index_add_(0, keys, torch.ones_like(keys)) - counts_per_expert = counts_all[:num_local_experts] - aligned_counts = align_tensor(counts_per_expert, expert_alignment) - expert_start = exclusive_cumsum(aligned_counts) - - # Stable sort by expert id => slots ordered by (expert, received row, top-k index). - perm = torch.argsort(keys, stable=True) - sorted_keys = keys[perm] - first_in_sorted = exclusive_cumsum(counts_all) - padded_start = torch.cat([expert_start, expert_start.new_zeros(1)]) - slot_sorted = padded_start[sorted_keys] + \ - (torch.arange(keys.numel(), device=device, dtype=torch.int64) - first_in_sorted[sorted_keys]) - slot_flat = torch.empty_like(slot_sorted) - slot_flat[perm] = slot_sorted - slot = torch.where(local_valid.reshape(-1), slot_flat, - torch.full_like(slot_flat, -1)).view(num_recv_tokens, num_topk) - - if do_cpu_sync: - # NEURON: host sync — exact expanded size. - num_expanded_tokens = int(aligned_counts.sum().item()) - else: - # Worst case, mirroring `csrc/buffers/ep.hpp:317-319`. - num_expanded_tokens = num_ranks * num_max_tokens_per_rank * min(num_topk, num_local_experts) - num_expanded_tokens += (expert_alignment - 1) * num_local_experts - num_expanded_tokens = align(num_expanded_tokens, expert_alignment) - - valid_flat = local_valid.reshape(-1).nonzero(as_tuple=True)[0] # NEURON: host sync - slot_rows = torch.div(valid_flat, num_topk, rounding_mode='floor') - + layout, aligned_counts, expert_start = self._expand_layout( + recv_topk_idx, num_experts, num_max_tokens_per_rank, expert_alignment, do_cpu_sync) plan = _RoutingPlan( - num_tokens=num_tokens, num_topk=num_topk, - num_local_experts=num_local_experts, expert_alignment=expert_alignment, + algo='a2a', num_tokens=num_tokens, num_topk=num_topk, send_counts=send_counts, recv_counts=recv_counts, num_send_tokens=num_send_tokens, num_recv_tokens=num_recv_tokens, send_token_idx=send_token_idx, send_row_rank=send_row_rank, - master_per_rank=master_per_rank, - local_valid=local_valid, slot=slot, slot_rows=slot_rows, valid_flat=valid_flat, - num_expanded_tokens=num_expanded_tokens, counts_per_expert=counts_per_expert) - return plan, recv_src_idx, recv_topk_idx, aligned_counts, expert_start + master_per_rank=master_per_rank, **layout) + psum_per_rank = torch.tensor(recv_counts, dtype=torch.int32, device=device).cumsum(0).to(torch.int32) + return plan, recv_src_idx, aligned_counts, expert_start, psum_per_rank + + # --- All-gather transport ----------------------------------------------------- + + @staticmethod + def _segments(parts): + """Byte layout of one packed row: `(offset, nbytes, dtype)` per part, 4-byte aligned segments.""" + segments, width = [], 0 + for t, _ in parts: + nbytes = t.size(1) * t.element_size() + segments.append((width, nbytes, t.dtype)) + width += align(nbytes, 4) + return segments, width + + def _gather_rows(self, parts, row_begin: int, row_end: int): + """All-gather rows `[row_begin, row_end)` of every rank's row-aligned tensors in ONE collective. + + `parts`: `(tensor [num_tokens, C_i], pad_byte)`; rows past a rank's `num_tokens` are padding (bytes = + `pad_byte`). Everything moves as raw bytes (FP8-safe: no float8 cat/index anywhere) in a + `[row_end - row_begin, W]` uint8 row buffer, gathered as a flat 1-D tensor (the shape verified on the + neuron backend; note int32 all-gathers are NOT exact there, see tests/neuron/deep_ep/test_collectives.py), + and split back into `[num_ranks * (row_end - row_begin), C_i]` tensors of the original dtypes + (row `r * n + i` = row `row_begin + i` of rank `r`). + """ + segments, width = self._segments(parts) + device = parts[0][0].device + num_tokens = parts[0][0].size(0) + n = row_end - row_begin + num_valid = max(0, min(row_end, num_tokens) - row_begin) + send = torch.zeros((n, width), dtype=torch.uint8, device=device) + for (t, pad_byte), (offset, nbytes, _) in zip(parts, segments): + if pad_byte != 0 and n > num_valid: + send[num_valid:, offset:offset + nbytes] = pad_byte + if num_valid and nbytes: + send[:num_valid, offset:offset + nbytes] = \ + t.contiguous().view(torch.uint8)[row_begin:row_begin + num_valid] + recv = torch.empty((self.num_ranks * n, width), dtype=torch.uint8, device=device) + if recv.numel(): + comm.all_gather_into_tensor(recv.view(-1), send.view(-1), self.group) + return [recv[:, offset:offset + nbytes].contiguous().view(dtype) for offset, nbytes, dtype in segments] + + def _rows_per_chunk(self, gather_rows: int, row_bytes: int) -> int: + """Rows per rank per collective so that one gathered / reduced buffer stays under the byte budget.""" + if gather_rows == 0: + return 0 + return max(1, min(gather_rows, self.ag_chunk_bytes // max(1, self.num_ranks * row_bytes))) + + def _chunks(self, plan: _RoutingPlan, rows_per_chunk: int): + """`(row_begin, row_end, pos, local)` per chunk of per-rank rows: `pos` = received rows (receive order) + whose source row falls in the chunk (None = all), `local` = their row in the gathered chunk.""" + cache = plan.chunk_cache + if rows_per_chunk in cache: + return cache[rows_per_chunk] + rows = plan.gather_rows + recv_rows = plan.recv_rows + chunks = [] + if rows_per_chunk >= rows: + if rows: + chunks.append((0, rows, None, recv_rows)) + else: + src_rank = torch.div(recv_rows, rows, rounding_mode='floor') + src_row = recv_rows - src_rank * rows + for begin in range(0, rows, rows_per_chunk): + end = min(rows, begin + rows_per_chunk) + pos = ((src_row >= begin) & (src_row < end)).nonzero(as_tuple=True)[0] # NEURON: host sync + chunks.append((begin, end, pos, src_rank[pos] * (end - begin) + (src_row[pos] - begin))) + cache[rows_per_chunk] = chunks + return chunks + + @staticmethod + def _payload_parts(x, sf, topk_weights): + parts = [(x, 0)] + if sf is not None: + parts.append((sf.contiguous(), 0)) + if topk_weights is not None: + parts.append((topk_weights.float().contiguous(), 0)) + return parts + + def _gather_received(self, plan: _RoutingPlan, parts): + """All-gather the payload `parts` and keep the rows this rank receives (receive order).""" + _, width = self._segments(parts) + rows_per_chunk = self._rows_per_chunk(plan.gather_rows, width) + chunks = self._chunks(plan, rows_per_chunk) + if len(chunks) == 1 and chunks[0][2] is None: + begin, end, _, local = chunks[0] + return [take_rows(g, local) for g in self._gather_rows(parts, begin, end)] + device = parts[0][0].device + outs = [empty_rows(plan.num_recv_tokens, t.size(1), t.dtype, device) for t, _ in parts] + for begin, end, pos, local in chunks: + gathered = self._gather_rows(parts, begin, end) # every rank joins every chunk's collective + if local.numel(): + for out, g in zip(outs, gathered): + put_rows(out, pos, take_rows(g, local)) + return outs + + def _split_payload(self, outs, sf, topk_weights): + recv_x = outs[0] + recv_sf = outs[1] if sf is not None else None + recv_w = outs[-1] if topk_weights is not None else None + return recv_x, recv_sf, recv_w + + def _global_num_tokens(self, num_tokens: int, device) -> int: + """Max token count over the group. NEURON: host sync (one 0-dim MAX all-reduce).""" + if self.num_ranks == 1: + return num_tokens + t = torch.tensor(num_tokens, dtype=torch.int32).to(device) + comm.all_reduce(t, self.group, op=dist.ReduceOp.MAX) + return int(t.item()) + + def _build_plan_allgather(self, x: torch.Tensor, sf: Optional[torch.Tensor], + topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor], + num_experts: int, num_max_tokens_per_rank: int, expert_alignment: int, + do_cpu_sync: bool): + """All-gather routing + payload (`DEEPINFER_NEURON_EP_ALGO=allgather`, issue #12). + + Every rank gathers every rank's tokens, keeps the ones with a selection on its local experts (in + gathered = (source rank, source token) order, identical to the all-to-all receive order) and builds the + same expanded layout as the a2a path. One packed all-gather (payload + routing) when the gathered buffer + fits `DEEPINFER_NEURON_EP_CHUNK_BYTES`; otherwise the routing first, then the payload in row chunks. + """ + device = x.device + num_ranks, rank_idx = self.num_ranks, self.rank_idx + num_tokens, num_topk = topk_idx.shape + num_local_experts = num_experts // num_ranks + assert num_experts % num_ranks == 0 + assert num_tokens <= num_max_tokens_per_rank + + if self.ag_rows_policy == 'cap': + gather_rows = num_max_tokens_per_rank + else: + gather_rows = self._global_num_tokens(num_tokens, device) # NEURON: host sync + plan = _RoutingPlan(algo='allgather', num_tokens=num_tokens, num_topk=num_topk, + gather_rows=gather_rows, chunk_cache={}) + + payload = self._payload_parts(x, sf, topk_weights) + idx_part = (topk_idx.to(torch.int32).contiguous(), 0xFF) # padding rows read as -1 (no selection) + _, packed_width = self._segments(payload + [idx_part]) + one_shot = self._rows_per_chunk(gather_rows, packed_width) >= gather_rows + if one_shot: + gathered = self._gather_rows(payload + [idx_part], 0, gather_rows) + g_idx = gathered.pop() + else: + g_idx, = self._gather_rows([idx_part], 0, gather_rows) + + g_idx = g_idx.to(torch.int64) # [R * rows, K] + expert_base = rank_idx * num_local_experts + recv_mask = ((g_idx >= expert_base) & (g_idx < expert_base + num_local_experts)).any(-1) + recv_rows = recv_mask.nonzero(as_tuple=True)[0] # NEURON: host sync + num_recv_tokens = int(recv_rows.numel()) + src_rank = torch.div(recv_rows, gather_rows, rounding_mode='floor') + recv_src_idx = src_rank * num_max_tokens_per_rank + (recv_rows - src_rank * gather_rows) + recv_topk_idx = g_idx[recv_rows] + recv_counts = recv_mask.view(num_ranks, gather_rows).sum(1) if gather_rows else \ + torch.zeros((num_ranks,), dtype=torch.int64, device=device) + + layout, aligned_counts, expert_start = self._expand_layout( + recv_topk_idx, num_experts, num_max_tokens_per_rank, expert_alignment, do_cpu_sync) + for key, value in layout.items(): + setattr(plan, key, value) + plan.recv_rows = recv_rows + plan.num_recv_tokens = num_recv_tokens + psum_per_rank = recv_counts.cumsum(0).to(torch.int32) + if one_shot: + outs = [take_rows(g, recv_rows) for g in gathered] + else: + outs = self._gather_received(plan, payload) + recv_payload = self._split_payload(outs, sf, topk_weights) + if _EP_DEBUG: # NEURON: host sync (debug only) + print(f'[deep_ep rank {rank_idx}] dispatch: tokens={num_tokens} gather_rows={gather_rows} ' + f'one_shot={one_shot} recv_tokens={num_recv_tokens} expanded={plan.num_expanded_tokens} ' + f'per_expert={plan.counts_per_expert.tolist()}', flush=True) + return plan, recv_src_idx, aligned_counts, expert_start, psum_per_rank, recv_payload # --- Dispatch ------------------------------------------------------------------ - def _move_payload(self, plan: _RoutingPlan, - x: torch.Tensor, sf: Optional[torch.Tensor], - topk_weights: Optional[torch.Tensor], - num_expanded_tokens: int, - use_tma_aligned_col_major_sf: bool): - """All-to-all the token payload and scatter it into the expanded layout.""" + def _a2a_payload(self, plan: _RoutingPlan, + x: torch.Tensor, sf: Optional[torch.Tensor], + topk_weights: Optional[torch.Tensor]): + """All-to-all the token payload (received rows in source-rank order).""" device = x.device num_topk = plan.num_topk hidden = x.size(1) @@ -356,12 +600,19 @@ class EPBuffer(BufferBase): recv_weights = torch.empty((plan.num_recv_tokens, num_topk), dtype=torch.float32, device=device) comm.all_to_all_single(recv_weights, send_weights, plan.recv_counts, plan.send_counts, self.group) + return recv_x, recv_sf, recv_weights + def _expand_payload(self, plan: _RoutingPlan, recv_x: torch.Tensor, + recv_sf: Optional[torch.Tensor], recv_weights: Optional[torch.Tensor], + num_expanded_tokens: int, use_tma_aligned_col_major_sf: bool): + """Scatter the received rows into the expanded layout.""" + device = recv_x.device + hidden = recv_x.size(1) # Expand: one row per valid (received token, top-k) pair, grouped by local expert. # The whole buffer is zero-initialized, so alignment padding is always bitwise zero # (Ascend only guarantees that with `do_zero_padding=True`). slots = plan.slot.reshape(-1)[plan.valid_flat] - expanded_x = zeros_rows(num_expanded_tokens, hidden, x.dtype, device) + expanded_x = zeros_rows(num_expanded_tokens, hidden, recv_x.dtype, device) if slots.numel(): put_rows(expanded_x, slots, take_rows(recv_x, plan.slot_rows)) @@ -387,7 +638,7 @@ class EPBuffer(BufferBase): if slots.numel(): expanded_weights[slots] = recv_weights.reshape(-1)[plan.valid_flat] - return recv_x, expanded_x, expanded_sf, expanded_weights + return expanded_x, expanded_sf, expanded_weights def dispatch(self, x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], @@ -455,21 +706,30 @@ class EPBuffer(BufferBase): plan = handle._plan assert plan is not None assert plan.num_tokens == x.size(0), 'Cached dispatch requires an unchanged token count' - _, expanded_x, expanded_sf, expanded_weights = self._move_payload( - plan, x, sf, topk_weights, handle.num_expanded_tokens, use_tma_aligned_col_major_sf) + if plan.algo == 'allgather': + recv_payload = self._split_payload( + self._gather_received(plan, self._payload_parts(x, sf, topk_weights)), sf, topk_weights) + else: + recv_payload = self._a2a_payload(plan, x, sf, topk_weights) + expanded_x, expanded_sf, expanded_weights = self._expand_payload( + plan, *recv_payload, handle.num_expanded_tokens, use_tma_aligned_col_major_sf) result_handle = handle else: - plan, recv_src_idx, recv_topk_idx, aligned_counts, expert_start = self._build_plan( - topk_idx, num_experts, num_max_tokens_per_rank, expert_alignment, do_cpu_sync) - _, expanded_x, expanded_sf, expanded_weights = self._move_payload( - plan, x, sf, topk_weights, plan.num_expanded_tokens, use_tma_aligned_col_major_sf) + if self.algo == 'allgather': + plan, recv_src_idx, aligned_counts, expert_start, psum_per_rank, recv_payload = \ + self._build_plan_allgather(x, sf, topk_idx, topk_weights, num_experts, + num_max_tokens_per_rank, expert_alignment, do_cpu_sync) + else: + plan, recv_src_idx, aligned_counts, expert_start, psum_per_rank = self._build_plan( + topk_idx, num_experts, num_max_tokens_per_rank, expert_alignment, do_cpu_sync) + recv_payload = self._a2a_payload(plan, x, sf, topk_weights) + expanded_x, expanded_sf, expanded_weights = self._expand_payload( + plan, *recv_payload, plan.num_expanded_tokens, use_tma_aligned_col_major_sf) device = x.device num_topk = plan.num_topk num_local_experts = plan.num_local_experts - psum_per_rank = torch.tensor(plan.recv_counts, dtype=torch.int32, - device=device).cumsum(0).to(torch.int32) psum_per_expert = (expert_start + plan.counts_per_expert).to(torch.int32) num_unaligned = plan.counts_per_expert.to(torch.int32) if cumulative_local_expert_recv_stats is not None: @@ -538,7 +798,12 @@ class EPBuffer(BufferBase): Matches `refs.py:combine(..., reduce_in_local=True)`: each receiving rank sums its slots for one source token in fp32 and stores bf16; the source rank then accumulates - the biases and the per-rank partials in fp32, in ascending master top-k slot order. + the biases and the per-rank partials in fp32. + * a2a: partials added in ascending master top-k slot order (bit-exact with the reference). + * allgather: the per-rank partials (bf16 values, upcast) are summed by an fp32 reduce-scatter in the + collective's own order and then added to the bias sum, i.e. `(b0 + b1) + (p_r0 + p_r1 + ...)` + instead of `((b0 + b1) + p_m0) + p_m1 ...`. Identical whenever a token's selections live on a + single rank; otherwise it can differ by fp32 rounding (<= ~1 bf16 ulp after the final cast). """ assert not self.destroyed assert handle is not None and handle.recv_src_metadata is not None @@ -552,50 +817,80 @@ class EPBuffer(BufferBase): hidden = x.size(1) bias_0, bias_1 = self._unpack_bias(bias) - # Per-rank reduction of this rank's slots (fp32 accumulation, bf16 result). - slot = plan.slot # [N, K], -1 where unused - valid = slot >= 0 - safe_slot = slot.clamp(min=0) - partial = torch.zeros((plan.num_recv_tokens, hidden), dtype=torch.float32, device=device) - for k in range(num_topk): # sequential adds to match the reference accumulation order - rows = x[safe_slot[:, k]].float() - partial += torch.where(valid[:, k:k + 1], rows, torch.zeros_like(rows)) - partial = partial.to(torch.bfloat16) - - send_weights = None if topk_weights is not None: assert topk_weights.dim() == 1 and topk_weights.size(0) == x.size(0) assert topk_weights.dtype == torch.float32 - gathered = topk_weights[safe_slot] - send_weights = torch.where(valid, gathered, torch.zeros_like(gathered)) - # Return the partials to the source ranks (reverse split sizes of dispatch). - back_x = torch.empty((plan.num_send_tokens, hidden), dtype=torch.bfloat16, device=device) - comm.all_to_all_single(byte_rows(back_x), byte_rows(partial.contiguous()), - plan.send_counts, plan.recv_counts, self.group) - back_weights = None - if send_weights is not None: - back_weights = torch.empty((plan.num_send_tokens, num_topk), dtype=torch.float32, device=device) - comm.all_to_all_single(back_weights, send_weights.contiguous(), - plan.send_counts, plan.recv_counts, self.group) + def local_partials(slot: torch.Tensor): + """This rank's per-token partials for the received rows `slot` (`[m, K]`, -1 = not here): fp32 + accumulation over the top-k slots in order, rounded to bf16 (`refs.py` reduce_in_local), plus the + top-k weights of the local slots (0 elsewhere).""" + valid = slot >= 0 + safe_slot = slot.clamp(min=0) + partial = torch.zeros((slot.size(0), hidden), dtype=torch.float32, device=device) + for k in range(num_topk): # sequential adds to match the reference accumulation order + rows = x[safe_slot[:, k]].float() + partial += torch.where(valid[:, k:k + 1], rows, torch.zeros_like(rows)) + weights = None + if topk_weights is not None: + gathered = topk_weights[safe_slot] + weights = torch.where(valid, gathered, torch.zeros_like(gathered)) + return partial.to(torch.bfloat16), weights - # Accumulate: biases first, then the partials in ascending master top-k slot order. acc = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=device) for bias_tensor in (bias_0, bias_1): if bias_tensor is not None: assert bias_tensor.shape == (num_tokens, hidden) and bias_tensor.dtype == x.dtype acc += bias_tensor.float() - row_master = plan.master_per_rank[plan.send_row_rank, plan.send_token_idx] - for k in range(num_topk): - selected = (row_master == k).nonzero(as_tuple=True)[0] # NEURON: host sync - if selected.numel(): - acc.index_add_(0, plan.send_token_idx[selected], back_x[selected].float()) - combined_x = acc.to(torch.bfloat16) combined_topk_weights = None - if back_weights is not None: - combined_topk_weights = torch.zeros((num_tokens, num_topk), dtype=torch.float32, device=device) - combined_topk_weights.index_add_(0, plan.send_token_idx, back_weights) + if plan.algo == 'allgather': + # Every gathered row gets this rank's partial (zero where it has no selection); an fp32 SUM + # reduce-scatter then hands each rank the sum over ranks for its own `gather_rows` rows. Row-chunked + # (`DEEPINFER_NEURON_EP_CHUNK_BYTES`): the partials are also computed per chunk, so every transient + # buffer stays bounded however many tokens this rank's experts received. + width = hidden + (num_topk if topk_weights is not None else 0) + if topk_weights is not None: + combined_topk_weights = torch.zeros((num_tokens, num_topk), dtype=torch.float32, device=device) + rows_per_chunk = self._rows_per_chunk(plan.gather_rows, width * 4) + for begin, end, pos, local in self._chunks(plan, rows_per_chunk): + n = end - begin + red_in = torch.zeros((self.num_ranks * n, width), dtype=torch.float32, device=device) + if local.numel(): + partial, weights = local_partials(plan.slot if pos is None else plan.slot[pos]) + row = partial.float() if weights is None else torch.cat([partial.float(), weights], dim=1) + red_in[local] = row + del partial, weights, row + red_out = torch.empty((n, width), dtype=torch.float32, device=device) + comm.reduce_scatter_tensor(red_out, red_in, self.group) # every rank joins every chunk + del red_in + num_valid = min(end, num_tokens) - begin + if num_valid > 0: + acc[begin:begin + num_valid] += red_out[:num_valid, :hidden] + if combined_topk_weights is not None: + combined_topk_weights[begin:begin + num_valid] = red_out[:num_valid, hidden:] + else: + partial, send_weights = local_partials(plan.slot) + # Return the partials to the source ranks (reverse split sizes of dispatch). + back_x = torch.empty((plan.num_send_tokens, hidden), dtype=torch.bfloat16, device=device) + comm.all_to_all_single(byte_rows(back_x), byte_rows(partial.contiguous()), + plan.send_counts, plan.recv_counts, self.group) + back_weights = None + if send_weights is not None: + back_weights = torch.empty((plan.num_send_tokens, num_topk), dtype=torch.float32, device=device) + comm.all_to_all_single(back_weights, send_weights.contiguous(), + plan.send_counts, plan.recv_counts, self.group) + + # Accumulate: biases first (above), then the partials in ascending master top-k slot order. + row_master = plan.master_per_rank[plan.send_row_rank, plan.send_token_idx] + for k in range(num_topk): + selected = (row_master == k).nonzero(as_tuple=True)[0] # NEURON: host sync + if selected.numel(): + acc.index_add_(0, plan.send_token_idx[selected], back_x[selected].float()) + if back_weights is not None: + combined_topk_weights = torch.zeros((num_tokens, num_topk), dtype=torch.float32, device=device) + combined_topk_weights.index_add_(0, plan.send_token_idx, back_weights) + combined_x = acc.to(torch.bfloat16) result = (combined_x, combined_topk_weights) event_overlap = EventOverlap(EventHandle()) diff --git a/deepinfer/neuron_ops/vendor/deep_ep/comm.py b/deepinfer/neuron_ops/vendor/deep_ep/comm.py index 824fff2..a0e049b 100644 --- a/deepinfer/neuron_ops/vendor/deep_ep/comm.py +++ b/deepinfer/neuron_ops/vendor/deep_ep/comm.py @@ -202,6 +202,50 @@ def all_gather_into_tensor(output: torch.Tensor, input_: torch.Tensor, return output +def all_reduce(tensor: torch.Tensor, group: dist.ProcessGroup, + op: dist.ReduceOp = dist.ReduceOp.SUM) -> torch.Tensor: + """In-place `dist.all_reduce` with a single-rank fast path.""" + if group.size() > 1: + dist.all_reduce(tensor, op=op, group=group) + return tensor + + +def _default_reduce_scatter_mode() -> str: + mode = os.environ.get('DEEPINFER_NEURON_EP_REDUCE', 'rs') + assert mode in ('rs', 'ar'), f'DEEPINFER_NEURON_EP_REDUCE must be rs|ar, got {mode!r}' + return mode + + +def reduce_scatter_tensor(output: torch.Tensor, input_: torch.Tensor, group: dist.ProcessGroup, + mode: Optional[str] = None) -> torch.Tensor: + """SUM reduce-scatter along dim 0: `output = sum_over_ranks(input_)[rank * n:(rank + 1) * n]`. + + Runs on flat 1-D views (the shapes verified on the neuron backend). `mode` + (default `DEEPINFER_NEURON_EP_REDUCE`, ``rs``): ``rs`` = `dist.reduce_scatter_tensor`, ``ar`` = `all_reduce` + of the whole input and a local slice (R x the traffic, kept as a fallback). gloo has no reduce-scatter, so + host groups always take ``ar``. + """ + num_ranks = group.size() + assert input_.numel() == num_ranks * output.numel() + if num_ranks == 1: + output.view(-1).copy_(input_.reshape(-1)) + return output + mode = mode or _default_reduce_scatter_mode() + flat_in = input_.contiguous().view(-1) + if mode == 'rs' and dist.get_backend(group) != 'gloo': + flat_out = output.view(-1) if output.is_contiguous() else torch.empty( + (output.numel(),), dtype=output.dtype, device=output.device) + dist.reduce_scatter_tensor(flat_out, flat_in, op=dist.ReduceOp.SUM, group=group) + if flat_out.data_ptr() != output.data_ptr(): + output.copy_(flat_out.view(output.shape)) + return output + buf = flat_in.clone() + dist.all_reduce(buf, op=dist.ReduceOp.SUM, group=group) + n = output.numel() + output.copy_(buf.narrow(0, group.rank() * n, n).view(output.shape)) + return output + + def all_gather_object(obj: Any, group: dist.ProcessGroup) -> List[Any]: if group.size() == 1: return [obj] diff --git a/tests/neuron/deep_ep/test_collectives.py b/tests/neuron/deep_ep/test_collectives.py new file mode 100644 index 0000000..644574d --- /dev/null +++ b/tests/neuron/deep_ep/test_collectives.py @@ -0,0 +1,86 @@ +"""The collectives the all-gather EP path relies on (`deep_ep.comm`), checked against host math. + +all_gather_into_tensor (uint8 / int32 / float32, 1-D), reduce_scatter_tensor and all_reduce (float32, SUM) on the +group handed to the buffers. DEEP_EP_TEST_DEVICE=neuron runs them on a torch-neuronx ``neuron`` group (see dev.py). +Run: python test_collectives.py --num-processes {2,4,...,64} +""" + +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")) +from deep_ep import comm # noqa: E402 +from dev import DEVICE, init_dist, to_dev, to_host # noqa: E402 + + +def rank_data(r: int, n: int, dtype: torch.dtype, step: int) -> torch.Tensor: + g = torch.Generator().manual_seed(1000 * step + r) + if dtype == torch.uint8: + return torch.randint(0, 256, (n,), generator=g, dtype=torch.int32).to(torch.uint8) + if dtype == torch.int32: + # |v| < 2**24 only: a full-range int32 all-gather is WRONG on a 16-rank neuron group once the output reaches + # 1 MiB (every chunk corrupted; consistent with an fp32 round trip), while uint8 / fp32 of the same size are + # exact. The EP path therefore moves everything as uint8 bytes. + return torch.randint(-2**24, 2**24, (n,), generator=g, dtype=torch.int64).to(torch.int32) + return torch.randn((n,), generator=g) + + +def partial_data(r: int, idx: slice, step: int) -> torch.Tensor: + """bf16-representable fp32 values, ~70% exact zeros, like the combine partials (closed form in the index).""" + i = torch.arange(idx.start, idx.stop, dtype=torch.int64) + v = (((i * 2654435761 + r * 97 + step * 13) % 255) - 127).float() / 64 + return v.masked_fill((i * 31 + r + step) % 10 < 7, 0.0) + + +def worker(local_rank: int, num_ranks: int, args): + group = init_dist(local_rank, num_ranks, args.port) + rank = group.rank() + sizes = (1, 7, 4096 + 3, 3 * 7168 + 4, 65536 * 4 + 3) + for step, n in enumerate([n for _ in range(args.repeat) for n in sizes]): + # all_gather + for dtype in (torch.uint8, torch.int32, torch.float32): + local = rank_data(rank, n, dtype, step) + out = to_dev(torch.empty((num_ranks * n,), dtype=dtype)) + comm.all_gather_into_tensor(out, to_dev(local), group) + ref = torch.cat([rank_data(r, n, dtype, step) for r in range(num_ranks)]) + got = to_host(out) + if not torch.equal(got, ref): + bad = (got != ref).view(num_ranks, n) + rows = bad.any(1).nonzero().flatten().tolist() + first = bad.flatten().nonzero().flatten()[:4].tolist() + raise AssertionError(f'rank {rank}: all_gather {dtype} n={n}: wrong chunks from ranks {rows}, ' + f'{int(bad.sum())} wrong elements, first flat idx {first}') + # reduce_scatter / all_reduce + slice (fp32 SUM). Inputs are a closed-form function of (rank, index), so + # the reference only materializes this rank's slice of every rank's input. + my = slice(rank * n, (rank + 1) * n) + terms = torch.stack([partial_data(r, my, step) for r in range(num_ranks)]).double() + my_ref, tol = terms.sum(0), 1e-6 * (terms.abs().sum(0).max().item() + 1) + for mode in ('rs', 'ar'): + out = to_dev(torch.empty((n,), dtype=torch.float32)) + comm.reduce_scatter_tensor(out, to_dev(partial_data(rank, slice(0, num_ranks * n), step)), group, + mode=mode) + err = (to_host(out).double() - my_ref).abs().max().item() + assert err <= tol, f'reduce_scatter[{mode}] n={n} err={err} tol={tol}' + if rank == 0: + print(f' ok: n={n}', flush=True) + if rank == 0: + print(f'collective tests passed (world_size={num_ranks}, device={DEVICE})', flush=True) + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num-processes', type=int, default=2) + parser.add_argument('--port', type=int, default=29541) + parser.add_argument('--repeat', type=int, default=1) + 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) diff --git a/tests/neuron/deep_ep/test_ep.py b/tests/neuron/deep_ep/test_ep.py index 60aaa7a..d134e6a 100644 --- a/tests/neuron/deep_ep/test_ep.py +++ b/tests/neuron/deep_ep/test_ep.py @@ -209,7 +209,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): 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=}' + check_combined(buffer, combined_x, ref_combined_x, ref_y, bias, f'{tag} {max_diff=}') assert torch.equal(combined_weights, topk_weights) if num_bias == 0 and not col_major: @@ -217,7 +217,7 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): buffer, 'combine', defer_epilogue, x=d_input_for_combine, handle=d_handle, bias=to_dev(bias), num_sms=7, async_with_compute_stream=True) combined_x2 = to_host(combined_x2) - assert torch.equal(combined_x2, ref_combined_x) + check_combined(buffer, combined_x2, ref_combined_x, ref_y, bias, f'{tag} (no weights)') assert combined_weights2 is None # --- Cached dispatch replay --- @@ -271,6 +271,32 @@ def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync): assert cached_w is None +def check_combined(buffer, combined_x, ref_combined_x, ref_y, bias, tag): + """`a2a` reproduces the reference reduction order exactly (bit-exact check). + + `allgather` sums the per-rank bf16 partials with an fp32 reduce-scatter (collective order) and adds the bias + afterwards: `(b0 + b1) + (p_r0 + p_r1 + ...)` instead of `((b0 + b1) + p_m0) + p_m1 + ...` in master top-k slot + order. Tokens whose selections all live on one rank stay bit-exact; the rest may differ by fp32 rounding, + bounded by one bf16 ulp of the result plus the fp32 error of re-associating the terms. + """ + if buffer.algo == 'a2a': + assert torch.equal(combined_x, ref_combined_x), tag + return + terms = ref_y.float().abs().sum(1) + for b in (bias if isinstance(bias, tuple) else (bias,)): + if b is not None: + terms = terms + b.float().abs() + out, ref = combined_x.float(), ref_combined_x.float() + tol = torch.maximum(out.abs(), ref.abs()) * 2 ** -7 + terms * 2 ** -20 + diff = (out - ref).abs() + assert bool((diff <= tol).all()), f'{tag} max_excess={(diff - tol).max().item()}' + CHECK_STATS[0] += int((diff > 0).sum()) + CHECK_STATS[1] += diff.numel() + + +CHECK_STATS = [0, 0] # (non-identical elements, compared elements) for the allgather combine + + 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=HOST), @@ -300,7 +326,8 @@ def worker(local_rank: int, num_ranks: int, args: argparse.Namespace): buffer.destroy() if rank == 0: - print(f'EP tests passed (world_size={num_ranks})', flush=True) + print(f'EP tests passed (world_size={num_ranks}, algo={buffer.algo}, ' + f'combine non-identical elements on rank 0: {CHECK_STATS[0]}/{CHECK_STATS[1]})', flush=True) dist.destroy_process_group()