CUDA IPC Utilities¶

DLPack Utilities¶

Mapping Utilities¶

All-Gather Matmul¶

all_gather_matmul keeps its architecture-based default routing when backend="auto". On SM100 and SM103, backend="cake" explicitly selects the Cake fused backend for bfloat16 or float16 operands with K=8192: a contiguous [M, 8192] input with any positive M, a [8192, N] weight with N a positive multiple of 256 that is either contiguous or the transposed view of a contiguous [N, 8192] parameter (each layout has its own generated kernel; no copy is made), the process’s torch symmetric-memory backend being the default CUDA backend or NVSHMEM (the Cake backend never selects one itself), and a two-, four- or eight-rank NCCL process group. The local input may be an ordinary contiguous CUDA tensor because Cake uses internal symmetric scratch and flags for remote access and synchronization; the scratch of a group grows to the largest M seen, and that growth is the only collective after the first call. Unsupported explicit Cake requests raise instead of silently falling back. One generated source per kernel serves both architectures; the JIT loader compiles it with the exact flag set of the device it runs on.

prepare_all_gather_matmul binds the weight, the process group and a row capacity max_rows (default: the rows of the sample input) once, sizing the symmetric scratch in that single collective, and returns a callable that accepts any contiguous input with the same dtype, device and K and at most max_rows rows. Both backend="auto" and backend="cake" select this prepared Cake route; the same operand rules as the one-shot Cake route apply. The returned launcher is CUDA-graph capturable: a captured call allocates nothing but its output, and a call that would need a larger scratch raises inside capture instead of running a collective. Unsupported configurations raise during preparation instead of falling back.

Fused FP8 Row-wise GEMM + All-Reduce¶

fp8_rowwise_gemm_allreduce computes, for an NCCL process group of two, four or eight ranks on one node, Y = sum_r (A_r @ W_r.T) * (a_scale_r[:, None] * w_scale_r[None, :]) with the per-rank FP8 (e4m3) operands A_r of shape [M, K_l] and W_r of shape [N, K_l], fp32 row and column scales applied before the cross-rank sum, and a bfloat16 result [M, N] on every rank. backend="cake" (the only backend; "auto" selects it) runs one fused kernel per call: the GEMM epilogue stores each owned tile’s bf16 partial into this rank’s output and publishes every other tile’s partial into the chunk owner’s symmetric receive slot, the owner sums the partials in fixed rank order and broadcasts the result (through the NVLink multicast address on the eight-rank Blackwell carrier), so no separate all-reduce kernel or collective runs. The delivered architectures and world sizes are recorded in the loader’s route tables (flashinfer.jit.cake_fp8_rowwise_gemm_allreduce); a device or group outside them raises instead of falling back. Operands are contiguous CUDA tensors with N a positive multiple of the delivered route’s column tile (256 on SM100 / SM103), K_l a positive multiple of 128, any positive M with M * N <= 2**30, and the process’s torch symmetric-memory backend being the default CUDA backend or NVSHMEM (the Cake backend never selects one itself). The symmetric workspace of a group grows to the largest M seen; that growth is the only collective after the first call.

prepare_fp8_rowwise_gemm_allreduce binds the process group, N, K_l and a row capacity max_rows (default: the rows of the sample input) once, sizing the symmetric workspace in that single collective, and returns a callable launcher(a, a_scale, w, w_scale, *, out=None) that accepts any operands of the bound N and K_l with at most max_rows rows and never runs another collective. The launcher is CUDA-graph capturable: a captured call allocates nothing but its output and replays interleave with eager calls; a call that exceeds the capacity raises before any launch. max_rows must be identical on every rank of the group: the preparation checks it in that collective and raises ValueError on every rank when the values differ.

TensorRT-LLM AllReduce¶

Unified AllReduce Fusion API¶

All-reduce workspaces backed by SymmDeviceMemory preserve their CUDA virtual addresses across process checkpoint/restore. After quiescing all work, release the physical handles and restore them with a fresh communication backend before replaying a captured CUDA graph:

workspace.checkpoint_prepare()
workspace.checkpoint_restore(comm_backend)

Both methods are collective. Every rank must call them in the same order, and comm_backend must reproduce the original rank and world size. Repeated calls are no-ops after the workspace reaches the requested state. If an exception occurs after detach or reattach begins, do not retry or reuse the workspace; restart the affected rank. Workspaces backed by torch symmetric memory do not support this lifecycle.

class flashinfer.comm.TRTLLMAllReduceFusionWorkspace(tp_size: int, tp_rank: int, max_token_num: int, hidden_dim: int, dtype: dtype = torch.float16, comm_backend: CommBackend | None = None, group: ProcessGroup | None = None)¶

Bases: AllReduceFusionWorkspace

TensorRT-LLM workspace for AllReduce fusion.

__init__(tp_size: int, tp_rank: int, max_token_num: int, hidden_dim: int, dtype: dtype = torch.float16, comm_backend: CommBackend | None = None, group: ProcessGroup | None = None)¶

Create TensorRT-LLM AllReduce fusion workspace.

Parameters:

• tp_size – Tensor parallel size (world size)

• tp_rank – Tensor parallel rank

• max_token_num – Maximum number of tokens

• hidden_dim – Hidden dimension size

• dtype – Data type

• comm_backend – Communication backend

• group – Process group for symmetric memory rendezvous. Defaults to torch.distributed.group.WORLD.

property backend: str¶

Return backend name.

checkpoint_prepare() → None¶

Detach physical backing; repeated successful calls are no-ops.

checkpoint_restore(comm_backend: CommBackend) → None¶

Restore physical backing; repeated successful calls are no-ops.

Parameters:

comm_backend (CommBackend) – Communication backend used to recreate and exchange workspace memory handles. It must have the same rank and world size as the original allocation.

destroy() → None¶

Destroy workspace and free resources.

class flashinfer.comm.MNNVLAllReduceFusionWorkspace(mapping: Mapping, max_num_tokens: int | None = None, hidden_dim: int | None = None, dtype: dtype | None = None, buffer_size_in_bytes: int | None = None, comm_backend: CommBackend | None = None)¶

Bases: AllReduceFusionWorkspace

__init__(mapping: Mapping, max_num_tokens: int | None = None, hidden_dim: int | None = None, dtype: dtype | None = None, buffer_size_in_bytes: int | None = None, comm_backend: CommBackend | None = None)¶

