[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.
This commit is contained in:
shaoxiongduan
2026-09-01 09:52:02 +00:00
parent a159b63c67
commit 3eabb7b40b
@@ -6,6 +6,8 @@ load-store accessible NVLink mesh: same layout, byte-identical results, fewer
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 @@ class UlyssesA2AHelper:
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 @@ class UlyssesA2AHelper:
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 @@ class UlyssesA2AHelper:
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 @@ class UlyssesA2AHelper:
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()