flashinfer.comm.moe_a2a_get_workspace_size_per_rank

flashinfer.comm.moe_a2a_get_workspace_size_per_rank(ep_size: int, max_num_tokens: int, total_dispatch_payload_size_per_token: int, combine_payload_size_per_token: int, eplb_stats_num_experts: int = 0, *, backend: Literal['trtllm', 'cake'] = 'trtllm')

Compute the per-rank workspace size for the MoE all-to-all primitive.

backend="trtllm" is the default. Pass backend="cake" to size a workspace for the explicitly selected Blackwell implementation.

Parameters:

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

  • max_num_tokens (int) – Maximum number of tokens across all ranks.

  • total_dispatch_payload_size_per_token (int) – Sum (in bytes) of all per-token payloads sent during the dispatch phase.

  • combine_payload_size_per_token (int) – Per-token payload size (in bytes) sent back during the combine phase.

  • eplb_stats_num_experts (int) – Number of experts reserved for the EPLB gathered-stats region (0 disables it).

  • backend (MoeAlltoAllBackend) – Communication backend used to size the workspace. Defaults to "trtllm"; pass "cake" for the Cake backend.

Returns:

Required workspace size per rank, in bytes.

Return type:

int