"""
Copyright (c) 2025-2026 by FlashInfer team.

The raw all-to-all entry points (merged from the former ulysses_a2a.py) wrap
a CUDA kernel adapted from ThunderKittens' NVLink all-to-all:
https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/parallel/all_to_all/all_to_all.cu

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import contextlib
import ctypes
import functools
import re
from types import SimpleNamespace
from typing import Any, List, Optional, Tuple, Union

import torch
import torch.distributed as dist
from torch.distributed import ProcessGroup

from ..api_logging import flashinfer_api
from ..jit.comm import gen_ulysses_a2a_module
from ..utils import register_custom_op
from .ulysses_topology import (
    SUPPORTED_WORLD_SIZES,
    UlyssesBackendDecision,
    resolve_ulysses_backend,
)

_INT32_MAX = 2**31 - 1
_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)

# communicator lifecycle states; CLOSED is only reached after a fully
# successful teardown so a failed close() can be retried
_OPEN, _CLOSING, _CLOSED = "open", "closing", "closed"


class UlyssesCommunicator:
    r"""Ulysses context-parallelism all-to-all communicator.

    Provides the two layout transforms of Ulysses attention over the 4-D
    layout ``[B, S, H, D]`` (a typical attention layer makes four collective
    calls: q/k/v through :meth:`scatter_heads`, the output through
    :meth:`gather_heads`):

    - :meth:`scatter_heads`: ``[B, S_local, H, D] -> [B, S_global, H_local, D]``
      (each rank keeps a head slice of the *full* sequence)
    - :meth:`gather_heads`:  ``[B, S_global, H_local, D] -> [B, S_local, H, D]``
      (each rank gets all heads of its *local* sequence shard back)

    where ``H`` is the global head count, ``H_local = H // world_size`` and
    ``S_global = S_local * world_size``. Both backends produce bit-identical
    results.

    Backend selection happens in the constructor, strictly before any IPC
    allocation or JIT compilation (see
    :func:`~flashinfer.comm.resolve_ulysses_backend`):

    - ``backend="auto"``: the fused-transpose NVLink-P2P kernel when the group
      is a verified single-node all-pairs NVLink mesh with a supported world
      size (2/4/6/8); NCCL otherwise — including when NVLink runtime
      initialization fails after a positive topology decision. Inspect
      :attr:`backend` and :attr:`fallback_reason` for the outcome.
    - ``backend="nvlink"``: force the fused kernel; raises on every rank
      (before any IPC/JIT for topology failures) when it cannot be used.
    - ``backend="nccl"``: force the ``dist.all_to_all_single`` path; skips
      the topology/NVML probe and all IPC/JIT entirely (the constructor
      still resolves and guards the CUDA device and performs CUDA-backed
      metadata collectives over ``group``). Supports any world size.

    All ranks must request the same ``backend``. The NCCL path with
    ``world_size > 1`` requires ``group`` to support CUDA all-to-all (an
    NCCL process group); this is checked at construction.
    ``world_size == 1`` is a passthrough: both collectives return the input
    tensor unchanged (no copy).

    Constraints
    -----------
    - The constructor is always collective: every rank of ``group`` must
      call it together. :meth:`close` is collective only when the NVLink
      backend was armed (its resources are IPC-shared); for the pure NCCL
      backend, ``world_size == 1``, or an auto fallback whose NVLink cleanup
      already completed, ``close`` is local and idempotent. Rank-local
      failures inside the constructor's NVLink initialization or inside a
      collective ``close`` are exchanged as group outcomes, so all ranks
      jointly clean up and raise (or fall back) instead of deadlocking; a
      failed ``close`` may be retried by all ranks.
    - Collectives run on the *current* CUDA stream of this rank; every rank
      must issue the same sequence of calls with consistently-shaped operands
      (a shape or call-order mismatch across ranks is a collective failure:
      expect hangs or garbage, exactly as with any collective library). At
      most one collective may be in flight per communicator at a time (the
      NVLink signal buffers assume serialized calls); do not call one
      communicator concurrently from multiple streams or threads.
    - Operand tensors must be contiguous 4-D CUDA tensors of the construction
      ``dtype`` (float16 / bfloat16 / float32) on the construction device,
      with every dim positive and total elements at most ``max_elems``;
      :meth:`scatter_heads` additionally requires ``H % world_size == 0`` and
      :meth:`gather_heads` requires ``S_global % world_size == 0``.
    - Each rank may use a different CUDA device (e.g. ``cuda:rank``); ranks
      must agree on ``max_elems``, ``dtype`` and ``backend``.

    Parameters
    ----------
    group : torch.distributed.ProcessGroup, optional
        Process group of the Ulysses ranks. Defaults to ``dist.group.WORLD``.
    max_elems : int
        Capacity: the largest element count of any single all-to-all operand
        (input and output have equal ``numel``, so this is ``B*S_local*H*D``
        for the largest call). Must be at most ``2**31 - 1`` (the kernel's
        int32 index range). Sizes the NVLink staging buffer once at
        construction.
    dtype : torch.dtype
        Element type of all operands (float16 / bfloat16 / float32); enforced
        on every call.
    backend : str
        ``"auto"`` | ``"nvlink"`` | ``"nccl"`` (see above).
    device : torch.device or str or int, optional
        CUDA device of this rank; normalized to an explicit index (bare
        ``"cuda"`` means the current device, an int is a CUDA ordinal).
        Defaults to the current CUDA device.

    Examples
    --------
    >>> with UlyssesCommunicator(group, max_elems=B*S*H*D, dtype=torch.bfloat16) as comm:
    ...     q_ = comm.scatter_heads(q)   # [B,S_local,H,D] -> [B,S_global,H_local,D]
    ...     ...
    ...     o = comm.gather_heads(o_)    # [B,S_global,H_local,D] -> [B,S_local,H,D]
    """

    @flashinfer_api
    def __init__(
        self,
        group: Optional[ProcessGroup] = None,
        *,
        max_elems: int,
        dtype: torch.dtype,
        backend: str = "auto",
        device: Optional[Union[torch.device, str, int]] = None,
    ):
        r"""Construct a Ulysses communicator.

        Parameters
        ----------
        group : Optional[ProcessGroup], optional
            Process group spanning the participating ranks. ``None`` uses
            ``torch.distributed.group.WORLD``.
        max_elems : int
            Per-rank upper bound on the number of elements communicated by a
            single collective call. Used to size the backend workspace.
        dtype : torch.dtype
            Element dtype for collective operands. Must be one of
            ``torch.float16``, ``torch.bfloat16``, or ``torch.float32``.
        backend : str, default = "auto"
            Backend selection policy. ``"auto"`` probes topology and prefers
            NVLink when supported, otherwise falls back to NCCL. ``"nvlink"``
            forces the NVLink backend and raises if unavailable. ``"nccl"``
            forces the NCCL path.
        device : Optional[Union[torch.device, str, int]], optional
            CUDA device bound to this rank. ``None`` uses the current CUDA
            device. Strings and integers are normalized to an explicit CUDA
            ordinal.
        """
        self._state = _CLOSED  # flipped to OPEN only when construction succeeds
        self._nvlink_armed = False  # joint property: set on all ranks or none
        # opaque NVLink-backend handle (a C++ UlyssesA2A* as an int) from
        # init_ulysses_a2a; None until armed and after teardown
        self._fa: Optional[int] = None
        self._out_ptrs: Optional[List[int]] = None
        self._sig_ptrs: Optional[List[int]] = None
        # rank-local resource tracking for staged init/teardown
        self._exports: List[int] = []  # device ptrs this rank cudaMalloc'ed
        self._imports: List[int] = []  # peer ptrs this rank IpcOpen'ed

        if group is None:
            group = dist.group.WORLD
        self.group = group
        self.rank = dist.get_rank(group=group)
        self.world_size = dist.get_world_size(group=group)

        # ---- collective-safe config validation ------------------------------
        # The bound device must be resolved BEFORE the first collective: NCCL
        # object collectives stage through a tensor on the *current* device,
        # so an explicit device="cuda:rank" without a prior set_device(rank)
        # would otherwise land every rank's metadata collective on GPU 0.
        # _resolve_device never raises; an unparsable input yields the current
        # device as a safe gather guard and the joint config validation right
        # after rejects it on every rank together.
        self.device = self._resolve_device(device)
        # Encode the local config with zero user code (exact type checks and
        # interpreter/torch-provided names only), gather, then validate the
        # identical list jointly so an invalid single-rank config raises the
        # same error on every rank instead of hanging peers in a later gather.
        # Devices are validated per rank but may legitimately differ across
        # ranks (cuda:rank); only max_elems and dtype must match.
        config = self._encode_config(max_elems, dtype, device)
        configs = self._gather(config)
        self._validate_configs_jointly(configs)

        self.max_elems = max_elems
        self.dtype = dtype

        # ---- backend selection: strictly before any IPC/JIT -----------------
        # topology_decision is what the probe concluded; decision is the
        # *effective* backend after runtime initialization (they differ only
        # when NVLink init failed at runtime and auto fell back to NCCL).
        self.topology_decision: UlyssesBackendDecision = resolve_ulysses_backend(
            backend, group=group, device=self.device
        )
        self.decision: UlyssesBackendDecision = self.topology_decision
        self.backend = self.decision.backend
        self.fallback_reason = (
            self.decision.reason
            if self.backend == "nccl" and backend != "nccl"
            else None
        )

        if self.backend == "nvlink":
            err = self._nvlink_init_transaction()
            if err is not None:
                # all ranks cleaned up (verified group-wide by the staged
                # cleanup) and hold the same joint error
                if backend == "nvlink":
                    raise RuntimeError(f"NVLink backend initialization failed: {err}")
                self.backend = "nccl"
                self.fallback_reason = f"nvlink init failed: {err}"
                self.decision = UlyssesBackendDecision("nccl", self.fallback_reason)

        # NCCL fallback needs a group that can move CUDA tensors; deterministic
        # in the (identical) group object, so a plain raise is group-uniform.
        if self.backend == "nccl" and self.world_size > 1:
            supported, observed = self._group_supports_cuda_alltoall()
            if not supported:
                raise ValueError(
                    "the Ulysses NCCL backend requires a process group "
                    f"supporting CUDA all-to-all (nccl), got '{observed}'"
                )

        self._state = _OPEN

    # ---- collective helpers ---------------------------------------------------

    def _group_supports_cuda_alltoall(self) -> Tuple[bool, str]:
        """A plain get_backend substring check would reject legitimate
        multi-backend groups (init_process_group(backend=None) reports
        "undefined" while its CUDA backend is ProcessGroupNCCL); check the
        backend actually bound to the CUDA device in that case. Deterministic
        in the group, so the resulting raise is group-uniform."""
        try:
            observed = str(dist.get_backend(self.group))
        except Exception as e:  # noqa: BLE001
            observed = f"<error: {type(e).__name__}: {e}>"
        if "nccl" in observed.lower():
            return True, observed
        try:
            cuda_backend = self.group._get_backend(torch.device("cuda"))
            if (
                cuda_backend is not None
                and "nccl" in type(cuda_backend).__name__.lower()
            ):
                return True, f"{observed} (cuda: {type(cuda_backend).__name__})"
        except Exception:  # noqa: BLE001 — no CUDA backend bound
            pass
        return False, observed

    def _gather(self, payload: Any) -> List[Any]:
        out: List[Any] = [None] * self.world_size
        # once the communicator device is resolved, metadata collectives must
        # not run on whatever the caller's current device happens to be
        device = getattr(self, "device", None)
        if device is not None:
            with torch.cuda.device(device):
                dist.all_gather_object(out, payload, group=self.group)
        else:
            dist.all_gather_object(out, payload, group=self.group)
        return out

    @staticmethod
    def _parse_cuda_ordinal(device) -> Tuple[Optional[int], Optional[str]]:
        """Strictly parse a device spec into a CUDA ordinal.

        Returns ``(index_or_None_for_current, error_or_None)``. torch.device
        wraps ordinals into a signed byte (``cuda:256`` silently becomes
        ``cuda:0``), so raw str/int ordinals are validated BEFORE any torch
        normalization; pre-built torch.device objects can only be checked for
        the surviving (possibly wrapped) index range.
        """
        count = torch.cuda.device_count()
        if device is None:
            return None, None
        if isinstance(device, bool):
            return None, f"invalid type: {type(device).__name__}"
        if isinstance(device, int):
            if 0 <= device < count:
                return device, None
            return None, f"ordinal {device} outside visible device count {count}"
        if isinstance(device, str):
            m = re.fullmatch(r"\s*cuda(?::(\d+))?\s*", device)
            if m is None:
                try:
                    parsed = torch.device(device)
                except (RuntimeError, ValueError, TypeError) as e:
                    return None, f"unparsable device: {e}"
                if parsed.type != "cuda":
                    return None, f"device must be a CUDA device, got {parsed}"
                return parsed.index, None
            if m.group(1) is None:
                return None, None  # bare "cuda" == current device
            idx = int(m.group(1))
            if 0 <= idx < count:
                return idx, None
            return None, f"ordinal {idx} outside visible device count {count}"
        if isinstance(device, torch.device):
            if device.type != "cuda":
                return None, f"device must be a CUDA device, got {device}"
            if device.index is None:
                return None, None
            if 0 <= device.index < count:
                return device.index, None
            return None, f"index {device.index} outside visible device count {count}"
        return None, f"invalid type: {type(device).__name__}"

    @classmethod
    def _resolve_device(cls, device) -> torch.device:
        """Never raises: yields the bound device for valid input and a safe
        gather-guard device (the current one) otherwise — the joint config
        validation rejects the invalid input right after."""
        try:
            index, err = cls._parse_cuda_ordinal(device)
            if err is not None:
                index = None
            if index is None:
                index = torch.cuda.current_device()
            return torch.device("cuda", index)
        except Exception:  # noqa: BLE001
            return torch.device("cuda", 0)

    @classmethod
    def _encode_config(cls, max_elems, dtype, device) -> Tuple[str, str, str]:
        if type(max_elems) is not int:  # bool is an int subclass: reject it too
            elems = f"<invalid type: {type(max_elems).__name__}>"
        else:
            elems = str(max_elems)
        if isinstance(dtype, torch.dtype):
            dt = str(dtype)
        else:
            dt = f"<invalid type: {type(dtype).__name__}>"
        try:
            index, err = cls._parse_cuda_ordinal(device)
        except Exception as e:  # noqa: BLE001
            index, err = None, f"{type(e).__name__}: {e}"
        if err is not None:
            dev = f"<invalid device: {err}>"
        elif index is None:
            dev = "cuda"
        else:
            dev = f"cuda:{index}"
        return (elems, dt, dev)

    def _validate_configs_jointly(self, configs) -> None:
        supported = tuple(str(d) for d in _SUPPORTED_DTYPES)
        problems = {}
        for r, (elems, dt, dev) in enumerate(configs):
            errs = []
            if not elems.isdigit() or int(elems) <= 0:
                errs.append(f"max_elems must be a positive int, got {elems}")
            elif int(elems) > _INT32_MAX:
                errs.append(
                    f"max_elems must be at most {_INT32_MAX} (int32 kernel "
                    f"index range), got {elems}"
                )
            if dt not in supported:
                errs.append(f"dtype must be one of {supported}, got {dt}")
            if not dev.startswith("cuda"):
                errs.append(f"device must be a CUDA device, got {dev}")
            if errs:
                problems[r] = "; ".join(errs)
        if problems:
            raise ValueError(f"invalid UlyssesCommunicator config by rank: {problems}")
        shared = {(elems, dt) for (elems, dt, _dev) in configs}
        if len(shared) > 1:
            raise ValueError(
                f"inconsistent UlyssesCommunicator configs across ranks: "
                f"(max_elems, dtype) = {sorted(shared)}; all ranks must agree"
            )

    # ---- staged NVLink initialization (collective-safe transaction) -----------
    #
    # Every stage ends with an outcome all-gather, so a rank-local failure at
    # any point (JIT compile, cudaMalloc, IPC get-handle, IPC open, kernel
    # init) is seen by all ranks together; they then run the same staged
    # cleanup (close imports -> gather -> free exports -> gather) and return
    # the same joint error. No bare barrier is ever reached by only a subset
    # of ranks.

    def _nvlink_init_transaction(self) -> Optional[str]:
        # stage J: JIT compile / load both modules and read the signal size.
        # Every import (including cudart below) lives inside a stage envelope:
        # an import failing on one rank must become a gathered outcome, not an
        # exception escaping before a gather.
        try:
            from .vllm_ar import meta_size

            with torch.cuda.device(self.device):
                get_ulysses_a2a_module()
                sig_bytes = int(meta_size())
            outcome: Tuple[str, ...] = ("ok", str(sig_bytes))
        except Exception as e:  # noqa: BLE001
            outcome = ("err", f"rank {self.rank} JIT/meta: {type(e).__name__}: {e}")
        err = self._first_error(self._gather(outcome))
        if err is not None:
            return err  # nothing allocated anywhere yet

        # stage A: allocate this rank's export buffers and IPC handles
        out_bytes = self.max_elems * self.dtype.itemsize
        handles: Optional[Tuple[Any, Any]] = None
        try:
            from .cuda_ipc import cudart

            with torch.cuda.device(self.device):
                out_ptr = cudart.cudaMalloc(out_bytes)
                self._exports.append(out_ptr.value)
                out_handle = cudart.cudaIpcGetMemHandle(out_ptr)
                sig_ptr = cudart.cudaMalloc(sig_bytes)
                self._exports.append(sig_ptr.value)
                sig_handle = cudart.cudaIpcGetMemHandle(sig_ptr)
            handles = (out_handle, sig_handle)
            outcome = ("ok",)
        except Exception as e:  # noqa: BLE001
            outcome = ("err", f"rank {self.rank} alloc: {type(e).__name__}: {e}")
        gathered = self._gather((outcome, handles))
        err = self._first_error([o for (o, _h) in gathered])
        if err is not None:
            return self._staged_cleanup(err)

        # stage B: open every peer's handles
        all_handles = [h for (_o, h) in gathered]
        out_ptrs: List[int] = [0] * self.world_size
        sig_ptrs: List[int] = [0] * self.world_size
        try:
            from .cuda_ipc import cudart

            with torch.cuda.device(self.device):
                for i, pair in enumerate(all_handles):
                    if i == self.rank:
                        out_ptrs[i] = self._exports[0]
                        sig_ptrs[i] = self._exports[1]
                        continue
                    p = cudart.cudaIpcOpenMemHandle(pair[0])
                    self._imports.append(p.value)
                    out_ptrs[i] = p.value
                    p = cudart.cudaIpcOpenMemHandle(pair[1])
                    self._imports.append(p.value)
                    sig_ptrs[i] = p.value
            outcome = ("ok",)
        except Exception as e:  # noqa: BLE001
            outcome = ("err", f"rank {self.rank} IPC open: {type(e).__name__}: {e}")
        err = self._first_error(self._gather(outcome))
        if err is not None:
            return self._staged_cleanup(err)

        # stage C: create the kernel handle (zeroes this rank's signal buffer)
        # and synchronize the bound device before reporting success — the
        # zeroing uses cudaMemset, which is asynchronous with respect to the
        # host, so neither the API returning nor the following host-side
        # gather is a CUDA completion fence on its own.
        try:
            with torch.cuda.device(self.device):
                self._fa = init_ulysses_a2a(
                    out_ptrs, sig_ptrs, self.rank, self.world_size, True
                )
                torch.cuda.synchronize()
            outcome = ("ok",)
        except Exception as e:  # noqa: BLE001
            outcome = ("err", f"rank {self.rank} init: {type(e).__name__}: {e}")
        # once every rank passes this gather, every rank's signal buffer is
        # both zeroed on-device (explicit synchronize above) and visible.
        err = self._first_error(self._gather(outcome))
        if err is not None:
            return self._staged_cleanup(err)

        self._out_ptrs = out_ptrs
        self._sig_ptrs = sig_ptrs
        self._nvlink_armed = True
        return None

    @staticmethod
    def _first_error(outcomes: List[Tuple[str, ...]]) -> Optional[str]:
        errs = [o[1] for o in outcomes if o and o[0] == "err"]
        return "; ".join(errs) if errs else None

    def _staged_cleanup(self, err: str) -> str:
        """Joint init-failure cleanup: all ranks arrive here together (they
        all saw the same failed outcome gather) and run the full teardown
        protocol. Cleanup completion is *verified* group-wide; if it cannot
        be completed the constructor fails jointly on every rank (auto is not
        allowed to fall back to NCCL while NVLink resources may linger)."""
        cleanup_err = self._teardown_protocol(sync_first=True)
        if cleanup_err is not None:
            raise RuntimeError(
                f"NVLink backend initialization failed ({err}) and cleanup "
                f"could not be completed: {cleanup_err}"
            )
        return err

    # ---- staged teardown protocol ---------------------------------------------
    #
    # Fixed stage sequence executed by EVERY rank whenever it runs, regardless
    # of how many resources the rank still holds locally (a rank with nothing
    # left still participates in every gather — otherwise a retry after a
    # partial failure deadlocks the ranks that do have work left). Each stage
    # drains with bounded retries; the retry/stop decision is taken from the
    # gathered remaining-counts, so every rank takes the same branch.

    _TEARDOWN_ATTEMPTS = 3

    def _teardown_protocol(self, *, sync_first: bool) -> Optional[str]:
        stages = []
        if sync_first:
            # collectives/memsets are async enqueues: never unmap while the
            # bound device may still be executing one
            stages.append(("synchronize device", self._try_sync))
        stages.append(("dispose kernel handle", self._try_dispose))
        stages.append(("close peer mappings", self._try_close_imports))
        # exports are freed only after the gathered remaining-import count is
        # zero on EVERY rank: freeing a buffer a peer still has mapped is
        # undefined behavior
        stages.append(("free exports", self._try_free_exports))

        for name, step in stages:
            for attempt in range(1, self._TEARDOWN_ATTEMPTS + 1):
                # broad envelope around the WHOLE step: a helper that raises
                # (module import, device-guard enter/exit, anything) must
                # become a nonzero remaining-count, never skip the gather and
                # strand the peers
                try:
                    remaining, detail = step()
                except Exception as e:  # noqa: BLE001
                    remaining = 1
                    detail = (
                        f"rank {self.rank} stage '{name}' raised: "
                        f"{type(e).__name__}: {e}"
                    )
                outcomes = self._gather((remaining, detail))
                if all(r == 0 for (r, _d) in outcomes):
                    break  # stage complete on every rank
                if attempt == self._TEARDOWN_ATTEMPTS:
                    per_rank = {r: d for r, (n, d) in enumerate(outcomes) if n > 0}
                    return f"stage '{name}' incomplete after {attempt} attempts: {per_rank}"
        return None

    def _try_sync(self) -> Tuple[int, Optional[str]]:
        try:
            with torch.cuda.device(self.device):
                torch.cuda.synchronize()
            return (0, None)
        except Exception as e:  # noqa: BLE001
            return (1, f"rank {self.rank} synchronize: {type(e).__name__}: {e}")

    def _try_dispose(self) -> Tuple[int, Optional[str]]:
        if self._fa is None:
            return (0, None)
        try:
            with torch.cuda.device(self.device):
                dispose_ulysses_a2a(self._fa)
                # ledger update inside the guard: a __exit__ raise after a
                # successful dispose must not lead to a double-delete on retry
                self._fa = None
            return (0, None)
        except Exception as e:  # noqa: BLE001
            return (
                0 if self._fa is None else 1,
                f"rank {self.rank} dispose: {type(e).__name__}: {e}",
            )

    # The release helpers update the resource ledger immediately after each
    # successful release (inside the per-pointer try, device guard included),
    # so a later failure — even a device-guard __exit__ raising — can never
    # lead to a double-close/double-free on the next attempt.

    def _try_close_imports(self) -> Tuple[int, Optional[str]]:
        from .cuda_ipc import cudart

        last = None
        for ptr in list(self._imports):
            try:
                with torch.cuda.device(self.device):
                    cudart.cudaIpcCloseMemHandle(ctypes.c_void_p(ptr))
                    # ledger update inside the guard: even a __exit__ raise
                    # after a successful close cannot cause a double-close
                    self._imports.remove(ptr)
            except Exception as e:  # noqa: BLE001 — keep for retry
                last = f"rank {self.rank} close import: {type(e).__name__}: {e}"
        return (len(self._imports), last)

    def _try_free_exports(self) -> Tuple[int, Optional[str]]:
        from .cuda_ipc import cudart

        last = None
        for ptr in list(self._exports):
            try:
                with torch.cuda.device(self.device):
                    cudart.cudaFree(ctypes.c_void_p(ptr))
                    self._exports.remove(ptr)
            except Exception as e:  # noqa: BLE001 — keep for retry
                last = f"rank {self.rank} free export: {type(e).__name__}: {e}"
        return (len(self._exports), last)

    # ---- lifecycle -----------------------------------------------------------

    def close(self) -> None:
        r"""Release the communicator. Idempotent once fully closed.

        Collective when the NVLink backend was armed: every rank must call
        ``close`` together, and every rank runs the same fixed teardown stage
        sequence even if it holds no resources locally — synchronize the
        bound device (collectives are asynchronous kernel launches; unmapping
        a peer buffer still in use would be undefined behavior), dispose the
        kernel handle, close peer mappings, and only after the group confirms
        all mappings are closed, free the exports. Each stage drains with
        bounded group-coordinated retries. If teardown still cannot complete,
        the call raises the same error on **all** ranks and the state stays
        CLOSING; every rank may retry ``close()``. The state becomes CLOSED
        only after a fully successful group-wide teardown. The pure-NCCL
        backend holds no resources and closes locally.
        """
        if self._state == _CLOSED:
            return
        self._state = _CLOSING

        if not getattr(self, "_nvlink_armed", False):
            # never held NVLink resources on ANY rank (armed is a joint
            # property: the init transaction either succeeds or cleans up on
            # every rank together), so closing locally cannot desync peers
            self._state = _CLOSED
            return

        err = self._teardown_protocol(sync_first=True)
        if err is not None:
            raise RuntimeError(
                f"UlyssesCommunicator.close failed (retry close() on all ranks): {err}"
            )
        self._out_ptrs = None
        self._sig_ptrs = None
        self._nvlink_armed = False
        self._state = _CLOSED

    def __enter__(self) -> "UlyssesCommunicator":
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        self.close()

    # ---- collectives -----------------------------------------------------------

    @flashinfer_api
    def scatter_heads(self, x: torch.Tensor) -> torch.Tensor:
        r"""``[B, S_local, H, D] -> [B, S_global, H_local, D]``.

        Scatter the global heads across ranks and gather the full sequence:
        afterwards this rank holds head slice
        ``[rank * H_local, (rank+1) * H_local)`` of every token. Runs on the
        current CUDA stream. Returns the input unchanged when
        ``world_size == 1``.

        Parameters
        ----------
        x : torch.Tensor
            Contiguous 4-D CUDA tensor with shape ``[B, S_local, H, D]``.

        Returns
        -------
        torch.Tensor
            Tensor with shape ``[B, S_global, H_local, D]`` on the same device
            and dtype as ``x``.
        """
        self._validate(x, "scatter_heads")
        B, S_local, H, D = x.shape
        if H % self.world_size != 0:
            raise ValueError(
                f"scatter_heads requires the global head count (dim 2) to be "
                f"divisible by world size {self.world_size}, got shape "
                f"{tuple(x.shape)}"
            )
        if self.world_size == 1:
            return x
        if self.backend == "nvlink":
            out = torch.empty(
                B,
                S_local * self.world_size,
                H // self.world_size,
                D,
                dtype=x.dtype,
                device=x.device,
            )
            ulysses_a2a(self._fa, x, out, B, S_local, H, D, 0)
            return out
        return self._nccl_scatter_heads(x)

    @flashinfer_api
    def gather_heads(self, x: torch.Tensor) -> torch.Tensor:
        r"""``[B, S_global, H_local, D] -> [B, S_local, H, D]``.

        Inverse of :meth:`scatter_heads`: gather all head slices for this
        rank's local sequence shard. Runs on the current CUDA stream. Returns
        the input unchanged when ``world_size == 1``.

        Parameters
        ----------
        x : torch.Tensor
            Contiguous 4-D CUDA tensor with shape ``[B, S_global, H_local, D]``.

        Returns
        -------
        torch.Tensor
            Tensor with shape ``[B, S_local, H, D]`` on the same device and
            dtype as ``x``.
        """
        self._validate(x, "gather_heads")
        B, S_global, H_local, D = x.shape
        if S_global % self.world_size != 0:
            raise ValueError(
                f"gather_heads requires the global sequence length (dim 1) to "
                f"be divisible by world size {self.world_size}, got shape "
                f"{tuple(x.shape)}"
            )
        if self.world_size == 1:
            return x
        S_local = S_global // self.world_size
        H = H_local * self.world_size
        if self.backend == "nvlink":
            out = torch.empty(B, S_local, H, D, dtype=x.dtype, device=x.device)
            ulysses_a2a(self._fa, x, out, B, S_local, H, D, 1)
            return out
        return self._nccl_gather_heads(x)

    # ---- NCCL fallback ---------------------------------------------------------
    # The conventional all_to_all_single path with explicit permute/contiguous
    # glue before and after (exactly the data movement the fused NVLink kernel
    # folds into its cross-GPU writes). Bit-identical to the NVLink backend.

    def _nccl_scatter_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, S_local, H, D = x.shape
        W = self.world_size
        H_local = H // W
        xt = x.reshape(B, S_local, W, H_local, D).permute(2, 0, 1, 3, 4).contiguous()
        recv = torch.empty_like(xt)
        dist.all_to_all_single(recv, xt, group=self.group)
        # chunk j == rank j's contribution to sequence block j
        return recv.permute(1, 0, 2, 3, 4).reshape(B, W * S_local, H_local, D)

    def _nccl_gather_heads(self, x: torch.Tensor) -> torch.Tensor:
        B, S_global, H_local, D = x.shape
        W = self.world_size
        S_local = S_global // W
        xt = x.reshape(B, W, S_local, H_local, D).permute(1, 0, 2, 3, 4).contiguous()
        recv = torch.empty_like(xt)
        dist.all_to_all_single(recv, xt, group=self.group)
        # chunk p == this rank's sequence block, head slice p
        return (
            recv.permute(1, 2, 0, 3, 4).reshape(B, S_local, W * H_local, D).contiguous()
        )

    # ---- validation ------------------------------------------------------------

    def _validate(self, x, op: str) -> None:
        if self._state != _OPEN:
            raise RuntimeError(
                f"{op} called on a {self._state} UlyssesCommunicator (use-after-close)"
            )
        if not isinstance(x, torch.Tensor):
            raise TypeError(f"{op} expects a torch.Tensor, got {type(x).__name__}")
        if x.dim() != 4:
            raise ValueError(
                f"{op} expects a 4-D [B, S, H, D] tensor, got {x.dim()}-D shape "
                f"{tuple(x.shape)}"
            )
        if x.device != self.device:
            raise ValueError(
                f"{op} tensor is on {x.device}, but this communicator is bound "
                f"to {self.device}"
            )
        if x.dtype != self.dtype:
            raise ValueError(
                f"{op} tensor dtype {x.dtype} does not match the communicator "
                f"dtype {self.dtype}"
            )
        if not x.is_contiguous():
            raise ValueError(f"{op} tensor must be contiguous")
        if any(s <= 0 for s in x.shape):
            raise ValueError(
                f"{op} tensor dims must all be positive, got shape {tuple(x.shape)}"
            )
        if x.numel() > self.max_elems:
            raise ValueError(
                f"{op} tensor has {x.numel()} elements, exceeding the "
                f"communicator capacity max_elems={self.max_elems} "
                f"(which is capped at the int32 index range {_INT32_MAX})"
            )


# =============================================================================
# Raw (advanced) kernel entry points
# =============================================================================
# Merged from the former flashinfer/comm/ulysses_a2a.py submodule: the
# function `ulysses_a2a` exported from flashinfer.comm used to shadow that
# submodule of the same name, breaking attribute-based module access. Custom
# op names, lazy JIT timing, the post-init memset fence and handle ownership
# are unchanged. The underlying CUDA kernel is adapted from ThunderKittens'
# NVLink all-to-all:
# https://github.com/HazyResearch/ThunderKittens/blob/main/kernels/parallel/all_to_all/all_to_all.cu


@functools.cache
def get_ulysses_a2a_module():
    module = gen_ulysses_a2a_module().build_and_load()

    @register_custom_op(
        "flashinfer::init_ulysses_a2a",
        mutates_args=[],
    )
    def init_ulysses_a2a(
        out_ipc_ptrs: List[int],
        signal_ipc_ptrs: List[int],
        rank: int,
        world_size: int,
        full_nvlink: bool,
    ) -> int:
        return module.init_ulysses_a2a(
            out_ipc_ptrs, signal_ipc_ptrs, rank, world_size, full_nvlink
        )

    @register_custom_op("flashinfer::dispose_ulysses_a2a", mutates_args=[])
    def dispose_ulysses_a2a(fa: int) -> None:
        module.dispose_ulysses_a2a(fa)

    @register_custom_op("flashinfer::ulysses_a2a", mutates_args=["out"])
    def ulysses_a2a(
        fa: int,
        inp: torch.Tensor,
        out: torch.Tensor,
        B: int,
        S_local: int,
        H: int,
        D: int,
        mode: int,
    ) -> None:
        module.ulysses_a2a(fa, inp, out, B, S_local, H, D, mode)

    return SimpleNamespace(
        init_ulysses_a2a=init_ulysses_a2a,
        dispose_ulysses_a2a=dispose_ulysses_a2a,
        ulysses_a2a=ulysses_a2a,
    )


@flashinfer_api
def init_ulysses_a2a(
    out_ipc_ptrs: List[int],
    signal_ipc_ptrs: List[int],
    rank: int,
    world_size: int,
    full_nvlink: bool,
) -> int:
    r"""Initialize the fused-transpose Ulysses NVLink-P2P all-to-all backend.

    .. note::
        Advanced / internal API. Prefer
        :class:`~flashinfer.comm.UlyssesCommunicator`, which selects the
        backend from the actual GPU topology before any IPC allocation or JIT
        compilation, owns the IPC workspace lifecycle, and validates operands.
        This raw entry point assumes the caller has already verified all-pairs
        NVLink P2P.

    The kernel is a *push* model: each rank writes the head/sequence blocks
    destined for its peers directly into the peers' IPC-shared output staging
    buffers over NVLink, with the Ulysses layout permutation folded into the
    write addresses. Only the output staging buffers and the signal buffers must
    be IPC-shared (allocate them with
    :func:`flashinfer.comm.create_shared_buffer`); the input tensor is read
    locally and needs no registration.

    Parameters
    ----------
    out_ipc_ptrs : list[int]
        Per-rank device pointers (opened via CUDA IPC) to the output staging
        buffers, ordered by rank. Each must be at least as large as the
        all-to-all output for this group.
    signal_ipc_ptrs : list[int]
        Per-rank device pointers to the signal buffers used for the inter-GPU
        barrier. Each buffer must be :func:`flashinfer.comm.vllm_meta_size`
        bytes (same ``Signal`` layout as the vLLM custom all-reduce).
    rank : int
        Current rank within the Ulysses group.
    world_size : int
        Ulysses group size; must be one of ``(2, 4, 6, 8)``.
    full_nvlink : bool
        ``True`` when every pair of ranks is connected via NVLink. The push
        kernel requires all-pairs P2P access; callers must gate on this.

    Returns
    -------
    int
        Opaque handle (``fa``) to pass to subsequent ``ulysses_a2a`` calls.
        Free it with :func:`dispose_ulysses_a2a`.

    Note
    ----
    ``init`` zeroes this rank's own signal buffer with a ``cudaMemset``, which
    is asynchronous with respect to the host. This wrapper therefore
    synchronizes the *current* CUDA device before returning (call it with the
    target device current). Callers still must issue a process-group barrier
    (e.g. ``torch.distributed.barrier``) after all ranks return from init and
    before the first all-to-all call — the barrier alone is not a CUDA
    completion fence, and the device sync alone is not group-wide.
    """
    if world_size not in SUPPORTED_WORLD_SIZES:
        raise ValueError(
            f"ulysses a2a only supports world size in {SUPPORTED_WORLD_SIZES}, got {world_size}"
        )
    if not full_nvlink:
        raise ValueError(
            "full_nvlink=False is not supported: the fused kernel pushes over "
            "all-pairs NVLink P2P and has no non-P2P path. Use "
            "UlyssesCommunicator(backend='auto') for topology-aware NCCL "
            "fallback instead."
        )
    module = get_ulysses_a2a_module()
    fa = module.init_ulysses_a2a(
        out_ipc_ptrs, signal_ipc_ptrs, rank, world_size, full_nvlink
    )
    # make the signal zeroing a real completion fence on this device
    try:
        torch.cuda.synchronize()
    except Exception:
        # the caller never receives fa on this path and could not dispose it:
        # ownership stays here, so release the handle before re-raising
        with contextlib.suppress(Exception):  # surface the sync error, not this
            module.dispose_ulysses_a2a(fa)
        raise
    return fa


@flashinfer_api
def dispose_ulysses_a2a(fa: int) -> None:
    r"""Release a handle returned by :func:`init_ulysses_a2a`.

    Parameters
    ----------
    fa : int
        The opaque backend handle previously returned by
        :func:`init_ulysses_a2a`. It is a C++ ``UlyssesA2A*`` reinterpreted as
        an integer (``fptr_t``), not a device pointer or a Python object, so it
        is only meaningful to this module. After this call the handle is
        dangling and must not be passed to :func:`ulysses_a2a` again.
    """
    get_ulysses_a2a_module().dispose_ulysses_a2a(fa)


@flashinfer_api
def ulysses_a2a(
    fa: int,
    inp: torch.Tensor,
    out: torch.Tensor,
    B: int,
    S_local: int,
    H: int,
    D: int,
    mode: int,
) -> None:
    r"""Fused-transpose Ulysses all-to-all.

    .. note::
        Advanced / internal API. Prefer
        :meth:`UlyssesCommunicator.scatter_heads` (``mode == 0``) and
        :meth:`UlyssesCommunicator.gather_heads` (``mode == 1``), which derive
        the geometry from the tensor shapes and validate operands.

    ``fa`` is the opaque backend handle returned by :func:`init_ulysses_a2a`
    (a C++ ``UlyssesA2A*`` reinterpreted as an integer ``fptr_t``); it selects
    the all-to-all context to run on and is not a device pointer.

    The result for this rank is written into ``out`` (bit-identical to the
    equivalent NCCL all-to-all followed by the layout permutation).

    ``mode == 0`` (input a2a): ``inp [B, S_local, H, D] -> out [B, S_global, H_local, D]``

    ``mode == 1`` (output a2a): ``inp [B, S_global, H_local, D] -> out [B, S_local, H, D]``

    where ``H`` is the *global* head count, ``H_local = H // world_size`` and
    ``S_global = S_local * world_size``. Both tensors must be contiguous CUDA
    tensors of the same dtype (float32/float16/bfloat16). All ranks must call
    with consistent geometry in the same order; a mismatch is a collective
    failure (hang or corruption), as with any collective.

    Parameters
    ----------
    fa : int
        Opaque backend handle returned by :func:`init_ulysses_a2a`.
    inp : torch.Tensor
        Contiguous 4-D CUDA input tensor.
    out : torch.Tensor
        Contiguous 4-D CUDA output tensor written in place.
    B : int
        Batch size.
    S_local : int
        Local sequence length per rank.
    H : int
        Global head count.
    D : int
        Head dimension.
    mode : int
        ``0`` for scatter-heads input all-to-all, ``1`` for gather-heads
        output all-to-all.
    """
    if type(fa) is not int or fa == 0:
        raise ValueError(
            f"fa must be a nonzero handle returned by init_ulysses_a2a, got {fa!r}"
        )
    for v, vname in (
        (B, "B"),
        (S_local, "S_local"),
        (H, "H"),
        (D, "D"),
        (mode, "mode"),
    ):
        if type(v) is not int:  # bool is an int subclass: reject it too
            raise ValueError(f"{vname} must be an int, got {type(v).__name__}")
    for name, t in (("inp", inp), ("out", out)):
        if not (isinstance(t, torch.Tensor) and t.is_cuda):
            raise ValueError(f"{name} must be a CUDA tensor")
        if not t.is_contiguous():
            raise ValueError(f"{name} must be contiguous")
        if t.dim() != 4:
            raise ValueError(f"{name} must be 4-D, got shape {tuple(t.shape)}")
    if inp.device != out.device:
        raise ValueError(f"inp is on {inp.device} but out is on {out.device}")
    if inp.dtype != out.dtype:
        raise ValueError(f"inp dtype {inp.dtype} != out dtype {out.dtype}")
    if inp.dtype not in (torch.float16, torch.bfloat16, torch.float32):
        raise ValueError(f"dtype must be float16/bfloat16/float32, got {inp.dtype}")
    if mode not in (0, 1):
        raise ValueError(f"mode must be 0 or 1, got {mode}")
    if min(B, S_local, H, D) <= 0:
        raise ValueError(f"B/S_local/H/D must be positive, got {(B, S_local, H, D)}")
    # exact-shape checks: the [B, S_local, H, D]-layout operand of each mode is
    # fully determined by the geometry args; the other operand's split of
    # (S_global, H_local) depends on world_size (unknown here), so check its
    # batch/D dims and total size
    local_shape = (B, S_local, H, D)
    checked, other = (inp, out) if mode == 0 else (out, inp)
    if tuple(checked.shape) != local_shape:
        raise ValueError(
            f"{'inp' if mode == 0 else 'out'} shape {tuple(checked.shape)} does "
            f"not match [B, S_local, H, D] = {local_shape} for mode {mode}"
        )
    if other.shape[0] != B or other.shape[3] != D or other.numel() != checked.numel():
        raise ValueError(
            f"{'out' if mode == 0 else 'inp'} shape {tuple(other.shape)} is "
            f"inconsistent with [B, S_local, H, D] = {local_shape} "
            f"(batch/D dims and total size must match)"
        )
    for name, t in (("inp", inp), ("out", out)):
        if t.numel() > _INT32_MAX:
            raise ValueError(
                f"{name} has {t.numel()} elements, exceeding the int32 index "
                f"range {_INT32_MAX} supported by the ulysses_a2a kernel"
            )
    get_ulysses_a2a_module().ulysses_a2a(fa, inp, out, B, S_local, H, D, mode)
