[Doc] Add docs for mega moe (vllm-project#12926)

nanshen3000 · web-flow · commit 019d3e944e42 · 2026-07-29T14:07:19.000+08:00
### What this PR does / why we need it? Update docs for fused mc2. Now mega_moe is already supported in vllm-ascend for A2/A3. To let users know about this change, the docs should also be updated. This PR adds the description of the mega_moe operator to the original dispatch_ffn_combine document and comments. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@fe784ff Signed-off-by: nanshen3000 <2693653227@qq.com>
diff --git a/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/GLM5.2.po b/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/GLM5.2.po
@@ -220,10 +220,10 @@ msgstr "以下是一些用于优化的配置："
 
 msgid ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator."
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator."
 msgstr ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`：启用 FlashComm 优化，以减少预填充节点上的通信和计算开销。启用 FlashComm 后，layer_sharding 列表不能包含 o_proj 作为元素。\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine 融合算子。"
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine/mega_moe 融合算子。"
 
 msgid ""
 "Please refer to the following python file for further explanation and "
diff --git a/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/GLM5.po b/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/GLM5.po
@@ -259,11 +259,11 @@ msgstr "以下是一些用于优化的配置："
 
 msgid ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator.\n"
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator.\n"
 "- `VLLM_ASCEND_ENABLE_MLAPO`: Enable fused operator MlaPreprocessOperation."
 msgstr ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`：启用 FlashComm 优化，以减少预填充节点上的通信和计算开销。启用 FlashComm 后，layer_sharding 列表不能包含 o_proj 作为元素。\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine 融合算子。\n"
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine/mega_moe 融合算子。\n"
 "- `VLLM_ASCEND_ENABLE_MLAPO`：启用融合算子 MlaPreprocessOperation。"
 
 msgid ""
diff --git a/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/InternVL3.5.po b/docs/source/locale/zh_CN/LC_MESSAGES/tutorials/models/InternVL3.5.po
@@ -142,10 +142,10 @@ msgstr "一些用于优化的配置如下所示："
 
 msgid ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator."
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator."
 msgstr ""
 "- `VLLM_ASCEND_ENABLE_FLASHCOMM1`：启用 FlashComm 优化以减少 prefill 节点上的通信和计算开销。启用 FlashComm 后，layer_sharding 列表不能包含 o_proj 作为元素。\n"
-"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine 融合算子。"
+"- `VLLM_ASCEND_ENABLE_FUSED_MC2`：启用 dispatch_ffn_combine/mega_moe 融合算子。"
 
 msgid ""
 "Please refer to the following python file for further explanation and "
diff --git a/docs/source/locale/zh_CN/LC_MESSAGES/user_guide/configuration/additional_config.po b/docs/source/locale/zh_CN/LC_MESSAGES/user_guide/configuration/additional_config.po
@@ -342,14 +342,14 @@ msgstr ""
 
 msgid ""
 "| `mega_moe_max_tokens`               | int  | `65536` | Per-rank token "
-"capacity after dispatch in the mega moe (dispatch_ffn_combine) fused "
+"capacity after dispatch in the mega moe (dispatch_ffn_combine/mega_moe) fused "
 "operator. When load imbalance causes a rank to receive more tokens than this"
 " limit, the excess tokens are dropped and skipped from computation, "
 "degrading accuracy. Do not set this too large: workspace memory scales "
 "linearly with this value. |"
 msgstr ""
 "| `mega_moe_max_tokens`               | int  | `65536` | 在 mega "
-"moe（dispatch_ffn_combine）融合算子中，每个 rank 在分发后的 token 容量。当负载不均衡导致某个 rank 接收的 "
+"moe（dispatch_ffn_combine/mega_moe）融合算子中，每个 rank 在分发后的 token 容量。当负载不均衡导致某个 rank 接收的 "
 "token 超过此限制时，多余的 token 将被丢弃并跳过计算，从而降低精度。请勿将此值设置过大：工作空间内存与此值呈线性增长。 |"
 
 msgid ""
diff --git a/docs/source/tutorials/models/GLM5.2.md b/docs/source/tutorials/models/GLM5.2.md
@@ -1122,7 +1122,7 @@ python load_balance_proxy_server_example.py \
 Some configurations for optimization are shown below:
 
 - `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.
-- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator.
+- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator.
 
 Please refer to the following python file for further explanation and restrictions of the environment variables above: [envs.py](https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/envs.py)
 
diff --git a/docs/source/tutorials/models/GLM5.md b/docs/source/tutorials/models/GLM5.md
@@ -1291,7 +1291,7 @@ python load_balance_proxy_server_example.py \
 Some configurations for optimization are shown below:
 
 - `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.
-- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator.
+- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator.
 - `VLLM_ASCEND_ENABLE_MLAPO`: Enable fused operator MlaPreprocessOperation.
 
 Please refer to the following python file for further explanation and restrictions of the environment variables above: [envs.py](https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/envs.py)
diff --git a/docs/source/tutorials/models/InternVL3.5.md b/docs/source/tutorials/models/InternVL3.5.md
@@ -178,7 +178,7 @@ If you want to deploy multi-node environment, you need to set up environment on
 Some configurations for optimization are shown below:
 
 - `VLLM_ASCEND_ENABLE_FLASHCOMM1`: Enable FlashComm optimization to reduce communication and computation overhead on prefill node. With FlashComm enabled, layer_sharding list cannot include o_proj as an element.
-- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine fused operator.
+- `VLLM_ASCEND_ENABLE_FUSED_MC2`: Enable the dispatch_ffn_combine/mega_moe fused operator.
 
 Please refer to the following python file for further explanation and restrictions of the environment variables above: [envs.py](https://github.com/vllm-project/vllm-ascend/blob/main/vllm_ascend/envs.py)
 
diff --git a/vllm_ascend/ascend_forward_context.py b/vllm_ascend/ascend_forward_context.py
@@ -374,7 +374,7 @@ def select_moe_comm_method(num_tokens: int, vllm_config: VllmConfig) -> MoECommT
         moe_comm_type = MoECommType.ALLGATHER
     elif lora_config is not None and vllm_config.parallel_config.enable_expert_parallel:
         # LoRA + EP requires AlltoAll because the MC2/FusedMC2 paths
-        # Ascend MoE LoRA cannot patch FusedMC2 path for dispatch_ffn_combine
+        # Ascend MoE LoRA cannot patch FusedMC2 path for dispatch_ffn_combine/mega_moe
         # is a single fused C++ op. This covers both normal model
         # forward and _dummy_run during profile_run.
         moe_comm_type = MoECommType.ALLTOALL
diff --git a/vllm_ascend/envs.py b/vllm_ascend/envs.py
@@ -84,10 +84,11 @@
     "VLLM_ASCEND_ENABLE_NZ": lambda: int(os.getenv("VLLM_ASCEND_ENABLE_NZ", 1)),
     # Whether to anbale dynamic EPLB
     "DYNAMIC_EPLB": lambda: os.getenv("DYNAMIC_EPLB", "false").lower(),
-    # Whether to enable fused MC2 (`dispatch_ffn_combine`).
+    # Whether to enable fused MC2 (`dispatch_ffn_combine/mega_moe`).
     # 0, or not set: default ALLTOALL and MC2 will be used.
-    # 1: ALLTOALL and MC2 might be replaced by `dispatch_ffn_combine` operator.
+    # 1: ALLTOALL and MC2 might be replaced by `dispatch_ffn_combine/mega_moe` operator.
     # `dispatch_ffn_combine` can be used only for moe layer with W8A8, EP<=32, non-mtp, non-dynamic-eplb.
+    # `mega_moe` can be used only for moe layer with W8A8/W4A8/bf16(none quant), EP<=64, non-dynamic-eplb.
     "VLLM_ASCEND_ENABLE_FUSED_MC2": lambda: int(os.getenv("VLLM_ASCEND_ENABLE_FUSED_MC2", "0")),
     # DEPRECATED: VLLM_ASCEND_BALANCE_SCHEDULING env var will be removed in a future release.
     # Use --additional-config '{"enable_balance_scheduling": true}' instead.
diff --git a/vllm_ascend/lora/fused_moe.py b/vllm_ascend/lora/fused_moe.py
@@ -181,7 +181,7 @@ def _assert_ascend_moe_lora_supported(base_layer: nn.Module) -> None:
     if int(envs_ascend.VLLM_ASCEND_ENABLE_FUSED_MC2) != 0:
         raise AssertionError(
             "Ascend MoE LoRA cannot patch FusedMC2 path "
-            "(dispatch_ffn_combine is a single fused C++ op). "
+            "(dispatch_ffn_combine/mega_moe is a single fused C++ op). "
             "Set VLLM_ASCEND_ENABLE_FUSED_MC2=0."
         )
     if getattr(base_layer, "_shared_experts", None) is not None:
diff --git a/vllm_ascend/ops/fused_moe/fused_moe.py b/vllm_ascend/ops/fused_moe/fused_moe.py
@@ -137,9 +137,9 @@ def process_weights_after_loading(self, layer):
         w2_data = self._maybe_pad_weight(layer.w2_weight.data).transpose(1, 2).contiguous()
         layer.w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)
 
-        # TODO: Current dispatch_ffn_combine fusion operator ONLY supports NZ format.
+        # TODO: Current dispatch_ffn_combine/mega_moe fusion operator ONLY supports NZ format.
         # Therefore, we must cast weights to NZ when fusion is enabled.
-        # Once the underlying dispatch_ffn_combine operator is updated to support
+        # Once the underlying dispatch_ffn_combine/mega_moe operator is updated to support
         # ND format (or other formats), remove this specific 'if' check and the forced
         # npu_format_cast. At that point, the operator should be able to handle weights
         # in their native format without explicit casting here.
@@ -242,11 +242,11 @@ def apply(
         moe_comm_method = _EXTRA_CTX.moe_comm_method
         # NOTE: In the MoECommType.FUSED_MC2 branch, we wrap weights (w1, w2) into lists
         # and provide dummy scales (w1_scale, w2_scale). This is required because:
-        # The underlying Ascend fused operator (e.g., dispatch_ffn_combine) expects
+        # The underlying Ascend fused operator (e.g., dispatch_ffn_combine/mega_moe) expects
         # inputs in a list format.
         # TODO: Passing an empty tensor as scale for float (BF16) cases is semantically
         # incorrect. The ideal solution is to pass None. However, if the underlying
-        # dispatch_ffn_combine C++ operator does not support None for the scale argument
+        # dispatch_ffn_combine/mega_moe C++ operator does not support None for the scale argument
         # (due to signature constraints), we are forced to use a placeholder empty tensor.
         # This TODO tracks the requirement to update the C++ operator to accept Optional[Tensor]
         # or None for scales in non-quantized scenarios.

Original file line number	Diff line number	Diff line change
`@@ -181,7 +181,7 @@ def _assert_ascend_moe_lora_supported(base_layer: nn.Module) -> None:`
`181`	`181`	`if int(envs_ascend.VLLM_ASCEND_ENABLE_FUSED_MC2) != 0:`
`182`	`182`	`raise AssertionError(`
`183`	`183`	`"Ascend MoE LoRA cannot patch FusedMC2 path "`
`184`		`- "(dispatch_ffn_combine is a single fused C++ op). "`
	`184`	`+ "(dispatch_ffn_combine/mega_moe is a single fused C++ op). "`
`185`	`185`	`"Set VLLM_ASCEND_ENABLE_FUSED_MC2=0."`
`186`	`186`	`)`
`187`	`187`	`if getattr(base_layer, "_shared_experts", None) is not None:`