flashinfer.comm

This module provides communication primitives and utilities for distributed computing, including CUDA IPC, AllReduce operations, and memory management utilities.

CUDA IPC Utilities

CudaRTLibrary([so_file])
create_shared_buffer(size_in_bytes[, group])	Allocate a buffer and share it across the process group via CUDA IPC.
free_shared_buffer(pointers[, group])	Free a shared buffer previously created by create_shared_buffer().

DLPack Utilities

pack_strided_memory(ptr, segment_size, ...)

Pack a strided device allocation as a PyTorch tensor view.

Mapping Utilities

Mapping([world_size, rank, gpus_per_node, ...])

A node with 8 GPUs, tp_size = 4, cp_size = 1, pp_size = 2

All-Gather Matmul

all_gather_matmul keeps its architecture-based default routing when backend="auto". On SM100 and SM103, backend="cake" explicitly selects the source-built fused backend for contiguous bfloat16 or float16 inputs with K=8192, N=2048, a positive M divisible by 128, an NVSHMEM symmetric memory backend, and a two- or four-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. Unsupported explicit Cake requests raise instead of silently falling back. The packaged manifest carries the exact dynamic shared-memory requirement resolved for every generated main route; the loader validates that value against the packaged CUDA source before compiling the host launcher.

prepare_all_gather_matmul prepares the source-built packed-QKV route for SM103, bfloat16, four-rank NCCL groups, contiguous [M, 8192] inputs, and a contiguous [8192, 2560] weight, where M is a positive multiple of 128. It binds the weight and process group once and returns a callable that accepts a contiguous input with the same shape, dtype, and device. Both backend="auto" and backend="cake" select this prepared route. Unsupported configurations raise during preparation instead of falling back.

all_gather_matmul(inp, w, group, *[, ...])	Run push-wait all-gather matmul.
prepare_all_gather_matmul(inp, w, group, *)	Prepare the packed-QKV all-gather matmul launcher.

TensorRT-LLM AllReduce

Types and Enums

AllReduceFusionOp()
AllReduceFusionPattern()
AllReduceStrategyConfig()
AllReduceStrategyType()
QuantizationSFLayout()

Core Operations

