flashinfer.moe_ep.NVLinkTwoSidedAlltoAll

class flashinfer.moe_ep.NVLinkTwoSidedAlltoAll(bootstrap: BootstrapConfig, params: MoEEpCommParams, config: NVLinkTwoSidedConfig | None = None)

Bases: MoEEpCommunication

Rank-major dispatch/combine with NVLink two-sided all-to-all-v kernels.

Requires num_experts divisible by 4. EPLB statistics, when passed to dispatch(), must cover all num_experts experts.

classmethod is_platform_supported() → bool

Whether this backend can run on the current machine.

Callable before construction; performs no collective work.

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

Send each token to the ranks that own its selected experts.

Parameters:

  • hidden_states (torch.Tensor) – [local_num_tokens, *] token rows, unquantized or quantized.

  • topk_ids (torch.Tensor) – [local_num_tokens, top_k] global expert ids.

  • topk_weights (torch.Tensor, optional) – [local_num_tokens, top_k] routing weights, forwarded with the tokens so the expert computation can apply them.

  • hidden_states_scale (torch.Tensor, optional) – [local_num_tokens, *] per-token scale factors of quantized hidden_states, forwarded row-aligned with them.

  • runtime_max_tokens_per_rank (int, optional) – Largest local_num_tokens over all ranks for this step, at most MoEEpCommParams.max_tokens_per_rank; it sets tokens_per_rank of the result. Every rank must pass the same value. None uses MoEEpCommParams.max_tokens_per_rank, which keeps shapes static (e.g. for CUDA graphs).

  • eplb_local_stats (torch.Tensor, optional) – This rank’s expert-load statistics, all-gathered alongside the dispatch by backends that support it.

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

Return each token’s expert outputs, summed over ranks.

Parameters:

  • expert_output (torch.Tensor) – [ep_size * tokens_per_rank, hidden] (or the equivalent [ep_size, tokens_per_rank, hidden]) per-row outputs of the expert computation, already weighted and reduced over this rank’s local experts.

  • output (torch.Tensor, optional) – Preallocated [local_num_tokens, hidden] result buffer.

destroy() → None

Release backend resources. Collective on every rank, idempotent.