flashinfer.comm.init_ulysses_a2a

flashinfer.comm.init_ulysses_a2a(out_ipc_ptrs: List[int], signal_ipc_ptrs: List[int], rank: int, world_size: int, full_nvlink: bool) → int

Initialize the fused-transpose Ulysses NVLink-P2P all-to-all backend.

Note

Advanced / internal API. Prefer 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 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 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:

Opaque handle (fa) to pass to subsequent ulysses_a2a calls. Free it with dispose_ulysses_a2a().

Return type:

int

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.