trtllm_allreduce_fusion(allreduce_in, ...[, ...])	Parameters: - allreduce_in: the input tensor. [token_num, hidden_dim] - world_size: the size of the process group. - world_rank: the rank of the current process. - token_num: the number of tokens in the sequence. - hidden_dim: the dimension of the hidden states. - workspace_ptrs: the workspace pointers. - launch_with_pdl: whether to launch with pdl. - use_oneshot: whether to use oneshot. If None, internal heuristics will be used. - trigger_completion_at_end: whether to trigger completion at the end. - fp32_acc: whether to use fp32 accumulation. - pattern_code: the pattern code. - allreduce_out: the output tensor. [token_num, hidden_dim] - residual_in: the residual input tensor. [token_num, hidden_dim] - residual_out: the residual output tensor. [token_num, hidden_dim] - norm_out: the norm output tensor. [token_num, hidden_dim] - quant_out: the quant output tensor. [token_num, hidden_dim] - scale_out: the scale output tensor. Initialization referece: tests/comm/test_trtllm_allreduce_fusion.py - rms_gamma: the rms gamma tensor. [hidden_dim] - rms_eps: the rms epsilon value. - scale_factor: the scale factor. For cudaGraphs safety, it should be a tensor. - layout_code: the layout code. - metadata: optional workspace metadata dict from create_ipc_workspace_for_all_reduce_fusion. If provided, validates that token_num <= max_token_num, world_size == tp_size, and hidden_dim == workspace hidden_dim. Raises ValueError if validation fails. - block_quant_group_size: group size (in elements along hidden_dim) for per-token-group block-wise FP8 quantization patterns (e.g. kPerTokenGroupFP8Packed / DeepSeek-style FP8 with UE8M0 packed scales). Number of consecutive elements that share a single scale factor. Must be > 0 and divide hidden_dim when the pattern requires it; ignored (treated as 0 / unused) for patterns that do not perform block-quantization. - weight_bias: bias added to rms_gamma before scaling. None or 0.0 -> standard RMSNorm (out = gamma * x * rsqrt(...)). 1.0 -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(...)). Ignored for kAllReduce and quant-only patterns that don't apply RMSNorm.
trtllm_moe_allreduce_fusion(world_size, ...)	Parameters: - world_size: the size of the process group. - world_rank: the rank of the current process. - token_num: the number of tokens in the sequence. - hidden_dim: the dimension of the hidden states. - workspace_ptrs: the workspace pointers. - launch_with_pdl: whether to launch with pdl. - residual_in: the residual input tensor. [token_num, hidden_dim] - rms_gamma: the rms gamma tensor. [hidden_dim] - rms_eps: the rms epsilon value. - scale_factor: the scale factor. - moe_reduction_device_num_experts: the number of experts. - moe_reduction_scale_input: the scale input tensor. [token_num, hidden_dim] - moe_reduction_active_experts_token_input: the active experts token input tensor. [token_num, hidden_dim] - moe_reduction_token_input: the token input tensor. [token_num, hidden_dim] - layout_code: the layout code. - moe_allreduce_out: the moe allreduce output tensor. [token_num, hidden_dim] - residual_out: the residual output tensor. [token_num, hidden_dim] - norm_out: the norm output tensor. [token_num, hidden_dim] - quant_out: the quant output tensor. [token_num // 4, hidden_dim], fp16/bf16 -> fp4 - scale_out: the scale output tensor. Initialization referece: tests/comm/test_trtllm_moe_allreduce_fusion.py - weight_bias: bias added to rms_gamma before scaling. None or 0.0 -> standard RMSNorm (out = gamma * x * rsqrt(...)). 1.0 -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(...)). - backend: "trtllm" (default) or the constrained "cake" SM100/SM103 backend. The public Python API has 22 parameters; the isolated source module's run_reduction entry has an exact 18-argument FFI ABI. The optional backend supports contiguous FP16/BF16 tensors, world sizes 2, 4, and 8, hidden_dim=7168, token payloads within the existing Lamport MAX_COMM_SIZE byte limit, and residual plus norm outputs. It does not support quantization. weight_bias remains a runtime value; None is passed to the kernel as 0.0. On SM100 and SM103 with world_size 2, 4 or 8 and an moe_allreduce_out tensor, "cake" runs the verified source export of the Cake all-reduce union (flashinfer.jit.cake_trtllm_moe_allreduce_union); the route binds the workspace pointer table the way trtllm_create_ipc_workspace_for_all_reduce_fusion registers it and needs no device readback. Every other Cake configuration keeps the isolated source bundle.
trtllm_moe_finalize_allreduce_fusion(...[, ...])	Parameters: - allreduce_in: the permuted/padded MoE expert output tensor. Shape [num_permuted_rows, hidden_dim]. Rows are referenced by expanded_idx_to_permuted_idx; num_permuted_rows may be larger than token_num * top_k due to expert padding. - residual_in: the residual input tensor. [token_num, hidden_dim] - norm_weight: the norm weight tensor. [hidden_dim] - expanded_idx_to_permuted_idx: the expanded index to permuted index tensor. [token_num, top_k] - norm_out: the norm output tensor. [token_num, hidden_dim] - residual_out: the residual output tensor. [token_num, hidden_dim] - quant_out: the quant output tensor. [token_num // 4, hidden_dim], fp16/bf16 -> fp4 - scale_out: the scale output tensor. [token_num // SF_VEC_SIZE, hidden_dim], fp16/bf16 -> fp4 - workspace_ptrs: the workspace pointers. - launch_with_pdl: whether to launch with pdl. - world_rank: the rank of the current process. - world_size: the size of the process group. - eps: the epsilon value. - shared_expert_output: the shared expert output tensor. [token_num, hidden_dim] - expert_scale_factor: the expert scale factor tensor. [token_num, top_k] - routed_scaling_factor: the routed scaling factor. - weight_bias: bias added to rms_gamma before scaling. None or 0.0 -> standard RMSNorm (out = gamma * x * rsqrt(...)). 1.0 -> Gemma / Qwen3.5 RMSNorm (out = (1 + gamma) * x * rsqrt(...)). - backend: implementation to launch. "trtllm" preserves the existing backend. "cake" selects the source-built SM100/SM103 kernels for hidden size 7168 and tensor-parallel sizes 2, 4, or 8.

Workspace Management

trtllm_create_ipc_workspace_for_all_reduce_fusion(...)	Parameters: - tp_rank: the rank of the current process.
trtllm_destroy_ipc_workspace_for_all_reduce(...)	Release the symmetric-memory references held for an IPC workspace.
trtllm_destroy_ipc_workspace_for_all_reduce_fusion(...)	Destroy a workspace created by trtllm_create_ipc_workspace_for_all_reduce_fusion.

Initialization and Utilities

trtllm_lamport_initialize(buffer_ptr, size, ...)	Initialize a single Lamport-style buffer to negative zero.
trtllm_lamport_initialize_all(buffer_0_ptr, ...)	Initialize three Lamport buffers to negative zero.
compute_fp4_swizzled_layout_sf_size(...)	Compute the padded size (rows times columns) of the FP4 swizzled layout.

Unified AllReduce Fusion API

allreduce_fusion(input, workspace, pattern)	AllReduce + RMSNorm fusion operation, with optional FP8/NVFP4 quantization for supported backends.
create_allreduce_fusion_workspace([backend, ...])	Create workspace for AllReduce fusion operations.
AllReduceFusionWorkspace(world_size, rank)	Base class for AllReduce fusion workspaces.

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 max_num_tokens is larger than the one-shot threshold, the workspace will be created according to the max of required one-shot size at threshold, or the required two-shot size. Note that the workspace is not bind to the given problem size. It can be reused for different problem size 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 used 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.

FP8 Quantized AllReduce

quantized_all_reduce(inp, group, *[, ...])

FP8 quantized two-shot AllReduce with per-block scaling.

vLLM AllReduce

vllm_all_reduce(fa, inp, out, reg_buffer, ...)	Perform an out-of-place all-reduce via the vLLM custom kernel.
vllm_dispose(fa)	Release the resources held by a vLLM custom all-reduce handle.
vllm_init_custom_ar(ipc_tensors, rank_data, ...)	Initialize the vLLM custom all-reduce backend.
vllm_register_buffer(fa, fake_ipc_ptrs)	Register a peer's IPC-shared buffer with the local all-reduce handle.
vllm_register_graph_buffers(fa, handles, offsets)	Register graph-capture buffers across the all-reduce world.
vllm_get_graph_buffer_ipc_meta(fa)	Return IPC metadata for graph-capture buffers.
vllm_meta_size()	Return the size of the vLLM all-reduce metadata structure in bytes.

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.

Two data planes

Above roughly a megabyte the workspace can move the payload with the copy engines instead of with SM load/store. The reason is not that the copy engine is faster per hop — in isolation the two are comparable — but that it keeps its throughput when every rank transfers at once, where SM peer traffic loses bandwidth to the concurrency. So the gap widens with world size, which is what makes a second plane worth having at these payload sizes.

It is a second protocol rather than a flag on the existing kernels. The SM path encodes readiness in the payload — +0.0 means “not yet written” — which a copy engine can neither produce nor observe, so the ring synchronises through monotonic flags on side streams instead. Nothing about the SM kernels changes.

Both planes are ordinary tuner candidates, so nothing needs to be declared at the call site and there is no notion of a prefill or decode “phase”: the configuration follows the payload in bytes. PcieIpcVariant gains COPY_ENGINE_RING, a flat neighbour ring, and COPY_ENGINE_ISLAND, a 4+4 decomposition offered only at world size 8 and only on rootcplx-noswitch, where a ring’s two socket-crossing hops would otherwise set the pace. On that variant blocks carries the ring’s sub-chunk depth rather than a grid size.

Warning

Tuning covers the batch sizes it is given, and shapes above the largest one are served by whatever was measured there. Leave tune_batches unset and the ladder is derived from max_numel, which covers everything the workspace admits; pass it explicitly and the top of the range is yours to get right. tune() warns when an explicit list leaves a gap, and refuses to persist results measured below the default sample counts, since a table that is quietly wrong is worse than one that is missing.

PcieIpcAllReduceWorkspace(group, max_numel)	Shared workspace for the PCIe IPC all-reduce.
PcieIpcLaunchConfig	alias of IpcLaunchConfig
PcieIpcVariant	alias of IpcVariant
get_pcie_ipc_launch_config(world_size, ...)	Launch configuration for one shape, or None when unsupported.
probe_pcie_ipc_rank_topology(rank[, device])	Probe this rank's GPU against other visible GPUs, recording failures.
resolve_pcie_ipc_profile(group[, requested, ...])	Probe every rank and agree on one profile.

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.

PcieIpcAllGatherWorkspace(group, max_numel)	Workspace for out-of-place PCIe IPC all-gather.
PcieIpcAllGatherLaunchConfig(blocks, ...)
PcieIpcAllGatherVariant(value)	All-gather kernel selected across the FFI boundary.
get_pcie_ipc_all_gather_launch_config(...[, ...])	Return a deterministic all-gather launch seed.
PcieIpcReduceScatterWorkspace(group, max_numel)	Workspace for SUM reduce-scatter over PCIe IPC.
PcieIpcReduceScatterLaunchConfig(blocks, ...)
PcieIpcReduceScatterVariant(value)	Reduce-scatter kernel selected across the FFI boundary.
get_pcie_ipc_reduce_scatter_launch_config(...)	Return a deterministic reduce-scatter launch seed.

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

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

backend=	behavior
"auto"	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. The instance exposes .backend (effective), .fallback_reason and .decision / .topology_decision.
"nvlink"	force the fused kernel; raises on every rank (before IPC/JIT for topology failures) when it cannot be used.
"nccl"	force dist.all_to_all_single + permute; skips the topology/NVML probe and all IPC/JIT (the constructor still resolves/guards the CUDA device and performs CUDA-backed metadata collectives); any world size.

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 only when the NVLink 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 the NVLink 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_elems 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 backend (forced or fallen back to) requires group to support CUDA all-to-all (an NCCL process group), checked at construction. Operands must be contiguous 4-D CUDA tensors of the construction dtype (float16 / bfloat16 / float32 only) on the construction device, every dim positive, at most max_elems (≤ 2^31 − 1) elements; 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.

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):

with UlyssesCommunicator(group, max_elems=B * S_local * H * D,
                         dtype=torch.bfloat16) 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]

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)

