flashinfer.moe_ep.NVLinkOneSidedAlltoAll¶

class flashinfer.moe_ep.NVLinkOneSidedAlltoAll(bootstrap: BootstrapConfig, params: MoEEpCommParams, config: NVLinkOneSidedConfig | None = None)¶

Bases: MoEEpCommunication

Rank-major dispatch/combine with NVLink one-sided put/get kernels.

hidden_states, its optional scale factors, topk_ids and topk_weights travel as payloads of a single dispatch. Received rows beyond each source rank’s token count get invalid_expert_id routing; after a Lamport step, row source_rank * tokens_per_rank + token holds that token and rows of tokens not routed here get invalid_expert_id. See NVLinkOneSidedConfig for CFT counted writes, the Lamport transport and timeouts.

classmethod is_platform_supported() → bool¶

Whether this backend can run on the current machine.

Callable before construction; performs no collective work.

active_rank_mask(active_ranks: Sequence[int]) → torch.Tensor¶: CPU uint64 bitmask of active_ranks for active_rank_mask.

property lamport_enabled: bool¶: Whether the Lamport transport is available to this instance’s steps.

dispatch(hidden_states: torch.Tensor, topk_ids: torch.Tensor, topk_weights: torch.Tensor | None = None, *, hidden_states_scale: torch.Tensor | None = None, runtime_max_tokens_per_rank: int | None = None, eplb_local_stats: torch.Tensor | None = None, active_rank_mask: torch.Tensor | None = None) → MoEEpDispatchResult¶

See MoEEpCommunication.dispatch().

active_rank_mask is a CPU uint64 mask from active_rank_mask(); tokens routed to a masked-off rank are dropped. Requires enable_rank_mask=True.

A step takes the Lamport transport only if its runtime_max_tokens_per_rank (the capacity when None) is at most lamport_max_tokens; NVLinkOneSidedConfig lists the other conditions.

Every rank passes topk_weights in a step or none does (a rank without tokens passes an empty [0, top_k] tensor, not None), as every rank passes the same runtime_max_tokens_per_rank. A mismatch is not detected: on the fence path a rank that passes weights reads, for a peer that passed none, whatever an earlier step left in its weight rows; a Lamport step can also leave a stale weight row for a later step, or time out (timeout_sec).

get_combine_input_buffer(dtype: torch.dtype) → torch.Tensor¶: [ep_size * tokens_per_rank, hidden_size] view of the combine input region of this rank’s workspace.

combine(expert_output: torch.Tensor, *, output: torch.Tensor | None = None, active_rank_mask: torch.Tensor | None = None) → torch.Tensor¶

See MoEEpCommunication.combine().

active_rank_mask must match the mask passed to dispatch().

checkpoint_prepare() → None¶

Release the workspace’s physical memory and CFT endpoints before a process checkpoint.

A CUDA process checkpoint cannot save the workspace’s MNNVL physical memory, which peer processes import, nor its CFT logical endpoints. This method releases both but keeps the workspace’s virtual addresses reserved, so CUDA graphs captured on the workspace replay correctly once checkpoint_restore() has backed it with new memory:

communication.checkpoint_prepare()
# checkpoint and restore the process
communication.checkpoint_restore(comm_backend)

Collective over the EP group: every rank calls it in the same order, after all work on the workspace has completed and only between rounds. Instances sharing the workspace may call it again, which does nothing. If it raises once unmapping has begun, the workspace cannot be reused: restart the rank rather than retry.

checkpoint_restore(comm_backend: CommBackend) → None¶

Back the workspace with new memory after checkpoint_prepare(), re-arm its Lamport regions and recreate its CFT endpoints.

Collective over comm_backend, a fresh backend that must span the EP group with the original ranks and world size. Every rank calls it in the same order. Does nothing when the workspace is mapped. The CFT endpoints are recreated under the IDs they had, which CUDA graphs captured before the checkpoint use; raises if that is not possible. If it raises once remapping has begun, the workspace cannot be reused: restart the rank rather than retry.

destroy() → None¶

Release this instance’s share of the workspace.

The last instance on a workspace frees it, together with its CFT endpoints; every rank must have finished its last combine on it.