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

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_custom_all_reduce(inp, out, tp_size, ...)	Parameters: - inp: the input tensor. [token_num, hidden_dim] - out: the output tensor. [token_num, hidden_dim] - tp_size: the size of the process group. - tp_rank: the rank of the current process. - token_num: the number of tokens in the sequence. - fusion_op_code: the fusion operation code. - strategy_code: the strategy code. - config_code: the config code. - launch_with_pdl: whether to launch with pdl. - flag_value: monotonic counter, same value on all ranks for a given call. Calls that dispatch to the lamport kernel (see is_lamport_supported in the kernel header) need consecutive flag values across those calls: keep a separate counter for them, or re-run trtllm_lamport_initialize_all before each one. Sharing one counter with other calls silently corrupts the lamport output. - peer_comm_buffer_ptrs: the peer communication buffer pointers. - peer_barrier_ptrs_in: the peer barrier pointers in. - peer_barrier_ptrs_out: the peer barrier pointers out. - bias: the bias tensor. [hidden_dim] - residual: the residual tensor. [token_num, hidden_dim] - weight: the weight tensor. [hidden_dim] - weight_pre_residual_norm: the weight pre residual norm tensor. [hidden_dim] - eps: the epsilon value. - intermediate_buffer: the intermediate buffer tensor. - lamport_peer_comm_buffer_ptrs_0: the lamport peer communication buffer pointers 0. - lamport_peer_comm_buffer_ptrs_1: the lamport peer communication buffer pointers 1. - lamport_peer_comm_buffer_ptrs_2: the lamport peer communication buffer pointers 2.
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(...)).
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(...)).

Workspace Management

trtllm_create_ipc_workspace_for_all_reduce(...)	Parameters: - rank: the rank of the current process.
trtllm_create_ipc_workspace_for_all_reduce_fusion(...)	Parameters: - tp_rank: the rank of the current process.
trtllm_destroy_ipc_workspace_for_all_reduce(...)	Destroy a workspace created by trtllm_create_ipc_workspace_for_all_reduce.
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.

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.

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]

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

gather_heads(x: Tensor) → 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].

Returns:

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

Return type:

torch.Tensor

scatter_heads(x: Tensor) → 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].

Returns:

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

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: int | str | device | 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: int | str | device | 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, hidden_size: int = None, mnnvl_config: MnnvlConfig | None = None, eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False)

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.

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)

__init__(mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, workspace_size_per_rank: int = None, hidden_size: int = None, mnnvl_config: MnnvlConfig | None = None, eplb_stats_num_experts: int = 0, enable_rank_mask: bool = False)

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.

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

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

Returns:

Required workspace size per rank, in bytes.

Return type:

int

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.

Mixed Communication

MixedCommOp(value[, names, module, ...])	Enumeration of mixed communication operation types.
MixedCommMode(value[, names, module, ...])	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.