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)

将 token 和 payload 分发到专家 rank。

参数:

• token_selected_experts – [local_num_tokens, top_k] int32 张量

• input_payloads – 要分发的 [local_num_tokens, *] 张量列表

• workspace – [ep_size, size_per_rank] 工作区张量

• metainfo – 来自 initialize 的元数据张量

• runtime_max_tokens_per_rank – 此批次中每个 rank 的最大 token 数

• ep_rank – 当前专家并行 rank

• ep_size – 总专家并行大小

• top_k – 每个 token 的专家数量

• num_experts – 专家总数

返回值:

此 rank 的 payload 列表，由工作区中的数据支持 combine_payload_offset：将 combine payload 放置在工作区中的偏移量

返回值类型:

output_payloads