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, *, backend: Literal['trtllm', 'cake'] = 'trtllm', 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. Masking a peer does not advance that peer’s transport epoch. Before a skipped peer rejoins, all ranks must quiesce and coordinate workspace reinitialization; merely restoring its mask bit is insufficient.

  • backend ({"trtllm", "cake"}) – Defaults to "trtllm". Pass "cake" to opt in on CC 10.0/10.3; must match workspace initialization, combine, and all peer ranks.

  • recv_view_cache (dict, optional) – Opaque cache of receive views. Pass an empty dictionary to enable caching; None (default) disables it. 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 workspace or cached-view 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]]