flashinfer.comm.MoeAlltoAll

class flashinfer.comm.MoeAlltoAll(mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, workspace_size_per_rank: int = None, hidden_size: int = None, mnnvl_config: MnnvlConfig | None = None)

管理 MoE 全到全操作，并进行适当的工作区分配和同步。

此类提供吞吐量优化的后端，支持每个集体操作中的多个有效负载、显式调度/合并阶段以及基于工作区的张量。

示例

>>> moe_a2a = MoeAlltoAll(mapping, max_num_tokens=2048, top_k=2, num_experts=8)
>>> recv = moe_a2a.dispatch(experts, [hidden, ids, scales], batch_size)
>>> output = moe_a2a.combine(processed, batch_size)
__init__(mapping: Mapping, max_num_tokens: int, top_k: int, num_experts: int, workspace_size_per_rank: int = None, hidden_size: int = None, mnnvl_config: MnnvlConfig | None = None)

使用工作区分配初始化 MoeAlltoAll。

参数:

  • mapping – 包含等级信息的映射对象

  • max_num_tokens – 支持的最大 token 数量

  • top_k – 每个 token 的专家数量

  • num_experts – 专家总数

  • workspace_size_per_rank – 每个等级的工作区大小（字节），如果为 None，则必须提供 hidden_size

  • hidden_size – 用于计算工作区大小的隐藏维度大小，如果未提供 workspace_size_per_rank

  • mnnvl_config – 用于配置 MNNVL 内存对象的通信后端

方法

__init__(mapping, max_num_tokens, top_k, ...)

使用工作区分配初始化 MoeAlltoAll。

combine(payload, runtime_max_tokens_per_rank)

执行 MoE 全到全合并操作。

dispatch(token_selected_experts, ...[, ...])

执行 MoE 全到全调度操作。

get_combine_payload_tensor_in_workspace(...)

获取由工作区支持的合并有效负载张量（零拷贝）。

get_moe_workspace_size_per_rank(ep_size, ...)

方便的包装器，用于计算 MoE 全到全操作的每个等级的工作区大小。

get_workspace(workspace_size_per_rank, ...)