{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# **Dispatch&Combine 算子在 MoE 模型中的作用**\n",
"\n",
"**导航**：[← 上一章：MoE并行策略](02.03_parallel_strategy.ipynb) | [返回课程主页](02.01_chapter_intro.ipynb) | [下一章：输入输出对应关系 →](02.05_operator_logic_overview.ipynb)\n",
"\n",
"---\n",
"\n",
"本节聚焦 MoE（Mixture of Experts）模型中 Dispatch、Combine 算子的核心作用：理解 Dispatch、Combine 在 MoE 网络中的定位。\n",
"\n",
"**学习目标**：理解 Dispatch、Combine 算子在 MoE 网络架构中的角色及其必要性。\n",
"\n",
"**本章目录**（点击跳转）：\n",
"- [MoE 网络架构回顾](#1-moe-网络架构回顾)\n",
"- [Dispatch 算子的作用](#2-dispatch-算子的作用)\n",
"- [Dispatch 在 MoE 数据流中的位置](#3-dispatch-在-moe-数据流中的位置)\n",
"- [Token 按 Expert 分组的必要性](#4-token-按-expert-分组的必要性)\n",
"- [Combine 算子的核心作用](#5-combine-算子的核心作用)\n",
"- [Dispatch 与 Combine 的配合关系](#6-dispatch-与-combine-的配合关系)\n",
"- [小结](#7-小结)\n",
"- [课后练习](#课后练习)\n",
"\n",
"---\n",
"## **1. MoE 网络架构回顾**\n",
"\n",
"MoE（Mixture of Experts）模型的核心思想是：**每个 Token 只被路由到少数几个专家进行计算，而不是所有专家**。这使得模型可以在保持计算量可控的情况下大幅增加参数规模。\n",
"\n",
"> **前置知识**：本章需要了解 MoE 基本架构和并行策略，建议先阅读 [02.02 MoE架构概述](02.02_moe_architexture_overview.ipynb) 和 [02.03 MoE并行策略](02.03_parallel_strategy.ipynb)\n",
"\n",
"**MoE 层的基本结构**：\n",
"\n",
"<img src=\"./images/02.04_Moe_structure.png\" width=\"1200\">\n",
"\n",
"**Gate Network 的作用**：为每个 Token 计算对所有专家的亲和度分数，选出 Top-K 个专家作为该 Token 的路由目标。\n",
"\n",
"**问题**：Gate 决定了每个 Token 应该去哪些专家，但如何把这个决策落地执行？这就是 Dispatch 算子存在的意义。\n",
"\n",
"---\n",
"## **2. Dispatch 算子的作用**\n",
"\n",
"### **2.1 核心作用：执行路由决策**\n",
"\n",
"Dispatch 算子是 MoE 网络中**执行路由决策的算子**。它的核心作用是：\n",
"\n",
"**将 Gate Network 输出的路由决策转化为实际的 Token 分发操作。**\n",
"\n",
"<img src=\"./images/02.04_token_distribution.png\" width=\"1000\">\n",
"\n",
"---\n",
"### **2.2 为什么需要 Dispatch？**\n",
"\n",
"如果没有 Dispatch 算子，MoE 网络面临两个根本性问题：\n",
"\n",
"**问题一：Token 到专家的映射如何执行？**\n",
"\n",
"Gate 只输出「Token i 应去专家 E」，但没有实际执行这个映射。Token 数据本身还在原始位置，需要被搬运到专家处。\n",
"\n",
"**问题二：分布式场景下 Token 如何跨卡发送？**\n",
"\n",
"在专家并行场景下，不同专家分布在不同 NPU 上。Token 需要跨 NPU 发送到目标专家所在的卡。\n",
"\n",
"**Dispatch 的作用就是解决这两个问题**：\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">问题</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Dispatch 的解决方案</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Token 映射执行</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">按 Gate 的 <code>expertIds</code> 将 Token 发往对应专家</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">分布式跨卡发送</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">通过 AllToAll 通信将 Token 发送到目标专家所在 NPU</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"---\n",
"### **2.3 Dispatch 的具体作用**\n",
"\n",
"在 MoE 网络中，Dispatch 算子承担以下具体作用：\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">作用</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">说明</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>路由执行</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将 Gate 的 TopK 路由决策转化为实际 Token 分发</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>Token 分组</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将同一专家处理的 Token 连续排列，便于专家批量计算</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>跨卡分发</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">在专家并行场景下，通过 AllToAll 将 Token 发往目标专家所在卡</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>记录元信息</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">记录 Token 来源位置，供后续 Combine 还原使用</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"---\n",
"## **3. Dispatch 在 MoE 数据流中的位置**\n",
"\n",
"<img src=\"./images/02.04_dispatch_position.png\" width=\"1500\">\n",
"\n",
"**Dispatch 的上游**：Gate Network（提供路由决策 `expertIds`）\n",
"\n",
"**Dispatch 的下游**：Expert FFN（接收分发后的 Token 进行计算）\n",
"\n",
"---\n",
"## **4. Token 按 Expert 分组的必要性**\n",
"\n",
"Dispatch 的一个重要作用是将 Token 按目标专家重新分组排列。\n",
"\n",
"**原始状态**：Token 按序列顺序排列\n",
"\n",
"**Dispatch 后**：Token 按目标专家分组排列\n",
"\n",
"<img src=\"./images/02.04_grouped_by_expert.png\" width=\"600\">\n",
"\n",
"**为什么需要分组？**\n",
"\n",
"专家 FFN 是矩阵运算。如果同一专家的 Token 连续排列，专家可以一次性批量计算所有 Token，效率最大化。\n",
"\n",
"**不分组**：专家需要逐个处理分散的 Token，内存访问不连续，计算效率低。\n",
"\n",
"**分组后**：专家一次矩阵乘法处理所有 Token `[N_tokens, h] × [h, H]`，效率最大化。\n",
"\n",
"---\n",
"\n",
"## **5. Combine 算子的核心作用**\n",
"\n",
"### **5.1 聚合专家输出，还原 Token 序列**\n",
"\n",
"Combine 算子是 MoE 网络中**聚合专家输出的算子**。\n",
"\n",
"Token i 被分发到 K 个专家，产生 K 路输出。Combine 收集这 K 路输出，按 Gate 给出的权重加权融合，还原为 Token i 的最终输出。\n",
"\n",
"<img src=\"./images/02.04_combine.png\" width=\"800\">\n",
"\n",
"---\n",
"### **5.2 解决三个核心问题**\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">问题</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Combine 的解决方案</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">专家输出分散</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">通过 AllToAll 收集各专家输出</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">输出顺序需要还原</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">使用 Dispatch 记录的元信息还原 Token 顺序</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">多路输出需要融合</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">按 expertScales 对 K 路输出加权求和</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"---\n",
"### **5.3 TopK 加权求和**\n",
"\n",
"MoE 的核心思想：每个 Token 被 K 个专家处理，最终输出是 K 个专家输出的加权融合。\n",
"\n",
"**加权求和的作用**：\n",
"\n",
"- 高权重的专家对该 Token 的输出贡献更大\n",
"- 权重由 Gate Network 计算，反映 Token 对各专家的亲和度\n",
"- 权重归一化（Σwk = 1）保证输出尺度稳定\n",
"\n",
"---\n",
"\n",
"## **6. Dispatch 与 Combine 的配合关系**\n",
"\n",
"Dispatch 和 Combine 是 MoE 网络中对偶的两个算子，形成完整的 Token 路由流程。\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">算子</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">作用</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">方向</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>Dispatch</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Token → Expert（分发）</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将 Token 发往专家</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\"><strong>Combine</strong></td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Expert → Token（聚合）</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将专家输出还原到 Token</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"**配合机制**：\n",
"\n",
"<img src=\"./images/02.04_D&C cooperation mechanism.png\" width=\"1000\">\n",
"\n",
"Dispatch 在分发 Token 时，记录每个 Token 的原始位置信息（三元组<code>(rankId(本卡索引), tokenIndex(token索引), topkIdx(topK索引))</code>）。Combine 使用这些元信息将专家输出还原到正确的 Token 位置。\n",
"\n",
"---\n",
"## **7. 小结**\n",
"\n",
"Dispatch 算子在 MoE 网络中的作用总结：\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">作用</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">说明</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">执行路由决策</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将 Gate Network 的 TopK 路由决策转化为实际 Token 分发</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">Token 分组</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将同一专家处理的 Token 连续排列，使专家可以批量高效计算</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">跨卡分发</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">在专家并行场景下，通过 AllToAll 将 Token 发往目标专家所在 NPU</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">记录元信息</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">为 Combine 的还原操作提供必要的位置信息</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"Dispatch 是 MoE 网络中\"路由决策落地\"的关键算子，是 Token 从 Gate 到 Expert 的桥梁。\n",
"\n",
"Combine 算子在 MoE 网络中的作用：\n",
"\n",
"<table style=\"margin: 0; margin-right: auto; border-collapse: collapse;\">\n",
" <tr>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">作用</th>\n",
" <th style=\"border:1px solid #ccc; padding:8px; text-align:left;\">说明</th>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">聚合专家输出</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">收集各专家的 FFN 输出，从分散状态变为集中状态</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">还原 Token 顺序</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">将按 Expert 分组的输出还原为按 Token 分组</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">加权融合</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">对每个 Token 的 K 路专家输出按 Gate 权重加权求和</td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">输出组装</td>\n",
" <td style=\"border:1px solid #ccc; padding:8px; text-align:left;\">组装 MoE 层的最终输出，传递给下一层网络</td>\n",
" </tr>\n",
"</table>\n",
"\n",
"Combine 是 MoE 网络中\"专家输出聚合\"的关键算子。\n",
"\n",
"---\n",
"# **课后练习**\n",
"\n",
"1. （单选题）Dispatch 算子在 MoE 网络中的核心作用是什么？\n",
" A. 计算每个 Token 对各专家的亲和度分数。\n",
" B. 执行 Gate 的路由决策，将 Token 分发到对应专家。\n",
" C. 聚合各专家的输出。\n",
" D. 执行专家的 FFN 计算。\n",
"\n",
"2. （单选题）Dispatch 将 Token 按 Expert 分组排列的目的是什么？\n",
" A. 减少内存使用。\n",
" B. 使专家可以批量矩阵计算，提高效率。\n",
" C. 减少模型参数量。\n",
" D. 简化 Gate 计算。\n",
"\n",
"3. （单选题）在专家并行场景下，Dispatch 使用什么方式将 Token 发送到目标专家所在 NPU？\n",
" A. Broadcast。\n",
" B. AllReduce。\n",
" C. AllToAll。\n",
" D. Scatter。\n",
"\n",
"4. （单选题）Dispatch 记录的元信息（Token 原始位置）的作用是什么？\n",
" A. 供 Gate Network 使用。\n",
" B. 供 Expert FFN 使用。\n",
" C. 供 Combine 还原 Token 顺序使用。\n",
" D. 供下一层 MoE 使用。\n",
"\n",
"5. （单选题）Combine 对每个 Token 的 K 路专家输出做什么操作？\n",
" A. 简单拼接。\n",
" B. 按 Gate 权重加权求和。\n",
" C. 取最大值。\n",
" D. 取平均值。\n",
"\n",
"6. （单选题）Combine 使用什么方式收集各专家的输出？\n",
" A. Broadcast。\n",
" B. AllReduce。\n",
" C. AllToAll。\n",
" D. Gather。\n",
"\n",
"7. （多选题）关于 Dispatch 在 MoE 网络中的作用，以下说法正确的是？\n",
" A. Dispatch 执行 Gate Network 的路由决策。\n",
" B. Dispatch 使同一专家处理的 Token 连续排列，便于专家批量计算。\n",
" C. Dispatch 与 Combine 形成对偶操作：Dispatch 分发，Combine 聚合。\n",
" D. Dispatch 的上游是 Gate Network，下游是 Expert FFN。"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# **查看答案**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"vscode": {
"languageId": "plaintext"
}
},
"outputs": [],
"source": [
"!cat answer/02.04_answer.txt"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"---\n",
"\n",
"**导航**：[← 上一章：MoE并行策略](02.03_parallel_strategy.ipynb) | [返回课程主页](02.01_chapter_intro.ipynb) | [下一章：输入输出对应关系 →](02.05_operator_logic_overview.ipynb)\n",
"\n",
"---"
]
}
],
"metadata": {
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