Initialize the MNNVL Allreduce Fusion Workspace. The workspace will be allocated and initialized based on the provided problem size. If the maximum problem size in bytes exceeds the one-shot threshold, the workspace will be created according to the max of required one-shot size at the threshold or the required two-shot size. Note that the workspace is not bound to the given problem size. It can be reused for different problem sizes without reinitialization given the allocated size is sufficient.

If the buffer_size_in_bytes is provided, the workspace will be created according to the provided size. The user is expected to use the utility function get_required_buffer_size_bytes to calculate the required size. The actual allocation size may be larger due to alignment requirements. This covers the advanced use case, for example, the user may want to enforce oneshot strategy and ignore the heuristics.

Either max_num_tokens or buffer_size_in_bytes must be provided.

comm_backend will be used for creating the workspace and synchronization. If not provided, MPIBackend will be used which will use COMM_WORLD for synchronization.

Parameters:

• mapping – Mapping configuration containing rank info

• max_num_tokens – The maximum number of tokens in the input tensor.

• hidden_dim – The hidden dimension of the tensors to be reduced.

• dtype – The data type of the tensors to be reduced.

• buffer_size_in_bytes – The requested size in bytes for each lamport buffer. The actual allocation size may be larger due to alignment requirements. The actual usable size will be NUM_LAMPORT_BUFFERS * actual_buffer_size_per_lamport_buffer.

property backend: str¶

Return backend name.

checkpoint_prepare() → None¶

Detach physical backing; repeated successful calls are no-ops.

checkpoint_restore(comm_backend: CommBackend) → None¶

Restore physical backing; repeated successful calls are no-ops.

Parameters:

comm_backend (CommBackend) – Communication backend used to recreate and exchange MNNVL memory handles. It must have the same rank and world size as the original allocation.

destroy() → None¶

Destroy workspace and free resources.

static get_required_buffer_size_bytes(tp_size: int, num_tokens: int, hidden_dim: int, dtype: dtype, strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO) → int¶

Calculate the required buffer size for a given problem size.

is_buffer_size_sufficient(tp_size: int, num_tokens: int, hidden_dim: int, dtype: dtype, strategy: MNNVLAllreduceFusionStrategy = MNNVLAllreduceFusionStrategy.AUTO) → bool¶

Calculate the required buffer size for a given problem size.

ok = is_mnnvl_cutedsl_supported(
    4, 8192, torch.bfloat16, max_token_num=2048, apply_rms_norm=False
)
flag = torch.tensor(int(ok), device="cuda")
torch.distributed.all_reduce(flag, op=torch.distributed.ReduceOp.MIN)
if flag.item():
    workspace = create_allreduce_fusion_workspace(
        backend="mnnvl-cutedsl",
        world_size=4,
        rank=rank,
        max_token_num=2048,
        hidden_dim=8192,
        dtype=torch.bfloat16,
        apply_rms_norm=False,
        add_residual=False,
    )
    out = allreduce_fusion(
        input=x, workspace=workspace, pattern=AllReduceFusionPattern.kAllReduce
    )

Cake NVFP4-Transport AllReduce¶

cake_fp4_allreduce sums [rows, 6144] bfloat16 activations over a two- or four-rank NCCL process group through NVFP4 records: every rank quantises its rows (E2M1 values, one E4M3 scale per 16 values, one float32 row scale), pushes them into Lamport receive windows of its peers and decodes the gathered records in rank order with float32 accumulation, so the result is bitwise identical on every rank. World size 2 runs a one-shot exchange, world size 4 a two-shot exchange whose owner ranks reduce, re-encode and all-gather (one of three phase-B schedules is chosen per call from the rows each CTA owns). cake_fp4_allreduce_nvfp4 accepts pre-quantised [rows, 3456] records with one static scale per rank.

A CakeFp4AllReduceWorkspace is created collectively on every rank of the group with the same family ("allreduce" or "rmsnorm"), input format ("bf16" or "nvfp4") and row capacity; it shares the receive windows through CUDA IPC, arms their sentinels and publishes the peer window table. Calls accept any row count up to the capacity (a multiple of 4 for world size 4), perform no allocation when the output buffers are given and no host synchronisation, so they can be captured into CUDA graphs; the launch_with_pdl flag issues the kernel as a programmatic dependent launch. The programs are exact-architecture builds for SM100 (B200), SM103 (B300) and SM107; the build targets follow FLASHINFER_CUDA_ARCH_LIST when set, else the visible devices.

cake_fp4_allreduce_rmsnorm (family "rmsnorm", world size 2) fuses the exchange with the residual add and a Gemma RMSNorm: z = sum + residual in float32, residual_out = bf16(z), rinv = rsqrt(mean(z^2) + eps) from the float32 sum of squares and y = bf16(z) * rinv * (weight_bias + weight), with eps and weight_bias as runtime scalars of the CakeFp4AllReduceFusedNormSpec. The spec selects the output: bfloat16 y, or y quantised to NVFP4 (packed E2M1 values and 128x4-swizzled E4M3 block scales for a caller-supplied float32 global scale – the layout and arithmetic of fp4_quantize) or to MXFP8 (E4M3 values and swizzled UE8M0 scales, as mxfp8_quantize). cake_fp4_allreduce_rmsnorm_nvfp4 takes pre-quantised records with the same epilogue. Both return a CakeFp4AllReduceFusedNormResult.

FP8 Quantized AllReduce¶

vLLM AllReduce¶

PCIe IPC AllReduce¶

Custom all-reduce for intra-node PCIe machines without NVLink. Admission is a capability check — world size, dtype, workspace capacity, and enough payload for every rank to own a share — and shapes it rejects fall back to the caller’s own collective.

import flashinfer.comm as comm

# Collective: every rank builds the workspace with identical arguments.
# Size max_numel to the real workload -- an oversized one costs latency.
ws = comm.PcieIpcAllReduceWorkspace(group=group, max_numel=128 * 6144)

for x in activations:                     # same shapes, same order, all ranks
    if ws.supports(x):
        y = ws.all_reduce(x)              # out-of-place
    else:
        y = x.clone()
        dist.all_reduce(y, group=group)   # unsupported shape: fall back

ws.destroy()                              # collective; all ranks together

