From 3eabb7b40b2c009727d15093f760da19cb96fa02 Mon Sep 17 00:00:00 2001 From: shaoxiongduan Date: Tue, 1 Sep 2026 07:52:24 +0000 Subject: [PATCH] [bugfix] gate the fused Ulysses a2a on one host, and stop re-voting per call Two defects measured on 4x GB200, MiniMax-H3 geometry, per attention layer (NCCL baseline 2421us at sp=4, 1295us at sp=8 across two trays): as shipped sp=4 2035us (1.19x) sp=8 3102us (2.4x SLOWER than NCCL) with these sp=4 1590us (1.52x) sp=8 declines, 1300us ncclTeamLsa answers "addressable", not "fast". NCCL 2.29 extends the LSA team across a multi-node NVLink domain, so on a GB200 rack the gate passes for ranks on different trays and the kernel arms. Its fine-grained 16B remote stores are far slower there than NCCL's bulk transfers. Require a single host, which is the regime the slab decomposition was tuned for, and which matches flashinfer's own gate. torch 2.12 (the pin) bundles NCCL 2.29.7, so this is reachable today. _can_attempt now computes its local verdict without collectives and then runs one unconditional all_gather_object carrying (hostname, local_ok). A collective behind a rank-local early return hangs the group whenever ranks disagree -- exactly the case this gate exists to detect. Verified: with one rank reporting the kernel unavailable, the earlier ordering timed out at 180s while this completes with correct results on every rank. The per-call agreement cost a flat ~227us regardless of operand size -- a host-side gloo all_gather before every collective. The signature is architectural: two distinct values (scatter and gather shapes) across 50 layers x 4 steps. Cache the verdict so the collective runs twice per generation instead of ~400 times. The cache trades one property: a rank whose signature diverges mid-run now misses the cache and calls the collective alone, hanging rather than falling back. Re-voting every N calls would bound that if wanted. --- .../device_communicators/ulysses_a2a.py | 45 ++++++++++++++++--- 1 file changed, 39 insertions(+), 6 deletions(-) diff --git a/fastvideo/distributed/device_communicators/ulysses_a2a.py b/fastvideo/distributed/device_communicators/ulysses_a2a.py index 2e41a2a1c1..eec55c7d42 100644 --- a/fastvideo/distributed/device_communicators/ulysses_a2a.py +++ b/fastvideo/distributed/device_communicators/ulysses_a2a.py @@ -6,6 +6,8 @@ passes over local memory. Anything else falls back to the NCCL path. """ +import socket + import torch import torch.distributed as dist from torch.distributed import ProcessGroup @@ -76,6 +78,9 @@ def __init__(self, cpu_group: ProcessGroup, device_group: ProcessGroup, world_si self.pynccl_comm = pynccl_comm self._handle: int | None = None + # One signature per (shape, dtype, mode), so the control collective + # runs twice per generation instead of once per call. + self._verdicts: dict[tuple[int, ...], tuple[bool, bool, bool]] = {} self._nbytes = 0 self._disabled_reason: str | None = None @@ -95,15 +100,37 @@ def _comm_ptr(self) -> int: return int(getattr(comm, "value", comm)) def _can_attempt(self) -> tuple[bool, str]: - """Whether this rank could use the fused path, without allocating anything.""" + """Whether this rank could use the fused path, without allocating anything. + + The exchange is unconditional: a collective behind a rank-local early + return hangs the group exactly when ranks disagree. + """ + # LSA means addressable, not fast: NCCL 2.29 spans trays on a GB200 + # rack, where the kernel's 16B remote stores lose to NCCL. + local_ok = True + local_reason = "" try: from fastvideo_kernel import comm_ops if not comm_ops.is_available(): - return False, "fastvideo-kernel was built without the Ulysses a2a kernel" - if not comm_ops.lsa_covers_group(self._comm_ptr(), self.world_size): - return False, "the group is not a load-store-accessible (NVLink) mesh" + local_ok, local_reason = False, "fastvideo-kernel was built without the Ulysses a2a kernel" + elif not comm_ops.lsa_covers_group(self._comm_ptr(), self.world_size): + local_ok, local_reason = False, "the group is not a load-store-accessible (NVLink) mesh" except Exception as e: # noqa: BLE001 - return False, f"backend unavailable ({type(e).__name__}: {e})" + local_ok, local_reason = False, f"backend unavailable ({type(e).__name__}: {e})" + + try: + gathered: list[tuple[str, bool]] = [("", False)] * self.world_size + dist.all_gather_object(gathered, (socket.gethostname(), local_ok), group=self.cpu_group) + except Exception as e: # noqa: BLE001 + return False, f"topology exchange failed ({type(e).__name__}: {e})" + + hostnames = {host for host, _ in gathered} + if len(hostnames) > 1: + return False, f"ranks span multiple hosts: {sorted(hostnames)}" + if not local_ok: + return False, local_reason + if not all(ok for _, ok in gathered): + return False, "a peer rank cannot use the fused path" return True, "" def _agree(self, ok: bool) -> bool: @@ -286,6 +313,7 @@ def close(self) -> bool: state is leaked until process exit and permanently disabled instead of risking a distributed deadlock. """ + self._verdicts.clear() handle = self._handle all_armed = self._agree(handle is not None) all_unarmed = self._agree(handle is None) @@ -355,7 +383,12 @@ def try_all_to_all_4D(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) return None signature, reason = self._call_signature(x, scatter_dim, gather_dim) - use_fused, permanently_unavailable, lifecycle_consistent = self._agree_call(signature) + cached = self._verdicts.get(signature) + if cached is not None: + use_fused, permanently_unavailable, lifecycle_consistent = cached + else: + use_fused, permanently_unavailable, lifecycle_consistent = self._agree_call(signature) + self._verdicts[signature] = (use_fused, permanently_unavailable, lifecycle_consistent) if not use_fused: if not lifecycle_consistent: self.close()