Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / collectives.py: 96%
194 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2025-2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
16"""expert_parallel.collectives: backend-dispatched EP all_to_all.
18NCCL/HCCL use the ragged a2a (``_EPAllToAllUneven``, zero-padding); gloo and
19other backends that do not support ragged a2a use pad-to-max +
20``all_to_all_single`` (``_EPAllToAllPadded``). Both paths are numerically
21equivalent (padding only adds filler rows that do not participate in
22computation).
24When every per-peer count is equal -- which is exactly what a balanced routing
25plan produces -- the split sizes carry no information, so ``HP_EP_EQUAL_A2A``
26swaps the ragged exchange for the plain equal-length ``all_to_all_single``,
27whose backend kernel is cheaper than alltoallv. The rows, their order and the
28values are identical on both paths. On top of the kernel, the exchange has to
29stay *lazy*: the champion issues its token exchange through the async entry and
30hides it behind the shared-expert MLP, and a plain c10d collective orders the
31stream it is called on at issue time, which closes that window. ``=1``
32therefore means split-free *and* lazy (``_EPAllToAllEqualLazy``: the exchange is
33issued on a stream of its own and the caller waits on an event), while ``=eager``
34forces the split-free call back onto the caller's stream for an A/B of the two.
36Split out of components/distributed/ep_utils.py in stage 4e.
37"""
39import os
40import threading
41from typing import Any, Optional
42import torch
43import torch.distributed as dist
44from torch import _check, _check_is_size
45from torch.distributed import _functional_collectives
47from hyper_parallel.core.dtensor._utils import (
48 differentiable_all_to_all_single_async,
49 get_device_handle,
50)
51from hyper_parallel.core.utils.communication import _is_tracing_or_compiling
53_UNEVEN_A2A_BACKENDS = ("nccl", "hccl")
55# A ragged exchange reaches NCCL/HCCL as alltoallv, which takes a per-peer count
56# vector and runs a count-driven kernel. Under a balanced routing plan every
57# per-peer count is the same, so the counts add nothing and the plain
58# equal-length all_to_all_single (alltoall) can be issued instead. Opt-in while
59# it is being measured (HP_EP_EQUAL_A2A).
60_EQUAL_A2A_OFF = "off"
61_EQUAL_A2A_LAZY = "lazy"
62_EQUAL_A2A_EAGER = "eager"
63_EQUAL_A2A_MODES = {
64 "0": _EQUAL_A2A_OFF,
65 "1": _EQUAL_A2A_LAZY,
66 "eager": _EQUAL_A2A_EAGER,
67}
68_EQUAL_A2A_RAW = os.environ.get("HP_EP_EQUAL_A2A", "0")
70# Ragged (uneven) exchange kernel: the list form issues one
71# ``dist.all_to_all`` with per-peer tensors, which HCCL serves through
72# alltoallv; ``HP_EP_A2A_SINGLE=1`` issues the identical exchange as one
73# ``all_to_all_single`` with split sizes, the form the micro benchmark
74# measures at ~2x the bandwidth (106.4 vs 52.9 GB/s). Default off. Scope: this
75# only affects THIS synchronous entry; the training MoE path calls
76# ``ep_all_to_all_async``, whose ``_AsyncA2ALazyBwd`` already issues the split-sized
77# ``all_to_all_single``, so the knob is a no-op there (measured: neutral).
78_A2A_SINGLE = os.environ.get("HP_EP_A2A_SINGLE", "0") == "1"
80# Lazily created, then reused for every lazy exchange (see
81# :func:`_lazy_a2a_resources`): the stream the split-free collective is issued
82# on, the event that hands it the payload, and the event that hands the result
83# back to the consumer.
84_LAZY_A2A_STREAM = None
85_LAZY_A2A_READY_EVENT = None
86_LAZY_A2A_DONE_EVENT = None
87_LAZY_A2A_LOCK = threading.Lock()
90def ep_exchange_counts(counts: torch.Tensor, group: Any) -> torch.Tensor:
91 """Exchange one integer count per peer, waiting before reading the result.
93 Args:
94 counts: Non-differentiable send counts in process-group rank order.
95 group: Expert-parallel process group.
97 Returns:
98 Counts received from each peer in the same rank order.
99 """
100 if _is_tracing_or_compiling():
101 return _functional_collectives.wait_tensor(
102 _functional_collectives.all_to_all_single(counts, None, None, group)
103 )
104 received = torch.empty_like(counts)
105 dist.all_to_all_single(received, counts, group=group)
106 return received
109def ep_split_sizes(counts: torch.Tensor, total: Optional[int] = None) -> list[int]:
110 """Read per-call segment sizes, retaining data dependencies during capture.
112 Args:
113 counts: One-dimensional integer tensor of segment lengths.
114 total: Expected number of rows, when known from a payload's shape.
116 Returns:
117 Concrete eager sizes or symbolic sizes evaluated by the captured graph.
118 """
119 if not _is_tracing_or_compiling():
120 return counts.tolist()
121 sizes = counts.cpu().tolist()
122 for size in sizes:
123 _check_is_size(size)
124 if total is not None:
125 _check(sum(sizes) == total)
126 return sizes
129def _backend_supports_uneven_a2a(group) -> bool:
130 return dist.get_backend(group) in _UNEVEN_A2A_BACKENDS
133def _equal_a2a_mode() -> str:
134 """Resolve ``HP_EP_EQUAL_A2A`` to the split-free exchange's mode.
136 ``"0"`` (the default) keeps the ragged exchange, ``"1"`` takes the
137 split-free exchange and defers its wait, and ``"eager"`` takes the
138 split-free exchange while ordering the caller's stream on it at issue time
139 -- the behaviour ``"1"`` used to have, kept so the kernel's gain and the
140 lost overlap can be measured apart.
142 Returns:
143 One of ``"off"``, ``"lazy"``, ``"eager"``.
145 Raises:
146 ValueError: If the knob is none of the accepted values; a typo must fail
147 loudly instead of silently leaving the fast path off.
148 """
149 raw = _EQUAL_A2A_RAW.strip()
150 try:
151 return _EQUAL_A2A_MODES[raw]
152 except KeyError:
153 raise ValueError(
154 f"HP_EP_EQUAL_A2A must be one of "
155 f"{sorted(_EQUAL_A2A_MODES)}, but got {raw!r}"
156 ) from None
159class _EPAllToAllUneven(torch.autograd.Function): # pylint: disable=abstract-method
160 """Ragged all_to_all (NCCL/HCCL production path): split by send/recv counts.
162 forward: split(x, send_counts) -> dist.all_to_all(out_list, in_list) -> cat
163 backward: swap send/recv counts and run the ragged all_to_all again
164 (a2a is self-inverse).
165 """
167 @staticmethod
168 def forward(
169 ctx: Any,
170 x: torch.Tensor,
171 send_counts: list[int],
172 recv_counts: list[int],
173 group: Any,
174 ) -> torch.Tensor: # pylint: disable=arguments-differ
175 """Run the ragged all_to_all and retain the counts for backward."""
176 ctx.send_counts = send_counts
177 ctx.recv_counts = recv_counts
178 ctx.group = group
179 out = x.new_empty((sum(recv_counts),) + tuple(x.shape[1:]))
180 if _A2A_SINGLE:
181 dist.all_to_all_single(
182 out,
183 x.contiguous(),
184 output_split_sizes=[int(count) for count in recv_counts],
185 input_split_sizes=[int(count) for count in send_counts],
186 group=group,
187 )
188 else:
189 dist.all_to_all(list(out.split(recv_counts)),
190 list(x.split(send_counts)), group=group)
191 return out
193 @staticmethod
194 def backward(
195 ctx: Any,
196 grad_output: torch.Tensor,
197 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
198 """Swap send/recv counts and re-run the self-inverse ragged all_to_all."""
199 grad = _EPAllToAllUneven.apply(
200 grad_output.contiguous(), ctx.recv_counts, ctx.send_counts, ctx.group)
201 return grad, None, None, None
204def _equal_a2a_exchange(
205 x: torch.Tensor,
206 rows_per_peer: int,
207 ep_size: int,
208 group: Any,
209) -> torch.Tensor:
210 """Move ``rows_per_peer`` rows to each peer and return the peer-major result.
212 The split-free call: the backend derives the chunking from the tensor shape
213 alone, and the received chunks are concatenated in peer order. Both the
214 eager and the lazy equal-length exchange issue this, so a change here cannot
215 make the two disagree about the layout.
216 """
217 out = x.new_empty((rows_per_peer * ep_size,) + tuple(x.shape[1:]))
218 dist.all_to_all_single(out, x, group=group)
219 return out
222class _EPAllToAllEqual(torch.autograd.Function): # pylint: disable=abstract-method
223 """Equal-length ``all_to_all_single`` (NCCL/HCCL path when counts are uniform).
225 Without split sizes every peer gets ``rows_per_peer`` rows -- the backend
226 derives the chunking from the tensor shape alone -- and the received chunks
227 are concatenated in peer order. That is exactly the layout
228 :class:`_EPAllToAllUneven` produces when all counts are equal, so this path
229 is interchangeable with it (see :func:`_equal_a2a_rows` for when it applies).
231 forward: all_to_all_single(no splits) -> [ep_size * rows_per_peer, ...]
232 backward: the same exchange; an equal-length a2a is its own inverse, and
233 with a uniform plan the reverse chunks are the same size.
235 ``forward`` issues the exchange on the caller's stream, which orders that
236 stream on the transfer before it returns; :class:`_EPAllToAllEqualLazy` is
237 the same exchange with the ordering left to the caller.
238 """
240 @staticmethod
241 def forward(
242 ctx: Any,
243 x: torch.Tensor,
244 rows_per_peer: int,
245 ep_size: int,
246 group: Any,
247 ) -> torch.Tensor: # pylint: disable=arguments-differ
248 """Run the equal-length all_to_all and retain its geometry for backward.
250 Args:
251 ctx: Autograd context of this exchange.
252 x: Payload to exchange; it holds ``ep_size * rows_per_peer`` rows.
253 rows_per_peer: Rows handed to each peer.
254 ep_size: Number of ranks in ``group``.
255 group: Process group to exchange over.
257 Returns:
258 The received rows, concatenated in peer order.
259 """
260 ctx.rows_per_peer = rows_per_peer
261 ctx.ep_size = ep_size
262 ctx.group = group
263 return _equal_a2a_exchange(x, rows_per_peer, ep_size, group)
265 @staticmethod
266 def backward(
267 ctx: Any,
268 grad_output: torch.Tensor,
269 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
270 """Re-run the self-inverse equal-length exchange on the output gradient.
272 Args:
273 ctx: Autograd context of the forward exchange.
274 grad_output: Output gradient, one row per exchanged row.
276 Returns:
277 The input gradient, followed by ``None`` for the non-tensor arguments.
278 """
279 grad = _equal_a2a_exchange(
280 grad_output.contiguous(), ctx.rows_per_peer, ctx.ep_size, ctx.group)
281 return grad, None, None, None
284def _lazy_a2a_resources() -> tuple[Any, Any, Any]:
285 """Return the ``(stream, ready_event, done_event)`` the lazy exchange reuses.
287 One stream and two events for the whole process, which is deliberate: a
288 fresh event per exchange drains the runtime's event pool after a few hundred
289 exchanges (the run then dies in ``AclQueryEventRecordedStatus``), and one
290 stream is enough because every record and wait below is enqueued in host
291 order -- a wait therefore never observes a *stale* record, at worst a later
292 one, which only makes it wait longer than it strictly had to.
294 Returns:
295 The shared comm stream, the event that hands the payload over to it, and
296 the event that hands the result back.
297 """
298 global _LAZY_A2A_STREAM, _LAZY_A2A_READY_EVENT, _LAZY_A2A_DONE_EVENT # pylint: disable=global-statement
299 with _LAZY_A2A_LOCK:
300 if _LAZY_A2A_DONE_EVENT is None:
301 _LAZY_A2A_STREAM = get_device_handle().Stream()
302 _LAZY_A2A_READY_EVENT = get_device_handle().Event()
303 _LAZY_A2A_DONE_EVENT = get_device_handle().Event()
304 return _LAZY_A2A_STREAM, _LAZY_A2A_READY_EVENT, _LAZY_A2A_DONE_EVENT
307def _record_stream(tensor: torch.Tensor, stream: Any) -> None:
308 """Tell the caching allocator that ``stream`` touches ``tensor``'s storage.
310 Wrapped rather than called inline so the exchange can be driven with a fake
311 stream in unit tests: ``Tensor.record_stream`` rejects anything that is not a
312 real device stream.
313 """
314 tensor.record_stream(stream)
317def _issue_split_free_a2a(out: torch.Tensor, x: torch.Tensor, group: Any) -> Any:
318 """Issue the split-free exchange on the shared comm stream.
320 ``dist.all_to_all_single`` without split sizes is the one call shape whose
321 backend kernel is the plain alltoall, and it is a *blocking* collective:
322 c10d orders the stream it is called on after the transfer before returning,
323 which is what would close the caller's overlap window. Issuing it inside a
324 stream context moves that ordering onto the comm stream. c10d only orders
325 the collective against the stream it is called on, so the payload's
326 producers are re-connected to that stream explicitly through the ready
327 event.
329 Args:
330 out: Receive buffer, ``[ep_size * rows_per_peer, ...]``; written here.
331 x: Contiguous payload, ``[ep_size * rows_per_peer, ...]``; read here.
332 group: Process group to exchange over.
334 Returns:
335 The event recorded when the exchange has completed; the consumer's
336 stream waits on it (see :class:`_PendingEqualA2A`).
337 """
338 compute_stream = get_device_handle().current_stream()
339 stream, ready_event, done_event = _lazy_a2a_resources()
340 stream_context = get_device_handle().stream
341 ready_event.record(compute_stream)
342 with stream_context(stream):
343 ready_event.wait(stream)
344 dist.all_to_all_single(out, x, group=group)
345 # Both buffers outlive the call by different means: the receive buffer
346 # was allocated on the compute stream, and the payload is a temporary of
347 # the caller's dispatch, so its storage may be freed and reused while
348 # this stream is still reading it.
349 _record_stream(out, stream)
350 _record_stream(x, stream)
351 done_event.record(stream)
352 return done_event
355class _PendingEqualA2A:
356 """Handle to a split-free exchange that is still in flight.
358 ``wait()`` is the only way to read the result, and *where* it is called is
359 the caller's decision: the exchange is issued while the dispatch runs, but
360 the consumer's stream is only ordered on it at ``wait()``, so independent
361 work issued in between (the shared-expert MLP an ``overlap_fn`` runs) works
362 against the transfer instead of after it. Reading the exchange without
363 waiting is not an option -- the handle is not a tensor.
365 Attributes:
366 completed: Whether a stream has already been ordered on the exchange.
367 """
369 __slots__ = ("completed", "_event", "_tensor")
371 def __init__(self, tensor: torch.Tensor, event: Any, completed: bool = False) -> None:
372 """Wrap an issued-but-unwaited exchange result."""
373 self._tensor = tensor
374 self._event = event
375 self.completed = completed
377 def wait(self) -> torch.Tensor:
378 """Order the current stream after the exchange and return its result."""
379 if not self.completed:
380 self._event.wait(get_device_handle().current_stream())
381 self.completed = True
382 return self._tensor
384 def squeeze(self, dim: int) -> "_PendingEqualA2A":
385 """Return a handle to the squeezed result (a view, so the wait is pending).
387 The dispatched expert indices arrive as ``[rows, 1]`` and their consumers
388 want ``[rows]``. Squeezing is metadata-only, so it can be taken here
389 without materializing: a read here would enqueue the wait before the
390 caller reaches its independent work.
392 Args:
393 dim: Dimension to squeeze, as ``Tensor.squeeze`` takes it.
395 Returns:
396 A handle to the squeezed result; it still holds the exchange's result
397 as pending as this one does.
398 """
399 return _PendingEqualA2A(self._tensor.squeeze(dim), self._event, self.completed)
401 def __repr__(self) -> str:
402 """Name the pending exchange and what is known about it without reading it."""
403 return (f"PendingEqualA2A(shape={tuple(self._tensor.shape)}, "
404 f"completed={self.completed})")
407def wait_ep_all_to_all(value: Any) -> Any:
408 """Materialize a pending EP exchange.
410 Every consumer of an :func:`ep_all_to_all_async` result must pass it through
411 here before reading it: with ``HP_EP_EQUAL_A2A=1`` and a uniform plan that
412 entry returns a :class:`_PendingEqualA2A` instead of a tensor.
414 Args:
415 value: The pending handle, an already-materialized tensor from any other
416 path of :func:`ep_all_to_all_async`, or ``None``.
418 Returns:
419 The exchange result as a regular tensor; ``value`` unchanged when it is
420 not a pending handle.
421 """
422 if isinstance(value, _PendingEqualA2A):
423 return value.wait()
424 return value
427class _EPAllToAllEqualLazy(torch.autograd.Function): # pylint: disable=abstract-method
428 """Split-free equal-length exchange issued on the comm stream, waited lazily.
430 Same exchange as :class:`_EPAllToAllEqual` and just as differentiable, but
431 the forward leaves the caller's stream alone: it issues on the shared comm
432 stream and returns the buffer together with the event that says when the
433 buffer is final, so the caller can synchronize when it needs the values
434 rather than when it issues the exchange. Both travel out of ``forward`` --
435 the buffer as the autograd output, the event as a second, non-tensor output.
437 The backward is the plain split-free exchange on the caller's stream: the
438 engine consumes its result immediately (it accumulates it onto the input's
439 ``grad``), so there is nothing to defer.
440 """
442 @staticmethod
443 def forward( # pylint: disable=arguments-differ
444 ctx: Any,
445 x: torch.Tensor,
446 rows_per_peer: int,
447 ep_size: int,
448 group: Any,
449 ) -> tuple[torch.Tensor, Any]:
450 """Issue the split-free exchange and return ``(buffer, done event)``.
452 Args:
453 ctx: Autograd context of this exchange.
454 x: Contiguous payload, ``[ep_size * rows_per_peer, ...]``.
455 rows_per_peer: Rows handed to each peer.
456 ep_size: Number of ranks in ``group``.
457 group: Process group to exchange over.
459 Returns:
460 The receive buffer the exchange writes -- the output the autograd
461 graph is built on -- and the event that says when it is final.
462 """
463 ctx.rows_per_peer = rows_per_peer
464 ctx.ep_size = ep_size
465 ctx.group = group
466 out = x.new_empty((rows_per_peer * ep_size,) + tuple(x.shape[1:]))
467 return out, _issue_split_free_a2a(out, x, group)
469 @staticmethod
470 def backward( # pylint: disable=arguments-differ
471 ctx: Any,
472 grad_output: torch.Tensor,
473 _event_grad: Any, # pylint: disable=invalid-name
474 ) -> tuple[torch.Tensor, None, None, None]:
475 """Re-run the self-inverse exchange, on the caller's stream this time.
477 Args:
478 ctx: Autograd context of the forward exchange.
479 grad_output: Output gradient, one row per exchanged row.
480 _event_grad: Gradient slot of the forward's event output; unused.
482 Returns:
483 The input gradient, followed by ``None`` for the non-tensor arguments.
484 """
485 grad = _equal_a2a_exchange(
486 grad_output.contiguous(), ctx.rows_per_peer, ctx.ep_size, ctx.group)
487 return grad, None, None, None
490class _EPAllToAllPadded(torch.autograd.Function): # pylint: disable=abstract-method
491 """pad-to-max + all_to_all_single (gloo test path).
493 forward: pad each dest chunk to the global max(counts) (a2a_single
494 requires equal-length chunks per rank -> pad_to must be
495 globally consistent, obtained via all_reduce MAX)
496 -> a2a_single -> unpad by recv_counts;
497 backward: pad by recv_counts -> a2a_single (equal-length self-inverse)
498 -> unpad by send_counts.
499 """
501 @staticmethod
502 def _pad_and_exchange(x, counts, pad_to, group):
503 """Pad each chunk to pad_to, run equal-length a2a_single, return [ep*pad_to, ...]."""
504 chunks = []
505 for chunk, n in zip(x.split(counts), counts):
506 if n < pad_to:
507 pad = x.new_zeros((pad_to - n,) + tuple(x.shape[1:]))
508 chunk = torch.cat([chunk, pad])
509 chunks.append(chunk)
510 send = torch.cat(chunks).contiguous()
511 recv = torch.empty_like(send)
512 dist.all_to_all_single(recv, send, group=group)
513 return recv
515 @staticmethod
516 def _unpad(recv, counts, pad_to):
517 """Take the valid rows from the equal-length buffer per counts and cat."""
518 pieces = []
519 for i, n in enumerate(counts):
520 if n > 0:
521 pieces.append(recv[i * pad_to: i * pad_to + n])
522 if not pieces:
523 return recv.new_zeros((0,) + tuple(recv.shape[1:]))
524 return torch.cat(pieces)
526 @staticmethod
527 def forward(
528 ctx: Any,
529 x: torch.Tensor,
530 send_counts: list[int],
531 recv_counts: list[int],
532 group: Any,
533 ) -> torch.Tensor: # pylint: disable=arguments-differ
534 """Exchange padded expert-token chunks and retain counts for backward."""
535 ctx.send_counts = send_counts
536 ctx.recv_counts = recv_counts
537 ctx.group = group
538 local_max = max([*send_counts, *recv_counts, 1])
539 pad_to = x.new_tensor([local_max], dtype=torch.int64)
540 dist.all_reduce(pad_to, op=dist.ReduceOp.MAX, group=group)
541 ctx.pad_to = pad_to = int(pad_to.item())
542 recv = _EPAllToAllPadded._pad_and_exchange(x, send_counts, pad_to, group)
543 return _EPAllToAllPadded._unpad(recv, recv_counts, pad_to)
545 @staticmethod
546 def backward(
547 ctx: Any,
548 grad_output: torch.Tensor,
549 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
550 """Reverse the exchange: pad by recv_counts, a2a_single, unpad by send_counts."""
551 # backward = reversed a2a: pad by recv_counts -> a2a_single -> unpad by send_counts
552 recv = _EPAllToAllPadded._pad_and_exchange(
553 grad_output.contiguous(), ctx.recv_counts, ctx.pad_to, ctx.group)
554 grad = _EPAllToAllPadded._unpad(recv, ctx.send_counts, ctx.pad_to)
555 return grad, None, None, None
558def _equal_a2a_rows(
559 x: torch.Tensor,
560 send_counts: list[int],
561 recv_counts: list[int],
562 group: Any,
563) -> Optional[int]:
564 """Rows per peer when the exchange may take the equal-length path, else None.
566 The equal-length path is only equivalent to the ragged one when the plan is
567 uniform *and* the payload really holds one such chunk per peer: the split
568 sizes then carry no information, and the plain all-to-all -- which derives
569 its chunking from the tensor shape alone -- moves exactly the rows the
570 ragged call would. Anything else (unequal counts, a payload whose row count
571 disagrees with the counts, a non-contiguous payload, the knob off, a backend
572 without a ragged all-to-all) yields ``None`` so the caller keeps the path it
573 runs today.
575 Args:
576 x: Payload to exchange, split along dim 0.
577 send_counts: Rows sent to each EP rank.
578 recv_counts: Rows received from each EP rank.
579 group: EP process group.
581 Returns:
582 The uniform per-peer row count, or ``None`` when the fast path does not
583 apply.
584 """
585 if _equal_a2a_mode() == _EQUAL_A2A_OFF or not _backend_supports_uneven_a2a(group):
586 return None
587 ep_size = len(send_counts)
588 if ep_size == 0 or len(recv_counts) != ep_size or not x.is_contiguous():
589 return None
590 rows_per_peer = send_counts[0]
591 # The split-free call takes the chunk size from the tensor shape and needs
592 # the receive buffer to be the input's size, so a uniform plan is only
593 # equivalent when both count lists hold that same count.
594 uniform = (all(count == rows_per_peer for count in send_counts)
595 and all(count == rows_per_peer for count in recv_counts))
596 if not uniform or x.shape[0] != rows_per_peer * ep_size:
597 return None
598 return rows_per_peer
601def ep_all_to_all(
602 x: torch.Tensor,
603 send_counts: list[int],
604 recv_counts: list[int],
605 group: Any,
606) -> torch.Tensor:
607 """Unified entry for EP token exchange (autograd-differentiable).
609 send_counts/recv_counts: list[int], length ep_size, row counts per dest/src rank.
610 NCCL/HCCL -> ragged a2a (zero-padding); other backends (gloo test path) -> pad-to-max.
611 With ``HP_EP_EQUAL_A2A`` on, a uniform plan is exchanged with the equal-length
612 a2a instead (same rows, same order, cheaper kernel). This entry always
613 returns a materialized tensor in eager mode; graph capture records the
614 functional exchange and its consumer wait. :func:`ep_all_to_all_async`
615 provides the eager lazy entry.
616 """
617 if _is_tracing_or_compiling():
618 # Keep the autograd wrapper; its first consumer inserts the stream wait.
619 return _functional_collectives.all_to_all_single_autograd(x, recv_counts, send_counts, group)
620 if not _backend_supports_uneven_a2a(group):
621 return _EPAllToAllPadded.apply(x, send_counts, recv_counts, group)
622 rows_per_peer = _equal_a2a_rows(x, send_counts, recv_counts, group)
623 if rows_per_peer is not None:
624 return _EPAllToAllEqual.apply(x, rows_per_peer, len(send_counts), group)
625 return _EPAllToAllUneven.apply(x, send_counts, recv_counts, group)
628def ep_all_to_all_async(
629 x: torch.Tensor,
630 send_counts: list[int],
631 recv_counts: list[int],
632 group: Any,
633 *,
634 allow_pending: bool = True,
635) -> Any:
636 """Non-blocking variant of :func:`ep_all_to_all` (lazy wait).
638 On backends whose ``all_to_all_single`` accepts unequal splits the exchange
639 is issued through ``differentiable_all_to_all_single_async``, which returns
640 an ``AsyncCollectiveTensor``: the ``wait_tensor`` op is only enqueued when a
641 non-view op first consumes the result. Issuing independent work between the
642 exchange and that first read therefore overlaps with the in-flight transfer
643 (forward and backward alike).
645 Backends without that support fall back to the blocking path, so the result
646 always carries the same values — only the schedule differs.
648 With ``HP_EP_EQUAL_A2A=1`` and a uniform plan the exchange is the split-free
649 equal-length one instead -- the cheaper backend kernel -- and it stays lazy:
650 :class:`_EPAllToAllEqualLazy` issues it on a stream of its own, so the two
651 properties are no longer a trade. The result is then a
652 :class:`_PendingEqualA2A` *handle* rather than a tensor, and every consumer
653 must read it through :func:`wait_ep_all_to_all`; that call is where the
654 consumer's stream is ordered on the transfer, which is what the caller wants
655 to place after its overlapped work. With ``=eager`` the same split-free
656 exchange is issued on the caller's stream and a plain tensor comes back --
657 the behaviour ``=1`` used to have, kept for measuring the kernel's gain and
658 the lost overlap apart.
660 Args:
661 x: Input tensor, split along dim 0 by ``send_counts``.
662 send_counts: Rows sent to each EP rank.
663 recv_counts: Rows received from each EP rank.
664 group: EP process group.
665 allow_pending: Whether a uniform plan may return a pending handle. A
666 caller whose next step is a view chain over the exchanged buffer
667 (the fused states+indices dispatch) passes ``False`` and gets the
668 eager split-free tensor instead.
670 Returns:
671 The exchanged rows, materialized lazily on the async path.
672 """
673 if not _backend_supports_uneven_a2a(group):
674 return _EPAllToAllPadded.apply(x, send_counts, recv_counts, group)
675 rows_per_peer = _equal_a2a_rows(x, send_counts, recv_counts, group)
676 if rows_per_peer is not None:
677 ep_size = len(send_counts)
678 if _equal_a2a_mode() == _EQUAL_A2A_LAZY and allow_pending:
679 tensor, event = _EPAllToAllEqualLazy.apply(x, rows_per_peer, ep_size, group)
680 return _PendingEqualA2A(tensor, event)
681 return _EPAllToAllEqual.apply(x, rows_per_peer, ep_size, group)
682 return differentiable_all_to_all_single_async(
683 x, send_counts, recv_counts, group,
684 )