supports() is a pure function of shape and dtype, so every rank reaches the same answer without agreeing on one at runtime. It says nothing about speed: a supported shape runs a seed launch configuration — one crossover keyed on the payload in bytes, no per-machine constants — and warns once per workspace until tune() has measured the real one and persisted it, since the crossovers depend on the fabric. get_pcie_ipc_launch_config() exposes that seed for a (world_size, numel, elem_size) triple, and returns None only for shapes the kernels cannot run.

The kernels spin on peer flags with no timeout, so ranks that disagree on shape, dtype or call order hang rather than raise. One workspace serves one CUDA stream; use rebind_stream() after ordering the two if a move is genuinely needed.

PCIe IPC AllGather and ReduceScatter¶

These standalone workspaces support BF16, FP16, and FP32 tensors at world sizes 2, 4, and 8. max_numel sizes one rank’s shard: the AllGather input or the ReduceScatter output. Tensors must be contiguous, rank-2, 16-byte aligned, and contain complete 16-byte packs.

Their Python implementation lives in flashinfer.comm.pcie_ipc_collectives; the public entry points remain the flashinfer.comm exports shown below.

As with PcieIpcAllReduceWorkspace, construction, calls, and destruction are collective. Every rank must issue the same sequence and configuration, and one workspace belongs to one ordered CUDA stream. Keep the workspace alive until every captured CUDA graph replay has completed.

AllGather copies opaque 16-byte packs. ReduceScatter accumulates in FP32 and converts to the requested output dtype; its TP8 schedule converts one four-rank partial before the final accumulation, so bitwise agreement with NCCL is not part of the contract.

The TP8 CopyEngine AllGather and topology ReduceScatter schedules require live UUID/NVML evidence that logical ranks 0..3 and 4..7 form the expected two-island placement. AllGather falls back to recursive doubling when that proof is unavailable; TP8 ReduceScatter reports the input as unsupported.

import flashinfer.comm as comm

ag = comm.PcieIpcAllGatherWorkspace(group, max_local_rows * hidden, dtype=x.dtype)
rs = comm.PcieIpcReduceScatterWorkspace(group, max_local_rows * hidden, dtype=x.dtype)

gathered = ag.all_gather(x)            # [world_size * local_rows, hidden]
shard = rs.reduce_scatter(reduced_input)  # [local_rows, hidden]

ag.destroy()
rs.destroy()

launch_config() returns a conservative seed. tune([hidden, ...]) checks every candidate against NCCL, minimizes the maximum latency across ranks, and persists exact-shape results. The cache key includes the collective, world size, topology/rank-placement fingerprint, workspace limits, dtype, and shape.

The seed thresholds are heuristics; support does not imply a performance win. The variant names combine communication schedule, transport, and batching: COPY_ENGINE uses recursive doubling, and ONE_PACK means one 16-byte pack per thread in a grid-sized batch. See PCIe IPC AG/RS schedules for the seven variants, synchronization/traffic models, topology boundary, and measurement provenance.

Ulysses Context-Parallel All-to-All¶

Communication for Ulysses context parallelism over the 4-D layout [B, S, H, D]. Two layout transforms are provided; a typical attention layer makes four collective calls (q/k/v through scatter_heads, the output through gather_heads):

• scatter_heads: [B, S_local, H, D] -> [B, S_global, H_local, D] — each rank keeps head slice [rank * H_local, (rank+1) * H_local) of the full sequence;

• gather_heads: [B, S_global, H_local, D] -> [B, S_local, H, D] — the inverse, returning all heads of this rank’s sequence shard,

with H_local = H // world_size and S_global = S_local * world_size. All backends produce bit-identical results.

A third entry point, exchange_chunks, takes [1, 1, world_size, chunk] to the same shape: chunk r goes to peer r and lands in slot rank there, which is dist.all_to_all_single over an already-packed payload. It interleaves nothing, so it is for operands a preceding pack (a quantizer, say) already laid out destination-major — the head-axis transforms would make such a producer emit head-major and repack. On the "pcie" RDMA routes it is also the cheaper descriptor: each peer’s bytes are contiguous on both ends, so the NIC addresses them with the plain memory region and a linear offset instead of an interleaved MKey, which is what exempts it from the 65535-byte stride limit that bounds H*D*element_size for the other two.

Backend policy. UlyssesCommunicator selects its backend in the constructor, strictly before any IPC allocation or JIT compilation:

Typical fallback reasons reported by .fallback_reason (all conservative — anything unknown or unverifiable selects NCCL): unsupported world size (only 2/4/6/8 have fused-kernel instantiations; world_size == 1 is a no-copy passthrough), ranks spanning multiple hosts, missing pair-wise P2P or NVLink between any two concrete GPUs (verified per pair via NVML, not “some active link”), duplicate or unknown physical GPU identity, a topology probe error, inconsistent per-rank decisions, or a runtime NVLink initialization failure after a positive topology decision.

Constraints. The constructor is always collective (all ranks together); UlyssesCommunicator.close() is collective when the NVLink or PCIe backend was armed — 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 backend initialization or a collective close are exchanged as group outcomes so all ranks jointly clean up and raise (or fall back) instead of deadlocking, and a failed close may be retried. All ranks must request the same backend and agree on max_bytes and dtype; each rank may bind a different CUDA device (device accepts torch.device, str or an int ordinal, e.g. cuda:rank). With world_size > 1 the NCCL and NVLink backends (NVLink’s exchange_chunks runs all_to_all_single too) require group to support CUDA all-to-all (an NCCL process group), checked at construction before any resource is armed. Operands must be contiguous 4-D CUDA tensors of the construction dtype (float16 / bfloat16 / float32 only) on the construction device, every dim positive, and nbytes at most max_bytes; the element count of any one operand must additionally stay within the int32 kernel index range. scatter_heads requires H % world_size == 0 and gather_heads requires S_global % world_size == 0. Collectives run on the current CUDA stream; all ranks must issue the same call sequence with consistent shapes, one collective in flight per communicator at a time.

Destination-passing and workspace mode. Both collectives accept an optional preallocated out= tensor. The NCCL backend additionally accepts a reusable UlyssesWorkspace, which owns the packed send and receive buffers and removes their per-call allocations. Omitting both keywords keeps the original allocation path unchanged. The NVLink backend uses out= but does not consume the workspace because its IPC staging allocation is owned by the communicator.