class flashinfer.comm.UlyssesCommunicator(group: ProcessGroup | None = None, *, max_elems: int, dtype: dtype, backend: str = 'auto', device: device | str | 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. Both backends produce bit-identical results.

Backend selection 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); 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="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. 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; scatter_heads() additionally requires H % world_size == 0 and 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.

param group:

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

type group:

torch.distributed.ProcessGroup, optional

param max_elems:

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.

type max_elems:

int

param dtype:

Element type of all operands (float16 / bfloat16 / float32); enforced on every call.

type dtype:

torch.dtype

param backend:

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

type backend:

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_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]

__init__(group: ProcessGroup | None = None, *, max_elems: int, dtype: dtype, backend: str = 'auto', device: device | str | 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_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.

close() → None

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.

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.

Parameters:

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

Returns:

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

Return type:

UlyssesWorkspace

gather_heads(x: Tensor, *, out: Tensor | 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.

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

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

scatter_heads(x: Tensor, *, out: Tensor | 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.

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

Returns:

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

Return type:

torch.Tensor

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

Head-Chunk Layout and Transport Primitives

Head-chunk APIs are opt-in building blocks for framework-owned communication/compute pipelines. They do not choose a chunk schedule, create streams/events or process groups, or alter the ordinary scatter_heads / gather_heads path.

pack_ulysses_qkv_head_chunk accepts independent positive-strided Q/K/V views, selects the same destination-local head band from every Ulysses rank’s head slice, and fuses Q/K/V into the last dimension. The communicator’s scatter_qkv_head_chunk method combines that transform with one all-to-all. After attention, gather_output_head_chunk performs the reverse all-to-all and writes the compact band directly into the full output. A framework that overlaps the two communication directions must use separate communicators and workspaces:

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,
    )

Head chunking is not automatically profitable. In particular, tensor parallelism reduces the effective local head count to approximately H / (TP * Ulysses). Frameworks should enable a schedule only after all ranks agree that the attention backend supports every chunk head count and an offline benchmark shows a positive result. Small local-head counts, GQA/MLA with unequal Q/K/V heads, Ring Attention, and Attention2D should use the ordinary path unless separately implemented and measured.

flashinfer.comm.pack_ulysses_qkv_head_chunk(query: Tensor, key: Tensor, value: Tensor, *, world_size: int, head_offset: int, head_count: int, out: Tensor | None = None) → Tensor

Pack one destination-local head band of Q/K/V.

The three inputs use [B, S_local, H, D] and may be independent positive-strided views (including views produced by a fused projection). H must be divisible by world_size. The selected local band is applied inside every destination rank’s head slice, and the result is a contiguous fused payload [B, S_local, world_size * head_count, 3 * D].

The payload can be passed directly to flashinfer.comm.UlyssesCommunicator.scatter_heads(); its result has shape [B, S_global, head_count, 3 * D]. Slicing the final dimension yields Q/K/V views for the attention backend.

This operation runs on the caller’s current CUDA stream and allocates only when out is omitted.

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.

• world_size (int) – Number of Ulysses ranks. H must be divisible by world_size.

• head_offset (int) – Start of the selected band within every 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.

Returns:

Contiguous fused Q/K/V payload with shape [B, S_local, world_size * head_count, 3 * D].

Return type:

torch.Tensor

flashinfer.comm.merge_ulysses_output_head_chunk(received: Tensor, *, world_size: int, local_heads: int, head_offset: int, out: Tensor) → Tensor

Merge one compact Ulysses output head band into a full output.

received is the result of gathering a chunked attention output and has shape [B, S_local, world_size * head_count, D]. out is the full [B, S_local, world_size * local_heads, D] destination. Only the band [rank * local_heads + head_offset, + head_count) for each source rank is written; other output elements remain unchanged.

Calling this function once for every non-overlapping chunk in a schedule whose head counts sum to local_heads reconstructs the ordinary whole-head Ulysses result. It runs on the caller’s current CUDA stream and performs no allocation.

Parameters:

• received (torch.Tensor) – Contiguous gathered head band with shape [B, S_local, world_size * head_count, D].

• world_size (int) – Number of Ulysses ranks represented in received and out.

• local_heads (int) – Number of attention-output heads contributed by each rank to the full output.

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

• out (torch.Tensor) – Preallocated contiguous destination with shape [B, S_local, world_size * local_heads, D]. It must not alias received; elements outside the selected band remain unchanged.

Returns:

out after merging the selected head band.

Return type:

torch.Tensor

Topology Probing and Backend Selection

UlyssesBackendDecision(backend, reason)	Outcome of backend selection.
UlyssesRankTopology(rank, hostname, ...)	Per-rank topology probe result, exchanged across the group.
UlyssesBackendError	Raised when backend="nvlink" is forced but the topology cannot support it.

flashinfer.comm.resolve_ulysses_backend(backend: str = 'auto', group: ProcessGroup | None = None, device: device | str | int | None = None) → UlyssesBackendDecision

Group-consistent backend selection. Must run before any IPC allocation or JIT compilation. It allocates no IPC workspace and compiles nothing; the all_gather_object metadata collectives may themselves stage through CUDA buffers on NCCL process groups.

Collective-safe outcome protocol: every rank participates in the same fixed prefix of all_gather_object calls (at most three; a group-wide explicit NCCL request or an invalid/inconsistent request exits jointly after the first) no matter what fails locally — rank-local errors are encoded as serializable outcomes and re-raised (or turned into an NCCL fallback) jointly after the gather, so no rank can leave the collective sequence early and deadlock its peers. The only uncoordinated failure mode left is the process group itself failing.

Sequence:

1. gather every rank’s requested backend; jointly reject invalid or inconsistent requests. A group-wide explicit "nccl" request returns here, skipping the CUDA/NVML topology probe entirely (the gather itself may stage through CUDA on NCCL groups).

2. gather every rank’s probe outcome (the probe never raises; even a buggy probe implementation is caught into the outcome).

3. every rank evaluates the same pure decision on the same gathered list, catches the result into an outcome, gathers, and cross-checks. Any disagreement conservatively selects NCCL — or raises for backend="nvlink".

Parameters:

• backend (str, default = "auto") – Requested backend policy: "auto", "nvlink", or "nccl".

• group (torch.distributed.ProcessGroup, optional) – Process group whose ranks must make a consistent backend decision. Defaults to torch.distributed.group.WORLD.

• device (torch.device or str or int, optional) – CUDA device to use for probe and metadata collective guards.

Returns:

Group-consistent backend choice and selection or fallback reason.

Return type:

UlyssesBackendDecision

flashinfer.comm.decide_ulysses_backend(requested: str, topologies: List[UlyssesRankTopology], supported_world_sizes: Sequence[int] = (2, 4, 6, 8)) → UlyssesBackendDecision

Pure decision function: gathered per-rank probes -> (backend, reason).

Deterministic in its inputs, so every rank that gathers the same topology list computes the same decision. Conservative: anything unknown, inconsistent, or unverifiable selects NCCL. Raises UlyssesBackendError only when requested == "nvlink" and the NVLink path cannot be used.

Parameters:

• requested (str) – Requested backend policy: "auto", "nvlink", or "nccl".

• topologies (list[UlyssesRankTopology]) – Per-rank topology probe results gathered from the process group.

• supported_world_sizes (Sequence[int], optional) – World sizes for which the fused NVLink kernel is instantiated.

Returns:

Backend choice and selection or fallback reason.

Return type:

UlyssesBackendDecision

flashinfer.comm.probe_ulysses_rank_topology(device: device | str | int | None, rank: int) → UlyssesRankTopology

Probe this rank’s GPU identity and its P2P/NVLink reachability to every other CUDA device visible to this process.

Never raises: the whole probe (including hostname and device resolution) runs inside an exception envelope, so any failure lands in probe_error and the (conservative) decision layer falls back to NCCL.

Parameters:

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

• rank (int) – Rank id to record in the returned topology object.

Returns:

Per-rank topology information, including any probe error instead of raising locally.

Return type:

UlyssesRankTopology

Raw Kernel Entry Points (advanced)

Prefer UlyssesCommunicator; these assume the caller has already verified all-pairs NVLink P2P and owns the IPC workspace lifecycle.

init_ulysses_a2a(out_ipc_ptrs, ...)	Initialize the fused-transpose Ulysses NVLink-P2P all-to-all backend.
dispose_ulysses_a2a(fa)	Release a handle returned by init_ulysses_a2a().
ulysses_a2a(fa, inp, out, B, S_local, H, D, mode)	Fused-transpose Ulysses all-to-all.

MNNVL (Multi-Node NVLink)

Core Classes

MnnvlMemory(mapping, size)
McastGPUBuffer(buf_size, group_size, ...[, ...])	Wrapper class for SymmDeviceMemory to facilitate PyTorch tensor creation.

TensorRT-LLM MNNVL AllReduce

trtllm_mnnvl_all_reduce(inp, ...[, out])	Deprecated pointer-based MNNVL all-reduce API.
trtllm_mnnvl_allreduce(input, workspace, ...)	Perform an MNNVL all-reduce sum across tensor-parallel ranks.
trtllm_mnnvl_fused_allreduce_add_rmsnorm(...)	Performs MNNVL Allreduce + Residual + RMSNorm.
trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant(...)	Perform MNNVL AllReduce + Residual + RMSNorm + FP8/NVFP4 quantization.
trtllm_mnnvl_fused_allreduce_rmsnorm(...)	Performs MNNVL TwoShot Allreduce + RMSNorm.
mpi_barrier()

MNNVL A2A (Throughput Backend)

moe_a2a_initialize(workspace, ep_rank, ...)	Initialize the MoE all-to-all workspace and return a metainfo tensor.
moe_a2a_dispatch(token_selected_experts, ...)	Dispatch tokens and payloads to their target expert ranks.
moe_a2a_combine(payload, local_num_tokens, ...)	Combine per-expert outputs back to the originating ranks.
moe_a2a_sanitize_expert_ids(expert_ids, ...)	Sanitize invalid slots that contain no token routed to this rank by setting their expert IDs to invalid_expert_id.
moe_a2a_get_workspace_size_per_rank(ep_size, ...)	Compute the per-rank workspace size for the MoE all-to-all primitive.
moe_a2a_wrap_payload_tensor_in_workspace(...)	Wrap a slice of the shared workspace as a typed tensor view.

class flashinfer.comm.MoeAlltoAll(mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, workspace_size_per_rank: int | None = None, hidden_size: int | None = None, mnnvl_config: MnnvlConfig | None = None, eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False, *, backend: Literal['trtllm', 'cake'] = 'trtllm')

Bases: object

Manages MoE All-to-All operations with proper workspace allocation and synchronization.

This class provides the throughput-optimized backend that supports multiple payloads per collective operation, explicit dispatch/combine phases, and workspace-backed tensors.

backend="trtllm" is the default, including on Blackwell. To opt in to the generated Blackwell kernels, pass backend="cake" at construction on every rank. Unsupported architectures are rejected, not silently routed to a different backend.

Example

>>> moe_a2a = MoeAlltoAll(mapping, max_num_tokens=2048, top_k=2, num_experts=8)
>>> recv = moe_a2a.dispatch(experts, [hidden, ids, scales], batch_size)
>>> output = moe_a2a.combine(processed, batch_size)

Explicit Blackwell opt-in:

>>> moe_a2a = MoeAlltoAll(mapping, max_num_tokens=2048, top_k=2,
...                      num_experts=8, hidden_size=4096, backend="cake")

__init__(mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, workspace_size_per_rank: int | None = None, hidden_size: int | None = None, mnnvl_config: MnnvlConfig | None = None, eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False, *, backend: Literal['trtllm', 'cake'] = 'trtllm')

Initialize MoeAlltoAll and allocate the shared workspace.

Parameters:

• mapping (Mapping) – Mapping object describing the parallel layout (must expose moe_ep_rank and moe_ep_size).

• max_num_tokens (int) – Maximum number of tokens this rank will dispatch in any single call.

• top_k (int) – Number of experts assigned per token.

• num_experts (int) – Total number of experts (across all ranks).

• workspace_size_per_rank (int, optional) – Pre-computed workspace size in bytes per rank. When None, hidden_size must be provided and the workspace is sized via get_moe_workspace_size_per_rank().

• hidden_size (int, optional) – Hidden dimension size, used to derive workspace_size_per_rank when the latter is omitted.

• mnnvl_config (MnnvlConfig, optional) – Optional configuration for the underlying MNNVL communication backend.

• eplb_stats_num_experts (int) – Number of experts to reserve for the EPLB gathered-stats region. 0 (default) disables EPLB; when non-zero, pass an eplb_local_stats tensor of this length to dispatch().

• enable_rank_mask (bool) – Whether dispatch()/combine() may be called with an active_rank_mask. Fixed for the lifetime of this instance (mirrors the underlying kernel’s compile-time specialization): False (default) compiles out every rank-mask check and forbids passing active_rank_mask.

• backend ({"trtllm", "cake"}) – Kernel implementation, fixed for this instance. "trtllm" preserves the default implementation on all architectures. "cake" explicitly selects generated kernels and requires compute capability 10.0 or 10.3. All participating ranks must select the same backend.

checkpoint_prepare() → None

Unmap MNNVL handles for checkpointing; repeated calls are no-ops.

checkpoint_restore(comm_backend: CommBackend) → None

Remap MNNVL handles after restore; repeated 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.

combine(payload: Tensor, runtime_max_tokens_per_rank: int, payload_in_workspace: bool = False, output_dtype: dtype | None = None, output_scales: Tensor | None = None, output_scalar_scale: float = 1.0, sf_layout: SfLayout = SfLayout.layout_linear, output: Tensor | None = None, *, use_low_precision: bool = False, active_rank_mask: Tensor | None = None) → Tensor

Run the MoE all-to-all combine phase.

Parameters:

• payload (torch.Tensor) – [ep_size, runtime_max_tokens_per_rank, elements_per_token] output payload to scatter back to source ranks.

• runtime_max_tokens_per_rank (int) – Maximum tokens per rank in this batch (same value passed to dispatch()).

• payload_in_workspace (bool) – True if payload is already a workspace-backed view (skips the staging copy). Defaults to False.

• output_dtype (Optional[torch.dtype]) – Optional output data type. Currently supports torch.bfloat16, torch.float8_e4m3fn, and torch.uint8 (packed fp4).

• output_scales (Optional[torch.Tensor]) – Optional output scale tensor for quantized outputs. Currently supports UE8M0 (packed in torch.uint8) with vector size 32.

• output_scalar_scale (float) – Per-tensor global scale applied before FP4 block scaling (NVFP4 SFScaleVal). Defaults to 1.0; ignored by MXFP8/MXFP4 paths.

• sf_layout (SfLayout) – Output swizzle layout. Defaults to SfLayout.layout_linear.

• output (Optional[torch.Tensor]) – Caller-provided contiguous output tensor. Its shape and dtype must match the requested combine output, and it must be on the same device as payload.

• use_low_precision (bool) – If True, quantize the recv-buffer payload to FP8 (e4m3) before accumulating; the combine upcasts to a bf16 output.

• active_rank_mask (torch.Tensor, optional) – CPU uint64 tensor of shape [MOE_A2A_RANK_MASK_WORDS]. Should match the mask passed to the preceding dispatch() call (or be omitted from both). Requires the instance to have been constructed with enable_rank_mask=True.

Returns:

[local_num_tokens, elements_per_token] combined tensor.

Return type:

torch.Tensor

dispatch(token_selected_experts: Tensor, input_payloads: list[Tensor], runtime_max_tokens_per_rank: int, invalid_token_expert_id: int | None = None, expert_id_payload_index: int | None = None, eplb_local_stats: Tensor | None = None, active_rank_mask: Tensor | None = None) → list[Tensor]

Run the MoE all-to-all dispatch phase.

Parameters:

• token_selected_experts (torch.Tensor) – [local_num_tokens, top_k] int32 tensor of expert assignments.

• input_payloads (list[torch.Tensor]) – Per-token payload tensors, each shaped [local_num_tokens, *].

• runtime_max_tokens_per_rank (int) – Maximum tokens per rank in this batch. Must be <= max_num_tokens used at construction.

• invalid_token_expert_id (int, optional) – If supplied, expert IDs not owned by the current rank are rewritten to this value. Requires expert_id_payload_index.

• expert_id_payload_index (int, optional) – Index into input_payloads that holds the expert IDs to sanitize. Required when invalid_token_expert_id is set.

• eplb_local_stats (torch.Tensor, optional) – [eplb_stats_num_experts] int32 tensor of this rank’s local EPLB statistics to all-gather during dispatch. Requires the instance to have been constructed with eplb_stats_num_experts set. The gathered result is available afterwards via eplb_gathered_stats.

• active_rank_mask (torch.Tensor, optional) – CPU uint64 tensor of shape [MOE_A2A_RANK_MASK_WORDS] (see moe_a2a_active_rank_mask()). Tokens routed to a masked-off rank are dropped instead of hanging the collective. Requires the instance to have been constructed with enable_rank_mask=True.

Returns:

Workspace-backed receive tensors, one per input_payloads entry, each shaped [ep_size, runtime_max_tokens_per_rank, *].

Return type:

list[torch.Tensor]

property eplb_gathered_stats: Tensor | None

Gathered EPLB stats from the most recent dispatch().

Workspace-backed [ep_size, eplb_stats_num_experts] int32 view (row r holds rank r’s eplb_local_stats), or None when EPLB is disabled or no eplb_local_stats was passed. Valid only between dispatch() and combine() (combine resets state).

get_combine_payload_tensor_in_workspace(runtime_max_tokens_per_rank: int, hidden_size: int, dtype: dtype) → Tensor

Return a workspace-backed view to use as the combine payload.

Zero-copy variant of combine(): experts can write directly into the returned tensor and call combine() with payload_in_workspace=True. Must be called after a successful dispatch() and before combine().

Parameters:

• runtime_max_tokens_per_rank (int) – Maximum tokens per rank in this batch.

• hidden_size (int) – Hidden dimension size.

• dtype (torch.dtype) – Element dtype of the resulting view.

Returns:

[ep_size, runtime_max_tokens_per_rank, hidden_size] workspace-backed tensor.

Return type:

torch.Tensor

Raises:

RuntimeError – If called before a successful dispatch().

static get_moe_workspace_size_per_rank(ep_size: int, top_k: int, max_num_tokens: int, hidden_size: int, extra_payload_bytes_per_token: int = 0, eplb_stats_num_experts: int = 0, *, backend: Literal['trtllm', 'cake'] = 'trtllm') → int

Compute the per-rank workspace size for the MoE all-to-all primitive.

Convenience wrapper around moe_a2a_get_workspace_size_per_rank() that derives the dispatch / combine payload sizes from hidden_size and top_k assuming 16-bit hidden states. For a tighter bound on quantized models use moe_a2a_get_workspace_size_per_rank() directly.

backend="trtllm" is the default; pass backend="cake" when sizing a workspace for an opt-in MoeAlltoAll instance.

Parameters:

• ep_size (int) – Total expert-parallel world size.

• top_k (int) – Number of experts assigned per token.

• max_num_tokens (int) – Maximum number of tokens across all ranks.

• hidden_size (int) – Hidden dimension size.

• extra_payload_bytes_per_token (int) – Extra payload bytes per token to reserve (e.g. for quantization scales). Defaults to 0.

• eplb_stats_num_experts (int) – Number of experts reserved for the EPLB gathered-stats region (0 disables it).

• backend (MoeAlltoAllBackend) – Communication backend used to size the workspace. Defaults to "trtllm"; pass "cake" for the Cake backend.

Returns:

Required workspace size per rank, in bytes.

Return type:

int

classmethod get_workspace(workspace_size_per_rank: int, ep_rank: int, ep_size: int, max_num_tokens: int, mapping: Mapping, eplb_stats_num_experts: int = 0, *, backend: Literal['trtllm', 'cake'] = 'trtllm') → dict

Allocate/cache a workspace for the explicitly selected backend.

All ranks must choose the same backend; defaults to "trtllm". Workspaces belonging to different backends are never shared.

MoeAlltoAll preserves its CUDA virtual addresses across process checkpoint/restore. After quiescing all work, call checkpoint_prepare to release the non-checkpointable physical MNNVL handles. Then call checkpoint_restore with a fresh communication backend before replaying a captured CUDA graph:

moe_alltoall.checkpoint_prepare()
moe_alltoall.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 physical handle unmapping or remapping begins, do not retry or reuse the workspace; restart the affected rank.

MoeAlltoAll.checkpoint_prepare()	Unmap MNNVL handles for checkpointing; repeated calls are no-ops.
MoeAlltoAll.checkpoint_restore(comm_backend)	Remap MNNVL handles after restore; repeated calls are no-ops.

DCP All-to-All (Context-Parallel Attention Reduction)

decode_cp_a2a_workspace_size(cp_size)	Return the workspace size (in bytes) per rank for the given CP group size.
decode_cp_a2a_allocate_mnnvl_workspace(...)	Allocate an MNNVL-backed workspace of shape [cp_size, ws_elems_per_rank].
decode_cp_a2a_init_workspace(workspace, ...)	Initialize the workspace FIFO buffers (call once before the first alltoall).
decode_cp_a2a_alltoall(partial_o, ...[, ...])	Perform the DCP all-to-all exchange.

NCCL LSA DCP All-to-All + LSE Reduce

The fused path uses PyTorch NCCL symmetric memory and requires every context- parallel rank to be in one load/store-accessible NVLink domain. It does not use the MNNVL workspace accepted by decode_cp_a2a_alltoall.

decode_cp_a2a_lse_reduce_workspace_size(...)	Return the required NCCL symmetric workspace size in bytes.
decode_cp_a2a_lse_reduce_create_workspace(...)	Create and rendezvous the fused op's NCCL symmetric workspace.
decode_cp_a2a_lse_reduce(partial_o, ...[, ...])	Fuse an NCCL LSA DCP A2A exchange with the LSE-weighted reduce.

Mixed Communication

MixedCommOp(value)	Enumeration of mixed communication operation types.
MixedCommMode(value)	Enumeration of mixed communication execution modes.
MixedCommHandler(world_rank, world_size, ...)	An implementation for the combinations of all-reduce + all-gather and reduce-scatter + all-reduce.
run_mixed_comm(op, handler, x_in[, x_out, mode])	Execute a mixed communication operation.