Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / expert_parallel / collectives.py: 96%
174 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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
45from hyper_parallel.core.dtensor._utils import (
46 differentiable_all_to_all_single_async,
47 get_device_handle,
48)
50_UNEVEN_A2A_BACKENDS = ("nccl", "hccl")
52# A ragged exchange reaches NCCL/HCCL as alltoallv, which takes a per-peer count
53# vector and runs a count-driven kernel. Under a balanced routing plan every
54# per-peer count is the same, so the counts add nothing and the plain
55# equal-length all_to_all_single (alltoall) can be issued instead. Opt-in while
56# it is being measured (HP_EP_EQUAL_A2A).
57_EQUAL_A2A_OFF = "off"
58_EQUAL_A2A_LAZY = "lazy"
59_EQUAL_A2A_EAGER = "eager"
60_EQUAL_A2A_MODES = {
61 "0": _EQUAL_A2A_OFF,
62 "1": _EQUAL_A2A_LAZY,
63 "eager": _EQUAL_A2A_EAGER,
64}
65_EQUAL_A2A_RAW = os.environ.get("HP_EP_EQUAL_A2A", "0")
67# Ragged (uneven) exchange kernel: the list form issues one
68# ``dist.all_to_all`` with per-peer tensors, which HCCL serves through
69# alltoallv; ``HP_EP_A2A_SINGLE=1`` issues the identical exchange as one
70# ``all_to_all_single`` with split sizes, the form the micro benchmark
71# measures at ~2x the bandwidth (106.4 vs 52.9 GB/s). Default off. Scope: this
72# only affects THIS synchronous entry; the training MoE path calls
73# ``ep_all_to_all_async``, whose ``_AsyncA2ALazyBwd`` already issues the split-sized
74# ``all_to_all_single``, so the knob is a no-op there (measured: neutral).
75_A2A_SINGLE = os.environ.get("HP_EP_A2A_SINGLE", "0") == "1"
77# Lazily created, then reused for every lazy exchange (see
78# :func:`_lazy_a2a_resources`): the stream the split-free collective is issued
79# on, the event that hands it the payload, and the event that hands the result
80# back to the consumer.
81_LAZY_A2A_STREAM = None
82_LAZY_A2A_READY_EVENT = None
83_LAZY_A2A_DONE_EVENT = None
84_LAZY_A2A_LOCK = threading.Lock()
87def _backend_supports_uneven_a2a(group) -> bool:
88 return dist.get_backend(group) in _UNEVEN_A2A_BACKENDS
91def _equal_a2a_mode() -> str:
92 """Resolve ``HP_EP_EQUAL_A2A`` to the split-free exchange's mode.
94 ``"0"`` (the default) keeps the ragged exchange, ``"1"`` takes the
95 split-free exchange and defers its wait, and ``"eager"`` takes the
96 split-free exchange while ordering the caller's stream on it at issue time
97 -- the behaviour ``"1"`` used to have, kept so the kernel's gain and the
98 lost overlap can be measured apart.
100 Returns:
101 One of ``"off"``, ``"lazy"``, ``"eager"``.
103 Raises:
104 ValueError: If the knob is none of the accepted values; a typo must fail
105 loudly instead of silently leaving the fast path off.
106 """
107 raw = _EQUAL_A2A_RAW.strip()
108 try:
109 return _EQUAL_A2A_MODES[raw]
110 except KeyError:
111 raise ValueError(
112 f"HP_EP_EQUAL_A2A must be one of "
113 f"{sorted(_EQUAL_A2A_MODES)}, but got {raw!r}"
114 ) from None
117class _EPAllToAllUneven(torch.autograd.Function): # pylint: disable=abstract-method
118 """Ragged all_to_all (NCCL/HCCL production path): split by send/recv counts.
120 forward: split(x, send_counts) -> dist.all_to_all(out_list, in_list) -> cat
121 backward: swap send/recv counts and run the ragged all_to_all again
122 (a2a is self-inverse).
123 """
125 @staticmethod
126 def forward(
127 ctx: Any,
128 x: torch.Tensor,
129 send_counts: list[int],
130 recv_counts: list[int],
131 group: Any,
132 ) -> torch.Tensor: # pylint: disable=arguments-differ
133 """Run the ragged all_to_all and retain the counts for backward."""
134 ctx.send_counts = send_counts
135 ctx.recv_counts = recv_counts
136 ctx.group = group
137 out = x.new_empty((sum(recv_counts),) + tuple(x.shape[1:]))
138 if _A2A_SINGLE:
139 dist.all_to_all_single(
140 out,
141 x.contiguous(),
142 output_split_sizes=[int(count) for count in recv_counts],
143 input_split_sizes=[int(count) for count in send_counts],
144 group=group,
145 )
146 else:
147 dist.all_to_all(list(out.split(recv_counts)),
148 list(x.split(send_counts)), group=group)
149 return out
151 @staticmethod
152 def backward(
153 ctx: Any,
154 grad_output: torch.Tensor,
155 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
156 """Swap send/recv counts and re-run the self-inverse ragged all_to_all."""
157 grad = _EPAllToAllUneven.apply(
158 grad_output.contiguous(), ctx.recv_counts, ctx.send_counts, ctx.group)
159 return grad, None, None, None
162def _equal_a2a_exchange(
163 x: torch.Tensor,
164 rows_per_peer: int,
165 ep_size: int,
166 group: Any,
167) -> torch.Tensor:
168 """Move ``rows_per_peer`` rows to each peer and return the peer-major result.
170 The split-free call: the backend derives the chunking from the tensor shape
171 alone, and the received chunks are concatenated in peer order. Both the
172 eager and the lazy equal-length exchange issue this, so a change here cannot
173 make the two disagree about the layout.
174 """
175 out = x.new_empty((rows_per_peer * ep_size,) + tuple(x.shape[1:]))
176 dist.all_to_all_single(out, x, group=group)
177 return out
180class _EPAllToAllEqual(torch.autograd.Function): # pylint: disable=abstract-method
181 """Equal-length ``all_to_all_single`` (NCCL/HCCL path when counts are uniform).
183 Without split sizes every peer gets ``rows_per_peer`` rows -- the backend
184 derives the chunking from the tensor shape alone -- and the received chunks
185 are concatenated in peer order. That is exactly the layout
186 :class:`_EPAllToAllUneven` produces when all counts are equal, so this path
187 is interchangeable with it (see :func:`_equal_a2a_rows` for when it applies).
189 forward: all_to_all_single(no splits) -> [ep_size * rows_per_peer, ...]
190 backward: the same exchange; an equal-length a2a is its own inverse, and
191 with a uniform plan the reverse chunks are the same size.
193 ``forward`` issues the exchange on the caller's stream, which orders that
194 stream on the transfer before it returns; :class:`_EPAllToAllEqualLazy` is
195 the same exchange with the ordering left to the caller.
196 """
198 @staticmethod
199 def forward(
200 ctx: Any,
201 x: torch.Tensor,
202 rows_per_peer: int,
203 ep_size: int,
204 group: Any,
205 ) -> torch.Tensor: # pylint: disable=arguments-differ
206 """Run the equal-length all_to_all and retain its geometry for backward.
208 Args:
209 ctx: Autograd context of this exchange.
210 x: Payload to exchange; it holds ``ep_size * rows_per_peer`` rows.
211 rows_per_peer: Rows handed to each peer.
212 ep_size: Number of ranks in ``group``.
213 group: Process group to exchange over.
215 Returns:
216 The received rows, concatenated in peer order.
217 """
218 ctx.rows_per_peer = rows_per_peer
219 ctx.ep_size = ep_size
220 ctx.group = group
221 return _equal_a2a_exchange(x, rows_per_peer, ep_size, group)
223 @staticmethod
224 def backward(
225 ctx: Any,
226 grad_output: torch.Tensor,
227 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
228 """Re-run the self-inverse equal-length exchange on the output gradient.
230 Args:
231 ctx: Autograd context of the forward exchange.
232 grad_output: Output gradient, one row per exchanged row.
234 Returns:
235 The input gradient, followed by ``None`` for the non-tensor arguments.
236 """
237 grad = _equal_a2a_exchange(
238 grad_output.contiguous(), ctx.rows_per_peer, ctx.ep_size, ctx.group)
239 return grad, None, None, None
242def _lazy_a2a_resources() -> tuple[Any, Any, Any]:
243 """Return the ``(stream, ready_event, done_event)`` the lazy exchange reuses.
245 One stream and two events for the whole process, which is deliberate: a
246 fresh event per exchange drains the runtime's event pool after a few hundred
247 exchanges (the run then dies in ``AclQueryEventRecordedStatus``), and one
248 stream is enough because every record and wait below is enqueued in host
249 order -- a wait therefore never observes a *stale* record, at worst a later
250 one, which only makes it wait longer than it strictly had to.
252 Returns:
253 The shared comm stream, the event that hands the payload over to it, and
254 the event that hands the result back.
255 """
256 global _LAZY_A2A_STREAM, _LAZY_A2A_READY_EVENT, _LAZY_A2A_DONE_EVENT # pylint: disable=global-statement
257 with _LAZY_A2A_LOCK:
258 if _LAZY_A2A_DONE_EVENT is None:
259 _LAZY_A2A_STREAM = get_device_handle().Stream()
260 _LAZY_A2A_READY_EVENT = get_device_handle().Event()
261 _LAZY_A2A_DONE_EVENT = get_device_handle().Event()
262 return _LAZY_A2A_STREAM, _LAZY_A2A_READY_EVENT, _LAZY_A2A_DONE_EVENT
265def _record_stream(tensor: torch.Tensor, stream: Any) -> None:
266 """Tell the caching allocator that ``stream`` touches ``tensor``'s storage.
268 Wrapped rather than called inline so the exchange can be driven with a fake
269 stream in unit tests: ``Tensor.record_stream`` rejects anything that is not a
270 real device stream.
271 """
272 tensor.record_stream(stream)
275def _issue_split_free_a2a(out: torch.Tensor, x: torch.Tensor, group: Any) -> Any:
276 """Issue the split-free exchange on the shared comm stream.
278 ``dist.all_to_all_single`` without split sizes is the one call shape whose
279 backend kernel is the plain alltoall, and it is a *blocking* collective:
280 c10d orders the stream it is called on after the transfer before returning,
281 which is what would close the caller's overlap window. Issuing it inside a
282 stream context moves that ordering onto the comm stream. c10d only orders
283 the collective against the stream it is called on, so the payload's
284 producers are re-connected to that stream explicitly through the ready
285 event.
287 Args:
288 out: Receive buffer, ``[ep_size * rows_per_peer, ...]``; written here.
289 x: Contiguous payload, ``[ep_size * rows_per_peer, ...]``; read here.
290 group: Process group to exchange over.
292 Returns:
293 The event recorded when the exchange has completed; the consumer's
294 stream waits on it (see :class:`_PendingEqualA2A`).
295 """
296 compute_stream = get_device_handle().current_stream()
297 stream, ready_event, done_event = _lazy_a2a_resources()
298 stream_context = get_device_handle().stream
299 ready_event.record(compute_stream)
300 with stream_context(stream):
301 ready_event.wait(stream)
302 dist.all_to_all_single(out, x, group=group)
303 # Both buffers outlive the call by different means: the receive buffer
304 # was allocated on the compute stream, and the payload is a temporary of
305 # the caller's dispatch, so its storage may be freed and reused while
306 # this stream is still reading it.
307 _record_stream(out, stream)
308 _record_stream(x, stream)
309 done_event.record(stream)
310 return done_event
313class _PendingEqualA2A:
314 """Handle to a split-free exchange that is still in flight.
316 ``wait()`` is the only way to read the result, and *where* it is called is
317 the caller's decision: the exchange is issued while the dispatch runs, but
318 the consumer's stream is only ordered on it at ``wait()``, so independent
319 work issued in between (the shared-expert MLP an ``overlap_fn`` runs) works
320 against the transfer instead of after it. Reading the exchange without
321 waiting is not an option -- the handle is not a tensor.
323 Attributes:
324 completed: Whether a stream has already been ordered on the exchange.
325 """
327 __slots__ = ("completed", "_event", "_tensor")
329 def __init__(self, tensor: torch.Tensor, event: Any, completed: bool = False) -> None:
330 """Wrap an issued-but-unwaited exchange result."""
331 self._tensor = tensor
332 self._event = event
333 self.completed = completed
335 def wait(self) -> torch.Tensor:
336 """Order the current stream after the exchange and return its result."""
337 if not self.completed:
338 self._event.wait(get_device_handle().current_stream())
339 self.completed = True
340 return self._tensor
342 def squeeze(self, dim: int) -> "_PendingEqualA2A":
343 """Return a handle to the squeezed result (a view, so the wait is pending).
345 The dispatched expert indices arrive as ``[rows, 1]`` and their consumers
346 want ``[rows]``. Squeezing is metadata-only, so it can be taken here
347 without materializing: a read here would enqueue the wait before the
348 caller reaches its independent work.
350 Args:
351 dim: Dimension to squeeze, as ``Tensor.squeeze`` takes it.
353 Returns:
354 A handle to the squeezed result; it still holds the exchange's result
355 as pending as this one does.
356 """
357 return _PendingEqualA2A(self._tensor.squeeze(dim), self._event, self.completed)
359 def __repr__(self) -> str:
360 """Name the pending exchange and what is known about it without reading it."""
361 return (f"PendingEqualA2A(shape={tuple(self._tensor.shape)}, "
362 f"completed={self.completed})")
365def wait_ep_all_to_all(value: Any) -> Any:
366 """Materialize a pending EP exchange.
368 Every consumer of an :func:`ep_all_to_all_async` result must pass it through
369 here before reading it: with ``HP_EP_EQUAL_A2A=1`` and a uniform plan that
370 entry returns a :class:`_PendingEqualA2A` instead of a tensor.
372 Args:
373 value: The pending handle, an already-materialized tensor from any other
374 path of :func:`ep_all_to_all_async`, or ``None``.
376 Returns:
377 The exchange result as a regular tensor; ``value`` unchanged when it is
378 not a pending handle.
379 """
380 if isinstance(value, _PendingEqualA2A):
381 return value.wait()
382 return value
385class _EPAllToAllEqualLazy(torch.autograd.Function): # pylint: disable=abstract-method
386 """Split-free equal-length exchange issued on the comm stream, waited lazily.
388 Same exchange as :class:`_EPAllToAllEqual` and just as differentiable, but
389 the forward leaves the caller's stream alone: it issues on the shared comm
390 stream and returns the buffer together with the event that says when the
391 buffer is final, so the caller can synchronize when it needs the values
392 rather than when it issues the exchange. Both travel out of ``forward`` --
393 the buffer as the autograd output, the event as a second, non-tensor output.
395 The backward is the plain split-free exchange on the caller's stream: the
396 engine consumes its result immediately (it accumulates it onto the input's
397 ``grad``), so there is nothing to defer.
398 """
400 @staticmethod
401 def forward( # pylint: disable=arguments-differ
402 ctx: Any,
403 x: torch.Tensor,
404 rows_per_peer: int,
405 ep_size: int,
406 group: Any,
407 ) -> tuple[torch.Tensor, Any]:
408 """Issue the split-free exchange and return ``(buffer, done event)``.
410 Args:
411 ctx: Autograd context of this exchange.
412 x: Contiguous payload, ``[ep_size * rows_per_peer, ...]``.
413 rows_per_peer: Rows handed to each peer.
414 ep_size: Number of ranks in ``group``.
415 group: Process group to exchange over.
417 Returns:
418 The receive buffer the exchange writes -- the output the autograd
419 graph is built on -- and the event that says when it is final.
420 """
421 ctx.rows_per_peer = rows_per_peer
422 ctx.ep_size = ep_size
423 ctx.group = group
424 out = x.new_empty((rows_per_peer * ep_size,) + tuple(x.shape[1:]))
425 return out, _issue_split_free_a2a(out, x, group)
427 @staticmethod
428 def backward( # pylint: disable=arguments-differ
429 ctx: Any,
430 grad_output: torch.Tensor,
431 _event_grad: Any, # pylint: disable=invalid-name
432 ) -> tuple[torch.Tensor, None, None, None]:
433 """Re-run the self-inverse exchange, on the caller's stream this time.
435 Args:
436 ctx: Autograd context of the forward exchange.
437 grad_output: Output gradient, one row per exchanged row.
438 _event_grad: Gradient slot of the forward's event output; unused.
440 Returns:
441 The input gradient, followed by ``None`` for the non-tensor arguments.
442 """
443 grad = _equal_a2a_exchange(
444 grad_output.contiguous(), ctx.rows_per_peer, ctx.ep_size, ctx.group)
445 return grad, None, None, None
448class _EPAllToAllPadded(torch.autograd.Function): # pylint: disable=abstract-method
449 """pad-to-max + all_to_all_single (gloo test path).
451 forward: pad each dest chunk to the global max(counts) (a2a_single
452 requires equal-length chunks per rank -> pad_to must be
453 globally consistent, obtained via all_reduce MAX)
454 -> a2a_single -> unpad by recv_counts;
455 backward: pad by recv_counts -> a2a_single (equal-length self-inverse)
456 -> unpad by send_counts.
457 """
459 @staticmethod
460 def _pad_and_exchange(x, counts, pad_to, group):
461 """Pad each chunk to pad_to, run equal-length a2a_single, return [ep*pad_to, ...]."""
462 chunks = []
463 for chunk, n in zip(x.split(counts), counts):
464 if n < pad_to:
465 pad = x.new_zeros((pad_to - n,) + tuple(x.shape[1:]))
466 chunk = torch.cat([chunk, pad])
467 chunks.append(chunk)
468 send = torch.cat(chunks).contiguous()
469 recv = torch.empty_like(send)
470 dist.all_to_all_single(recv, send, group=group)
471 return recv
473 @staticmethod
474 def _unpad(recv, counts, pad_to):
475 """Take the valid rows from the equal-length buffer per counts and cat."""
476 pieces = []
477 for i, n in enumerate(counts):
478 if n > 0:
479 pieces.append(recv[i * pad_to: i * pad_to + n])
480 if not pieces:
481 return recv.new_zeros((0,) + tuple(recv.shape[1:]))
482 return torch.cat(pieces)
484 @staticmethod
485 def forward(
486 ctx: Any,
487 x: torch.Tensor,
488 send_counts: list[int],
489 recv_counts: list[int],
490 group: Any,
491 ) -> torch.Tensor: # pylint: disable=arguments-differ
492 """Exchange padded expert-token chunks and retain counts for backward."""
493 ctx.send_counts = send_counts
494 ctx.recv_counts = recv_counts
495 ctx.group = group
496 local_max = max([*send_counts, *recv_counts, 1])
497 pad_to = x.new_tensor([local_max], dtype=torch.int64)
498 dist.all_reduce(pad_to, op=dist.ReduceOp.MAX, group=group)
499 ctx.pad_to = pad_to = int(pad_to.item())
500 recv = _EPAllToAllPadded._pad_and_exchange(x, send_counts, pad_to, group)
501 return _EPAllToAllPadded._unpad(recv, recv_counts, pad_to)
503 @staticmethod
504 def backward(
505 ctx: Any,
506 grad_output: torch.Tensor,
507 ) -> tuple[torch.Tensor, None, None, None]: # pylint: disable=arguments-differ
508 """Reverse the exchange: pad by recv_counts, a2a_single, unpad by send_counts."""
509 # backward = reversed a2a: pad by recv_counts -> a2a_single -> unpad by send_counts
510 recv = _EPAllToAllPadded._pad_and_exchange(
511 grad_output.contiguous(), ctx.recv_counts, ctx.pad_to, ctx.group)
512 grad = _EPAllToAllPadded._unpad(recv, ctx.send_counts, ctx.pad_to)
513 return grad, None, None, None
516def _equal_a2a_rows(
517 x: torch.Tensor,
518 send_counts: list[int],
519 recv_counts: list[int],
520 group: Any,
521) -> Optional[int]:
522 """Rows per peer when the exchange may take the equal-length path, else None.
524 The equal-length path is only equivalent to the ragged one when the plan is
525 uniform *and* the payload really holds one such chunk per peer: the split
526 sizes then carry no information, and the plain all-to-all -- which derives
527 its chunking from the tensor shape alone -- moves exactly the rows the
528 ragged call would. Anything else (unequal counts, a payload whose row count
529 disagrees with the counts, a non-contiguous payload, the knob off, a backend
530 without a ragged all-to-all) yields ``None`` so the caller keeps the path it
531 runs today.
533 Args:
534 x: Payload to exchange, split along dim 0.
535 send_counts: Rows sent to each EP rank.
536 recv_counts: Rows received from each EP rank.
537 group: EP process group.
539 Returns:
540 The uniform per-peer row count, or ``None`` when the fast path does not
541 apply.
542 """
543 if _equal_a2a_mode() == _EQUAL_A2A_OFF or not _backend_supports_uneven_a2a(group):
544 return None
545 ep_size = len(send_counts)
546 if ep_size == 0 or len(recv_counts) != ep_size or not x.is_contiguous():
547 return None
548 rows_per_peer = send_counts[0]
549 # The split-free call takes the chunk size from the tensor shape and needs
550 # the receive buffer to be the input's size, so a uniform plan is only
551 # equivalent when both count lists hold that same count.
552 uniform = (all(count == rows_per_peer for count in send_counts)
553 and all(count == rows_per_peer for count in recv_counts))
554 if not uniform or x.shape[0] != rows_per_peer * ep_size:
555 return None
556 return rows_per_peer
559def ep_all_to_all(
560 x: torch.Tensor,
561 send_counts: list[int],
562 recv_counts: list[int],
563 group: Any,
564) -> torch.Tensor:
565 """Unified entry for EP token exchange (autograd-differentiable).
567 send_counts/recv_counts: list[int], length ep_size, row counts per dest/src rank.
568 NCCL/HCCL -> ragged a2a (zero-padding); other backends (gloo test path) -> pad-to-max.
569 With ``HP_EP_EQUAL_A2A`` on, a uniform plan is exchanged with the equal-length
570 a2a instead (same rows, same order, cheaper kernel). This entry always
571 returns a materialized tensor; :func:`ep_all_to_all_async` is the lazy one.
572 """
573 if not _backend_supports_uneven_a2a(group):
574 return _EPAllToAllPadded.apply(x, send_counts, recv_counts, group)
575 rows_per_peer = _equal_a2a_rows(x, send_counts, recv_counts, group)
576 if rows_per_peer is not None:
577 return _EPAllToAllEqual.apply(x, rows_per_peer, len(send_counts), group)
578 return _EPAllToAllUneven.apply(x, send_counts, recv_counts, group)
581def ep_all_to_all_async(
582 x: torch.Tensor,
583 send_counts: list[int],
584 recv_counts: list[int],
585 group: Any,
586 *,
587 allow_pending: bool = True,
588) -> Any:
589 """Non-blocking variant of :func:`ep_all_to_all` (lazy wait).
591 On backends whose ``all_to_all_single`` accepts unequal splits the exchange
592 is issued through ``differentiable_all_to_all_single_async``, which returns
593 an ``AsyncCollectiveTensor``: the ``wait_tensor`` op is only enqueued when a
594 non-view op first consumes the result. Issuing independent work between the
595 exchange and that first read therefore overlaps with the in-flight transfer
596 (forward and backward alike).
598 Backends without that support fall back to the blocking path, so the result
599 always carries the same values — only the schedule differs.
601 With ``HP_EP_EQUAL_A2A=1`` and a uniform plan the exchange is the split-free
602 equal-length one instead -- the cheaper backend kernel -- and it stays lazy:
603 :class:`_EPAllToAllEqualLazy` issues it on a stream of its own, so the two
604 properties are no longer a trade. The result is then a
605 :class:`_PendingEqualA2A` *handle* rather than a tensor, and every consumer
606 must read it through :func:`wait_ep_all_to_all`; that call is where the
607 consumer's stream is ordered on the transfer, which is what the caller wants
608 to place after its overlapped work. With ``=eager`` the same split-free
609 exchange is issued on the caller's stream and a plain tensor comes back --
610 the behaviour ``=1`` used to have, kept for measuring the kernel's gain and
611 the lost overlap apart.
613 Args:
614 x: Input tensor, split along dim 0 by ``send_counts``.
615 send_counts: Rows sent to each EP rank.
616 recv_counts: Rows received from each EP rank.
617 group: EP process group.
618 allow_pending: Whether a uniform plan may return a pending handle. A
619 caller whose next step is a view chain over the exchanged buffer
620 (the fused states+indices dispatch) passes ``False`` and gets the
621 eager split-free tensor instead.
623 Returns:
624 The exchanged rows, materialized lazily on the async path.
625 """
626 if not _backend_supports_uneven_a2a(group):
627 return _EPAllToAllPadded.apply(x, send_counts, recv_counts, group)
628 rows_per_peer = _equal_a2a_rows(x, send_counts, recv_counts, group)
629 if rows_per_peer is not None:
630 ep_size = len(send_counts)
631 if _equal_a2a_mode() == _EQUAL_A2A_LAZY and allow_pending:
632 tensor, event = _EPAllToAllEqualLazy.apply(x, rows_per_peer, ep_size, group)
633 return _PendingEqualA2A(tensor, event)
634 return _EPAllToAllEqual.apply(x, rows_per_peer, ep_size, group)
635 return differentiable_all_to_all_single_async(
636 x, send_counts, recv_counts, group,
637 )