Registered NVLink outputs. On the NVLink backend’s Cake route (SM100 / SM103, nvlink_kernel == "cake"), UlyssesCommunicator.allocate_output() is collective, as on the PCIe backend: every rank allocates max_bytes bytes and maps every peer’s allocation over CUDA IPC. A scatter_heads or gather_heads whose out= is such an output writes each rank’s blocks straight into the destination ranks’ outputs, with no pass through the staging buffer. All ranks pass their output of the same registration; outputs stay registered until close, which frees them. Caller-owned outputs keep the staged path, and the two may alternate on one communicator. Call allocate_output for every geometry before CUDA graph capture.

PCIe backend (experimental). It additionally accepts the 1-, 2- and 4-byte element types float16 / bfloat16 / float32, int8 / uint8 and float8_e4m3fn / float8_e5m2; its mlx5 routes — the eight-rank 4+4 NUMA hybrid and the FLASHINFER_ULYSSES_PCIE_ROUTE=rdma all-RDMA route — require batch size 1 and a head-row pitch H * D * element_size <= 65,535 bytes and refuse CUDA graph capture, while the all-P2P route takes any batch size and can be captured when every output comes from UlyssesCommunicator.allocate_output(). The JIT module links libibverbs, libmlx5 and the CUDA driver library even when topology selects the all-P2P route, so the rdma-core development packages must be present; it is the only FlashInfer JIT module that links them. Multi-rank PCIe calls require out= from UlyssesCommunicator.allocate_output(); the backend has no staging buffers, so workspace=, UlyssesCommunicator.create_workspace() and the head-chunk primitives are refused at world size above one. A rejection before native code runs (a bad operand, a refused capture, a second stream) leaves the RDMA routes usable, since every rank rejects the same call under the SPMD contract; the all-P2P route, whose barrier has no abort protocol, treats it as fail-stop. A failure inside the transport is fail-stop on every route: the communicator enters a BROKEN state that permits only a collective close. On the RDMA routes it publishes a sticky abort that releases the peers and close completes; on all-P2P close refuses to synchronize and the process must exit. See the PCIe Ulysses design note for the support matrix, output-lifetime rules, teardown protocol, routing controls, and benchmark.

Known limitations.

• PyTorch builds without torch.cuda.get_device_properties(...).uuid cannot establish physical GPU identity: auto conservatively falls back to NCCL (the reason names the missing attribute).

• When each process can only see its own GPU (e.g. one CUDA_VISIBLE_DEVICES entry per rank), peers are invisible to the P2P probe and auto falls back to NCCL.

• Out-of-range CUDA ordinals passed as strings or ints are rejected at construction; a pre-built torch.device object wraps its index into a signed byte before FlashInfer can see it (torch.device("cuda:256") is already cuda:0), so only the surviving index can be range-checked.

• Teardown metadata exchanges run bound to the communicator device; an extreme failure in the guard restore path after a completed collective can still desynchronize ranks (never observed in tests; tracked as a hardening note).

Example (Wan2.1-style attention; see the wan example for the full integration):

torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
with UlyssesCommunicator(group, max_bytes=q.nbytes,
                         dtype=torch.bfloat16, device=device) as comm:
    q_ = comm.scatter_heads(q)   # [B,S_local,H,D] -> [B,S_global,H_local,D]
    k_ = comm.scatter_heads(k)
    v_ = comm.scatter_heads(v)
    o_ = attention(q_, k_, v_)
    o = comm.gather_heads(o_)    # [B,S_global,H_local,D] -> [B,S_local,H,D]

The v0.7.2 spelling max_elems=... (a capacity in elements of dtype) is still accepted and converted to max_elems * dtype.itemsize bytes with a DeprecationWarning; it is deprecated since v0.7.3 and scheduled for removal in a future release. Pass exactly one capacity argument. comm.max_elems, the v0.7.2 attribute, remains available in the construction dtype; comm.max_bytes is the byte capacity used by the packed-chunk API.

Preallocated outputs and NCCL staging can be reused across serialized calls:

workspace = comm.create_workspace()
q_global = torch.empty(B, S_global, H_local, D,
                       dtype=q.dtype, device=q.device)
comm.scatter_heads(q, out=q_global, workspace=workspace)

Packed chunk exchange. exchange_chunks follows the byte-oriented communicator API: it accepts contiguous [1, 1, P, C] tensors, sends chunk r to rank r, and receives chunks in source-rank order. A packed Lowp payload uses dtype=torch.uint8; this describes its existing encoding and does not quantize or cast it. The NCCL and NVLink backends use NCCL for this operation; the fused NVLink kernel applies only to ordinary head-layout transforms. The experimental PCIe backend retains its registered-output path. A one-rank exchange without an explicit output is an identity; explicit outputs must not overlap the input.

For chunk exchange on NCCL or NVLink, allocate_output(x, "exchange_chunks", dtype=...) creates ordinary PyTorch receive storage. PCIe chunk outputs and NVLink Cake direct head-transform outputs follow the registration and lifetime rules described above. create_workspace(max_elems=...) and UlyssesWorkspace retain their element-based staging capacities:

# send_u8: destination-major, contiguous uint8 [P, chunk_bytes].
chunks = send_u8.view(1, 1, P, -1)
received = comm.allocate_output(chunks, "exchange_chunks", dtype=torch.uint8)
comm.exchange_chunks(chunks, out=received, dtype=torch.uint8)

class flashinfer.comm.UlyssesCommunicator(group: ProcessGroup | None = None, *, max_bytes: int | None = None, dtype: dtype, backend: str = 'auto', nvlink_kernel: str = 'auto', device: device | str | int | None = None, max_elems: int | None = None)¶

Bases: object

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 scatter_heads(), the output through gather_heads()):

• scatter_heads(): [B, S_local, H, D] -> [B, S_global, H_local, D] (each rank keeps a head slice of the full sequence)

• 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. All backends produce bit-identical results.

exchange_chunks() exchanges already-packed destination-major chunks. The NCCL and NVLink backends use NCCL for this operation. prepare_qkv() and scatter_qkv() additionally provide Sage2 INT8/FP8 QKV preparation, using this chunk exchange and a statistics AllGather on the same NCCL group.

