flashinfer.comm.moe_a2a_dispatch

flashinfer.comm.moe_a2a_dispatch(token_selected_experts: Tensor, input_payloads: list[Tensor], workspace: Tensor, metainfo: Tensor, runtime_max_tokens_per_rank: int, ep_rank: int, ep_size: int, top_k: int, num_experts: int, enable_pdl: bool | None = None, eplb_local_stats: Tensor | None = None, enable_rank_mask: bool = False, active_rank_mask: Tensor | None = None, recv_view_cache: dict | None = None)

Dispatch tokens and payloads to their target expert ranks.

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, *].

  • workspace (torch.Tensor) – [ep_size, size_per_rank] shared workspace.

  • metainfo (torch.Tensor) – Metainfo tensor returned by moe_a2a_initialize().

  • runtime_max_tokens_per_rank (int) – Maximum tokens per rank for this batch (must be <= the max_num_tokens used at initialize time).

  • ep_rank (int) – Current expert-parallel rank.

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

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

  • num_experts (int) – Total number of experts.

  • enable_pdl (Optional[bool]) – Whether to use programmatic dependent launch. None auto-detects from the device.

  • eplb_local_stats (Optional[torch.Tensor]) – Optional [eplb_stats_num_experts] int32 tensor of this rank’s local EPLB statistics. When provided, the dispatch all-gathers it across ranks and returns the result as eplb_gathered_stats. The length must match the eplb_stats_num_experts passed to moe_a2a_initialize().

  • enable_rank_mask (bool) – Whether to instantiate the kernel variant that checks active_rank_mask at all. False (default) compiles out every rank-mask check for the common no-fault-tolerance case and requires active_rank_mask to be omitted.

  • active_rank_mask (Optional[torch.Tensor]) – Optional CPU uint64 tensor of shape [MOE_A2A_RANK_MASK_WORDS] (see moe_a2a_active_rank_mask()). Bit i set means rank i is alive and participates in this collective; tokens routed to a masked-off rank are dropped instead of hanging the collective. Requires enable_rank_mask=True; the local ep_rank’s own bit must always be set.

  • recv_view_cache (dict, optional) – Opaque cache of receive views. Pass an empty dictionary to enable caching; None (default) disables it and preserves the previous behaviour exactly. The cache retains a reference to workspace and is cleared when a different workspace tensor is passed. Keep the backing allocation alive and do not change the workspace or the cached views’ shape, strides, or storage in place. Clear the dictionary to release its references; otherwise leave its contents unchanged.

Returns:

(output_payloads, combine_payload_offset, eplb_gathered_stats). output_payloads is a list of workspace-backed views, one per input_payloads entry, that contains the data routed to this rank. combine_payload_offset is the workspace offset reserved for the matching moe_a2a_combine() call. eplb_gathered_stats is a workspace-backed [ep_size, eplb_stats_num_experts] int32 view (row r holds rank r’s eplb_local_stats) when eplb_local_stats was provided, else None.

Return type:

Tuple[list[torch.Tensor], int, Optional[torch.Tensor]]