ulysses_a2a ¶

def close(self) -> bool:
    """Collectively destroy the device communicator and its window.

    Returns whether all ranks completed teardown. An armed/unarmed split
    cannot safely enter NCCL window deregistration, so that exceptional
    state is leaked until process exit and permanently disabled instead of
    risking a distributed deadlock.
    """
    handle = self._handle
    all_armed = self._agree(handle is not None)
    all_unarmed = self._agree(handle is None)
    if all_unarmed:
        self._nbytes = 0
        return True
    if not all_armed:
        self._handle = None
        self._nbytes = 0
        self._disable("ranks disagreed on whether a fused window was armed during teardown")
        return False

    assert handle is not None
    synchronize_ok = True
    try:
        torch.cuda.synchronize(self.device)
    except Exception:  # noqa: BLE001 - converted to a group verdict below
        synchronize_ok = False
        logger.warning("Ulysses pre-teardown synchronization failed", exc_info=True)
    if not self._agree(synchronize_ok):
        self._disable("a peer rank could not synchronize before fused-window teardown")
        return False

    dispose_ok = True
    try:
        self._dispose(handle, synchronize=False)
    except Exception:  # noqa: BLE001 - teardown must not mask a real error
        dispose_ok = False
        logger.warning("Ulysses window deregistration failed", exc_info=True)

    group_ok = self._agree(dispose_ok)
    # The native disposer consumes the handle even when a cleanup call
    # reports an error, so never retry a potentially dangling pointer.
    self._handle = None
    self._nbytes = 0
    if not group_ok:
        self._disable("fused-window teardown failed on a peer rank")
    return group_ok

def run_armed(self, x: torch.Tensor, mode: int, chunked: bool, blocks: int) -> torch.Tensor:
    """Run one call's collective(s) on an already-armed context."""
    assert self._handle is not None, "run_armed called on an unarmed helper"
    from fastvideo_kernel import comm_ops

    w = self.world_size
    if mode == 0:
        B, S_local, H, D = x.shape
        out = torch.empty(B, S_local * w, H // w, D, dtype=x.dtype, device=x.device)
    else:
        B, S_global, H_local, D = x.shape
        S_local, H = S_global // w, H_local * w
        out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device)
    if chunked:
        # Each copy completes on this stream before the next plane reuses
        # the registered window. The full result owns its storage; saved
        # activations never alias the reusable window.
        for plane in range(B):
            comm_ops.all_to_all(self._handle,
                                x[plane:plane + 1],
                                out[plane:plane + 1],
                                1,
                                S_local,
                                H,
                                D,
                                mode,
                                blocks=blocks)
    else:
        comm_ops.all_to_all(self._handle, x, out, B, S_local, H, D, mode, blocks=blocks)
    return out

def try_all_to_all_4D(self, x: torch.Tensor, scatter_dim: int, gather_dim: int) -> torch.Tensor | None:
    """Fused collective, or None to let the caller use the NCCL path."""
    if self._disabled_reason is not None:
        return None

    # Python lifecycle checks, votes, and pybind calls are not valid inside
    # a fullgraph region. The inherited NCCL path is compiler-visible, so
    # regional compile stays fullgraph by declining before any tensor read.
    if torch.compiler.is_compiling():
        return None

    signature, reason = self._call_signature(x, scatter_dim, gather_dim)
    use_fused, permanently_unavailable, lifecycle_consistent = self._agree_call(signature)
    if not use_fused:
        if not lifecycle_consistent:
            self.close()
            self._disable("ranks disagreed on the fused-window lifecycle")
        if permanently_unavailable:
            self._disable(reason or "a peer rank cannot use the fused path")
        return None

    mode = signature[2]
    nbytes = signature[8]
    if self._handle is None:
        if not self._build(nbytes):
            return None
    elif nbytes > self._nbytes:
        logger.info("Ulysses window grow: %d -> %d bytes", self._nbytes, nbytes)
        if not self.close():
            return None
        if not self._build(nbytes):
            return None

    return _FusedUlyssesA2A.apply(self, x, mode, bool(signature[10]), signature[11])