Backend selection for ordinary scatter/gather happens in the constructor, strictly before any IPC allocation or JIT compilation (see 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); where NVLink is unavailable and FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1 is set on every rank, the experimental PCIe transport below (it warns once); NCCL otherwise — including when NVLink runtime initialization fails after a positive topology decision. Inspect backend and 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="pcie": explicitly enable the experimental single-node PCIe transport at world size 1/2/4/8. One rank is an identity path, two ranks use CUDA P2P, and four/eight ranks prefer an all-RDMA route (every peer’s payload over the rank’s mlx5 QP) with all-P2P fallback; FLASHINFER_ULYSSES_PCIE_ROUTE (auto/p2p/rdma/hybrid) forces all-P2P, all-RDMA at any multi-rank world size, or the eight-rank 4+4 NUMA hybrid (same-NUMA CUDA P2P plus cross-NUMA mlx5). Multi-rank outputs are registered up front with allocate_output() (there are no staging workspaces); "auto" reaches this backend only behind the opt-in above.

• 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.

nvlink_kernel="portable" forces the portable CUDA kernel on NVLink; "cake" requires an available Cake export. All ranks must request the same backend and NVLink kernel policy. With world_size > 1 the NCCL and NVLink backends require group to support CUDA all-to-all (an NCCL process group); this is checked at construction before any resource is armed. world_size == 1 is a passthrough: both collectives return the input tensor unchanged (no copy).

The NVLink backend launches one of two kernels on the same IPC staging and signal allocations, decided once per communicator inside the NVLink initialization (every rank takes the same route; nvlink_kernel reports it): the exported Cake programs of flashinfer.cake_ulysses_backend when the device’s compute capability is one that package delivers for this dtype and world size (SM100 / SM103), else FlashInfer’s portable NVLink kernel. On the Cake route, a collective whose out was registered with allocate_output() writes every rank’s blocks straight into the peers’ registered outputs instead of passing them through the staging buffer; the two Cake kernels use disjoint words of one signal allocation and may alternate call by call. The raw entry points (init_ulysses_a2a / ulysses_a2a / dispose_ulysses_a2a) always drive the portable kernel; a communicator’s own collectives never mix the portable and the Cake kernels on its signal words.

Constraints¶

• The constructor is always collective: every rank of group must call it together. close() is collective when the NVLink or PCIe backend was armed (their resources are peer-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 and PCIe signal protocols assume serialized calls); do not call one communicator concurrently from multiple streams or threads. The experimental PCIe P2P route enqueues asynchronously and is CUDA Graph capturable when every output comes from allocate_output() (which is itself collective and not capturable); the hybrid and all-RDMA routes block on the host and refuse capture. PCIe collectives are bound to the stream of their first call.

• Operand tensors must be contiguous 4-D CUDA tensors of the construction dtype (float16 / bfloat16 / float32) on the construction device, with every dim positive, at most 2**31 - 1 elements and at most max_bytes bytes. exchange_chunks() also accepts packed uint8 with a per-call dtype override; scatter_heads() additionally requires H % world_size == 0 and gather_heads() requires S_global % world_size == 0. PCIe additionally accepts int8/uint8 and the float8 storage types (any 1-, 2- or 4-byte element); its mlx5 routes (hybrid and all-RDMA) additionally require batch size 1 and H * D * element_size <= 65_535 bytes, while the all-P2P routes take any batch size.

• Multi-rank PCIe calls require an explicit out returned by allocate_output(). This keeps registration lifetime and overwrite points explicit; pre-register one output per live result and geometry.

• A rejection before native code runs (a bad operand, a refused capture, a second stream) is decided from the call’s own arguments, so under the SPMD contract every rank rejects it: the RDMA routes raise the ordinary error and stay usable, while the all-P2P route, whose barrier has no abort protocol, poisons teardown and enters BROKEN. A failure inside the transport is fail-stop on every route: registered buffers, queue pairs and the GPU epoch counters may already be half-mutated, so the communicator enters BROKEN and permits only a collective close(). On the RDMA routes it publishes a sticky abort that releases the peers and close() completes; on all-P2P close() refuses to synchronize and reports that the process must exit.

• Each rank may use a different CUDA device (e.g. cuda:rank); ranks must agree on max_bytes, dtype and backend.

param group:

Process group of the Ulysses ranks. Defaults to dist.group.WORLD.

type group:

torch.distributed.ProcessGroup, optional

param max_bytes:

Per-rank byte capacity of one communication operand, shared across dtypes and transforms. Sizes the NVLink staging buffer once at construction. For quantized QKV this must cover the complete packed Q/K/V payload, including scales and alignment, rather than one floating input. Statistics and other workspace allocations are separate; this is not a total memory budget. The maximum is (2**31 - 1) * 4 bytes, and each call separately enforces the int32 element-index limit.

type max_bytes:

int, optional

param max_elems:

Deprecated since v0.7.3, scheduled for removal in a future release: the v0.7.2 capacity in elements of dtype, converted to max_elems * dtype.itemsize bytes with a DeprecationWarning. Pass exactly one of max_bytes and max_elems.

type max_elems:

int, optional

param dtype:

Default element type of operands (float16 / bfloat16 / float32; PCIe additionally int8 / uint8 / float8_e4m3fn / float8_e5m2); enforced on every call that does not name its own.

type dtype:

torch.dtype

param backend:

"auto" | "nvlink" | "pcie" | "nccl" (see above).

type backend:

str

param nvlink_kernel:

"auto" | "cake" | "portable"; relevant when NVLink is used. "portable" names the in-tree CUDA kernel’s architecture coverage, not a slow path: it uses the optimized vector mover when eligible.

type nvlink_kernel:

str

param device:

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.

type device:

torch.device or str or int, optional

Examples

>>> with UlyssesCommunicator(group, max_bytes=q.nbytes, 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]

__init__(group: ProcessGroup | None = None, *, max_bytes: int | None = None, dtype: dtype, backend: str = 'auto', nvlink_kernel: str = 'auto', device: device | str | int | None = None, max_elems: int | None = None)¶

Construct a Ulysses communicator.

Parameters:

• group (Optional[ProcessGroup], optional) – Process group spanning the participating ranks. None uses torch.distributed.group.WORLD.

• max_bytes (int, optional) – Per-rank upper bound in bytes on one communication operand. Quantized QKV requires the complete packed payload to fit. Used to size the backend workspace.

• max_elems (int, optional) – Deprecated since v0.7.3 (scheduled for removal in a future release): the v0.7.2 form of the capacity, in elements of dtype; converted with dtype.itemsize. Mutually exclusive with max_bytes.

