flashinfer.moe_ep.MoEEpCommunication

class flashinfer.moe_ep.MoEEpCommunication(bootstrap: BootstrapConfig, params: MoEEpCommParams)

Bases: ABC

Self-contained MoE expert-parallel dispatch/combine backend.

The peer of the Fleet/Handle transports; see the module docstring.

One instance serves one EP group and one in-flight dispatch/combine pair at a time: every dispatch() must be followed by exactly one combine() before the next dispatch(). All ranks of the group must issue the same sequence of calls.

abstract classmethod is_platform_supported() → bool

Whether this backend can run on the current machine.

Callable before construction; performs no collective work.

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

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

get_combine_input_buffer(dtype: torch.dtype) → torch.Tensor | None

Backend-owned buffer the expert computation may write its output into.

Passing the returned tensor to combine() saves a copy. Valid only between dispatch() and combine(). Backends without such a buffer return None.

destroy() → None

Release backend resources. Collective on every rank, idempotent.