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 <noreply@anthropic.com>

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

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

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]

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)

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