• dtype (torch.dtype) – Default element dtype for collective operands. Must be one of torch.float16, torch.bfloat16, or torch.float32; backend="pcie" additionally accepts torch.int8, torch.uint8, torch.float8_e4m3fn and torch.float8_e5m2.

• backend (str, default = "auto") – Backend selection policy. "auto" probes topology and prefers NVLink when supported, then the experimental PCIe transport where NVLink is unavailable and FLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1 is set on every rank, otherwise falls back to NCCL. "nvlink" forces the NVLink backend and raises if unavailable. "pcie" explicitly enables the experimental single-node 1/2/4/8-rank PCIe transport. "nccl" forces the NCCL path.

• nvlink_kernel (str, default = "auto") – "auto" uses Cake when an export exists, otherwise portable. "portable" forces the portable NVLink kernel; "cake" requires a Cake export. Portable includes the optimized vector mover when eligible. All ranks must select the same value.

• 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.

allocate_output(x: Tensor, op: str, *, dtype: dtype | None = None) → Tensor¶

Allocate a registered output for one Ulysses layout transform.

For backend="pcie" this is a collective cold-path operation: the fixed output uses native CUDA allocation, is registered with the selected transport, and remains registered until close(). This gives the caller a handle it owns for the communicator’s lifetime. Multi-rank PCIe calls require such an output; allocate every geometry during setup and pass it through out= on the exchange.

For backend="nvlink" on the exported Cake route (SM100 / SM103, nvlink_kernel "cake") the layout transforms register the output the same way: every rank allocates max_bytes bytes with native CUDA allocation and maps every peer’s allocation over CUDA IPC. A scatter_heads() or gather_heads() whose out= is such an output then writes every rank’s blocks straight into the destination ranks’ outputs, with no pass through the staging buffer. Every rank must pass its output of the same registration (the same position in the sequence of allocate_output calls). The memory belongs to the communicator and is freed by close(); the returned tensor must not be used afterwards. Any other output, and every output of the portable NVLink kernel, keeps the staged path.

Parameters:

• x (torch.Tensor) – Operand the output is sized for: a contiguous 4-D CUDA tensor with the shape op consumes.

• op (str) – "scatter_heads", "gather_heads" or "exchange_chunks". Outputs are registered per transform and are not interchangeable between them.

• dtype (torch.dtype, optional) – Element type for the calls this output serves, overriding the communicator dtype. For exchange_chunks, NCCL/NVLink also accept float16, bfloat16, float32 and uint8; other overrides require PCIe. The allocation keeps the communicator’s byte budget, so a narrower dtype holds proportionally more elements.

Returns:

Tensor with the shape op produces from x, on the same device and dtype as x. On multi-rank pcie, and on the multi-rank NVLink Cake route for scatter_heads and gather_heads, it stays registered until close().

Return type:

torch.Tensor

close() → None¶

Release the communicator. Idempotent once fully closed.

Collective when the NVLink or PCIe backend was armed: every rank must call close together. PCIe first requires every rank to report that native GPU/RDMA work is bounded; if any rank cannot prove that, no rank enters device synchronization or releases a registration, and process termination is required. A safe PCIe close then synchronizes, closes every output’s peer imports, disposes its MR/MKey registration, and finally disposes the transport. NVLink similarly synchronizes before releasing its kernel handle, peer mappings, and exports. Each stage has bounded group-coordinated retries. If a retryable teardown stage still cannot complete, every rank raises the same error and remains CLOSING; every rank may retry close(). CLOSED is reached only after complete group-wide teardown. Pure NCCL holds no resources and closes locally.

create_workspace(*, max_elems: int | None = None) → UlyssesWorkspace¶

Allocate reusable NCCL send/receive staging buffers.

This operation is rank-local and non-collective. max_elems defaults to the communicator capacity and may be smaller when a caller knows the maximum chunk size it will communicate. The multi-rank pcie backend has no staging buffers and rejects this call; its outputs are registered with allocate_output() instead.

Parameters:

max_elems (int, optional) – Capacity of each send and receive staging buffer in elements. None fills the communicator’s max_bytes capacity at the communicator dtype.

Returns:

A workspace on the communicator’s device with its dtype.

Return type:

UlyssesWorkspace

exchange_chunks(x: Tensor, *, out: Tensor | None = None, dtype: dtype | None = None) → Tensor¶

[1, 1, W, C] -> [1, 1, W, C]: equal-length chunk all-to-all.

Chunk r of this rank’s input goes to peer r, and lands in slot rank of that peer’s output – the semantics of torch.distributed.all_to_all_single on an already-packed payload.

This is the transform for a payload that was produced destination-major by a preceding pack (a quantizer, say), so its per-peer bytes are already contiguous. scatter_heads() and gather_heads() do the head-axis interleave instead, which forces a head-major layout on the producer; this entry point removes that constraint. On the RDMA routes it is also the cheaper descriptor – a single-row MKey, exempt from the 65535-byte interleaved-stride limit that bounds H*D* element_size for the other two.

Parameters:

• x (torch.Tensor) – Contiguous 4-D CUDA tensor shaped [1, 1, world_size, chunk]. chunk is in elements of x’s dtype; a packed uint8 payload passes dtype=torch.uint8 and chunk in bytes.

• out (torch.Tensor, optional) – Preallocated output of the same shape, not overlapping x. Multi-rank PCIe requires an output from allocate_output().

• dtype (torch.dtype, optional) – Element type for this call, overriding the communicator dtype. PCIe accepts its supported storage types; NCCL/NVLink accept float16, bfloat16, float32 and uint8. No conversion is performed.

Returns:

Tensor of the same shape, device and dtype as x.

Return type:

torch.Tensor

gather_heads(x: Tensor, *, out: Tensor | None = None, dtype: dtype | None = None, workspace: UlyssesWorkspace | None = None) → Tensor¶

[B, S_global, H_local, D] -> [B, S_local, H, D].

Inverse of 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].

• out (torch.Tensor, optional) – Preallocated contiguous output with shape [B, S_global // world_size, H_local * world_size, D]. Supplying it removes the public output allocation. It must not alias x when world_size > 1. Multi-rank PCIe requires an output returned by allocate_output(); on the NVLink Cake route such an output takes the direct path.

• dtype (torch.dtype, optional) – Element type for this call, overriding the communicator dtype. pcie backend only. max_bytes is denominated in bytes, so a narrower dtype here simply fits more elements in the same capacity.

• workspace (UlyssesWorkspace, optional) – Reusable NCCL pack/receive storage. Supplying it removes the two NCCL staging allocations. It is validated but unused by the NVLink backend, and rejected by the multi-rank pcie backend, which has no staging buffers.

Returns:

Tensor with shape [B, S_local, H, D] on the same device and dtype as x.

Return type:

torch.Tensor

gather_output_head_chunk(x: Tensor, *, local_heads: int, head_offset: int, out: Tensor, workspace: UlyssesWorkspace | None = None) → Tensor¶

Gather one attention-output head band into a full destination.

x has shape [B, S_global, head_count, D]. out is the preallocated full output [B, S_local, world_size * local_heads, D]. Only the selected head band is written; calling this method for every non-overlapping band in a complete schedule reconstructs ordinary whole-head Ulysses output communication.

The NCCL path packs the sequence split directly into reusable staging and merges the received head band directly into out. The NVLink path uses its fused-transpose collective followed by the same merge primitive. Runs on the current CUDA stream.

Parameters:

• x (torch.Tensor) – Positive-strided CUDA tensor with shape [B, S_global, head_count, D].

• local_heads (int) – Total number of attention-output heads owned by this rank before the output all-to-all.

• head_offset (int) – Start of this band within the rank’s local_heads.

• out (torch.Tensor) – Preallocated contiguous destination with shape [B, S_global // world_size, world_size * local_heads, D]. Only this head band is modified.

• workspace (UlyssesWorkspace, optional) – Reusable staging storage large enough for x. A workspace cannot service concurrent collectives.

Returns:

out after merging the gathered head band.

Return type:

torch.Tensor

input_buffer(out: Tensor, shape: Sequence[int]) → Tensor¶

The transport’s own input staging buffer behind a registered output.

The RDMA routes never let the NIC read caller memory: every operand is copied into a landing buffer registered once alongside out. Building the operand directly in the buffer returned here removes that copy – for a layer-scale scatter it is the largest device-to-device move on the path – and the caller’s own source allocation with it.

The result must reach scatter_heads() or gather_heads() unmodified: the fast path is exact pointer identity, and a slice or a re-view that moves the base pointer is rejected rather than silently staged, because on the scatter path the NIC reads the landing buffer whatever the caller passes.

Parameters:

• out (torch.Tensor) – An output returned by allocate_output().

• shape (Sequence[int]) – The [B, S, H, D] operand shape to view the buffer as.

Returns:

A view of transport-owned memory in the dtype out was registered under. It stays alive as long as out does, and every exchange on out overwrites it – it is where the next operand is built, not where one is kept.

Return type:

torch.Tensor

property max_elems: int¶

Legacy capacity in elements of the communicator’s construction dtype.

property nvlink_kernel: str | None¶

"cake" (the exported Cake programs of flashinfer.cake_ulysses_backend), "portable" (FlashInfer’s portable NVLink kernel) or None when the NVLink backend is not armed (NCCL backend, world_size == 1, closed).

Type:

Which kernel the NVLink backend launches

prepare_qkv(qkv_shape: Sequence[int], *, quantization: str = 'sage2', used_sequence: int | None = None) → UlyssesQKVWorkspace¶

Collectively prepare Sage2 QKV communication and reusable storage.

Every rank must call this method in the same order, outside CUDA graph capture. Shape/configuration, layout/JIT availability, and allocation outcomes are agreed in three separate stages. A recoverable preparation failure releases only this call’s storage; it does not close the communicator or invalidate existing workspaces. Missing participants and CUDA/process-group failures are not recoverable through this protocol.

This method does not change backend. Quantized QKV always uses the group’s CUDA NCCL backend, even when ordinary scatter/gather uses NVLink. The experimental PCIe backend is not supported by this prepared API. Supports homogeneous SM90 or SM120 groups of 2/4/8 ranks.

Parameters:

• qkv_shape (Sequence[int]) – Common local [B, L, H, D] shape, with positive dimensions, D equal to 64 or 128, and H divisible by the group size. Each Q/K/V and the combined byte payload must fit the int32 element-index limit. The complete packed QKV payload, including scales and alignment, must fit max_bytes. Statistics are allocated separately.

• quantization (str, default = "sage2") – QKV encoding policy. Only "sage2" is supported: grouped INT8 Q/K and channel-scaled FP8 E4M3 V, using the device’s Sage2 layout. The communicator dtype must be FP16 or BF16.

• used_sequence (int, optional) – Global live prefix length U in (0, world_size * L]. None uses the whole sequence. Callers must zero input rows outside this prefix; this method does not initialize inputs.

Returns:

Owner-bound, read-only geometry with reusable send/receive and statistics-gather buffers. Shape or live-length changes require another collective preparation call.

Return type:

UlyssesQKVWorkspace

scatter_heads(x: Tensor, *, out: Tensor | None = None, dtype: dtype | None = None, workspace: UlyssesWorkspace | None = None) → Tensor¶

[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].

• out (torch.Tensor, optional) – Preallocated contiguous output with shape [B, S_local * world_size, H // world_size, D]. Supplying it removes the public output allocation. It must not alias x when world_size > 1. Multi-rank PCIe requires an output returned by allocate_output(); on the NVLink Cake route such an output takes the direct path.

• dtype (torch.dtype, optional) – Element type for this call, overriding the communicator dtype. pcie backend only. max_bytes is denominated in bytes, so a narrower dtype here simply fits more elements in the same capacity.

• workspace (UlyssesWorkspace, optional) – Reusable NCCL pack/receive storage. Supplying it removes the two NCCL staging allocations. It is validated but unused by the NVLink backend, which owns IPC staging internally, and rejected by the multi-rank pcie backend, which has no staging buffers.

Returns:

Tensor with shape [B, S_global, H_local, D] on the same device and dtype as x.

Return type:

torch.Tensor

scatter_qkv(q: Tensor, k: Tensor, v: Tensor, *, workspace: UlyssesQKVWorkspace, out: UlyssesQKV | None = None) → UlyssesQKV¶

Quantize and exchange QKV for the prepared SageAttention2 layout.

Runs one FP32 statistics AllGather and one jointly packed uint8 AllToAll on group, using NCCL independently of backend. Execution uses the current CUDA stream; all ranks must issue the same sequence of calls with valid inputs. There is no per-call metadata collective, implicit preparation, or precision fallback. Local statistics still allocate temporary tensors.

Parameters:

• q (torch.Tensor) – FP16/BF16 CUDA tensors matching the prepared [B,L,H,D] shape, dtype and device. D must be contiguous; the base pointer and each outer stride must preserve 16-byte row alignment. Projection views satisfying those constraints do not require copies. Input rows outside the prepared global live prefix must already be zero.

• k (torch.Tensor) – FP16/BF16 CUDA tensors matching the prepared [B,L,H,D] shape, dtype and device. D must be contiguous; the base pointer and each outer stride must preserve 16-byte row alignment. Projection views satisfying those constraints do not require copies. Input rows outside the prepared global live prefix must already be zero.

• v (torch.Tensor) – FP16/BF16 CUDA tensors matching the prepared [B,L,H,D] shape, dtype and device. D must be contiguous; the base pointer and each outer stride must preserve 16-byte row alignment. Projection views satisfying those constraints do not require copies. Input rows outside the prepared global live prefix must already be zero.

• workspace (UlyssesQKVWorkspace) – Workspace returned by this communicator’s prepare_qkv. Must not be shared by concurrent operations or another communicator.

• out (UlyssesQKV, optional) – Reuse these six output tensors. Shapes, dtypes, device, layout, logical/live lengths and input dtype must match the preparation. Outputs must be contiguous and cannot overlap each other, the inputs, or workspace storage. None allocates independent final storage that remains valid after later workspace reuse.

Returns:

Pre-quantized Q/K/V and three FP32 scales, with consumer metadata. Returns out itself when supplied. Attention uses the live prefix; the caller restores a full-length floating output before invoking gather_heads.

Return type:

UlyssesQKV

scatter_qkv_head_chunk(query: Tensor, key: Tensor, value: Tensor, *, head_offset: int, head_count: int, out: Tensor | None = None, workspace: UlyssesWorkspace | None = None) → Tensor¶

Pack and scatter one Q/K/V head band with one collective.

query, key and value have shape [B, S_local, H, D] and may be independent positive-strided views. The output is a fused [B, S_global, head_count, 3 * D] payload; slicing its final dimension produces the three attention operands.

The method is a transport primitive, not a pipeline scheduler. A caller that overlaps input and output communication must use separate communicators and workspaces. Runs on the current CUDA stream.

When the NCCL backend, B == 1, out is None and a workspace is supplied, the returned tensor aliases the workspace receive buffer and remains valid only until that workspace is reused. Pass out for an explicit lifetime.

Parameters:

• query (torch.Tensor) – Positive-strided CUDA tensor with shape [B, S_local, H, D].

• key (torch.Tensor) – Positive-strided CUDA tensor with the same shape, dtype, and device as query.

• value (torch.Tensor) – Positive-strided CUDA tensor with the same shape, dtype, and device as query.

• head_offset (int) – Start of the selected band within each destination rank’s H // world_size local heads.

• head_count (int) – Number of consecutive local heads in the selected band.

• out (torch.Tensor, optional) – Preallocated contiguous result with shape [B, S_local * world_size, head_count, 3 * D]. It must not alias query, key, or value.

• workspace (UlyssesWorkspace, optional) – Reusable staging storage large enough for the fused Q/K/V payload. A workspace cannot service concurrent collectives.

Returns:

Fused Q/K/V payload with shape [B, S_global, head_count, 3 * D].

Return type:

torch.Tensor

class flashinfer.comm.UlyssesWorkspace(*, max_elems: int, dtype: dtype, device: device | str | int | None = None)¶

Bases: object

Reusable staging storage for NCCL Ulysses collectives.

A workspace owns two flat CUDA buffers: one for the packed all-to-all input and one for the receive result. Passing it to UlyssesCommunicator.scatter_heads() or UlyssesCommunicator.gather_heads() removes the per-call NCCL staging allocations. The NVLink backend already owns IPC staging storage, so it validates but otherwise does not use this object.

Workspaces are local allocations: they own no process group and create no streams. A workspace may be shared by serialized scatter/gather calls, but must not be used by concurrent collectives.

Parameters:

• max_elems (int) – Capacity of each staging buffer in elements.

• dtype (torch.dtype) – Element dtype; float16, bfloat16, or float32.

• device (torch.device or str or int, optional) – CUDA device. None uses the current CUDA device.

__init__(*, max_elems: int, dtype: dtype, device: device | str | int | None = None)¶

Initialize reusable NCCL send and receive staging buffers.

Parameters:

• max_elems (int) – Capacity of each staging buffer in elements.

• dtype (torch.dtype) – Element dtype; float16, bfloat16, or float32.

• device (torch.device or str or int, optional) – CUDA device. None uses the current CUDA device.

property recv_buffer: Tensor¶

Flat receive staging buffer (advanced/debugging use).

property send_buffer: Tensor¶

Flat send staging buffer (advanced/debugging use).

# The original floating Q/K/V total is a conservative wire-byte budget.
with UlyssesCommunicator(group, max_bytes=3 * q.nbytes, dtype=q.dtype,
                         device=q.device, backend="nccl") as comm:
    workspace = comm.prepare_qkv(q.shape, used_sequence=used)
    encoded = comm.scatter_qkv(q, k, v, workspace=workspace)
    # A framework adapter calls Sage with encoded.q/k/v and the scales.
    # It returns [B, P*L, H/P, D], in encoded.input_dtype, with a zero tail.
    output = comm.gather_heads(sage_adapter(encoded))
    # After the previous attention has consumed the buffers:
    encoded = comm.scatter_qkv(q_next, k_next, v_next,
                               workspace=workspace, out=encoded)

schedule = ((0, 3), (3, 8), (11, 3))  # local_heads == 14
full_output = torch.empty_like(q)
for head_offset, head_count in schedule:
    qkv = input_comm.scatter_qkv_head_chunk(
        q, k, v,
        head_offset=head_offset,
        head_count=head_count,
        out=qkv_chunk_buffers[head_count],
        workspace=input_workspace,
    )
    q_chunk, k_chunk, v_chunk = qkv.split(head_dim, dim=-1)
    o_chunk = attention(q_chunk, k_chunk, v_chunk)
    output_comm.gather_output_head_chunk(
        o_chunk,
        local_heads=14,
        head_offset=head_offset,
        out=full_output,
        workspace=output_workspace,
    )