[Feature] Support moe multi-stream for aclgraph. (#2946)

This PR puts the calculation of shared experts into a separate stream,
overlaping with routing experts.

- vLLM version: v0.10.2
- vLLM main:
fbd6523ac0

---------

Signed-off-by: whx-sjtu <2952154980@qq.com>

tests/e2e/multicard/test_offline_inference_distributed.py

@@ -66,8 +66,8 @@ def test_models_distributed_DeepSeek_multistream_moe():
             additional_config={
                 "torchair_graph_config": {
                     "enabled": True,
-                    "enable_multistream_moe": True,
                 },
+                "enable_multistream_moe": True,
                 "ascend_scheduler_config": {
                     "enabled": True,
                 },

tests/e2e/singlecard/test_multistream_overlap_shared_expert.py

@@ -0,0 +1,103 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+Compare the outputs of vLLM with multistream_overlap_shared_expert
+enabled and disabled.
+
+Run `pytest tests/e2e/singlecard/test_multistream_overlap_shared_expert.py`.
+"""
+
+import pytest
+from vllm import SamplingParams
+
+from tests.e2e.conftest import VllmRunner
+from tests.e2e.model_utils import check_outputs_equal
+
+MODELS = [
+    "Qwen/Qwen3-0.6B",
+]
+
+
+@pytest.mark.parametrize("model", MODELS)
+@pytest.mark.parametrize("max_tokens", [32])
+def test_models_with_multistream_overlap_shared_expert(
+    model: str,
+    max_tokens: int,
+) -> None:
+    prompts = [
+        "Hello, my name is", "The president of the United States is",
+        "The capital of France is", "The future of AI is"
+    ]
+
+    sampling_params = SamplingParams(max_tokens=max_tokens, temperature=0.0)
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            enforce_eager=True,
+            additional_config={
+                "multistream_overlap_shared_expert": True,
+            },
+    ) as runner:
+        vllm_moe_ms_eager_outputs = runner.model.generate(
+            prompts, sampling_params)
+
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            enforce_eager=False,
+            additional_config={
+                "multistream_overlap_shared_expert": True,
+            },
+    ) as runner:
+        vllm_moe_ms_aclgraph_outputs = runner.model.generate(
+            prompts, sampling_params)
+
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            enforce_eager=True,
+    ) as runner:
+        vllm_eager_outputs = runner.model.generate(prompts, sampling_params)
+
+    vllm_moe_ms_eager_outputs_list = []
+    for output in vllm_moe_ms_eager_outputs:
+        vllm_moe_ms_eager_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    vllm_moe_ms_aclgraph_outputs_list = []
+    for output in vllm_moe_ms_aclgraph_outputs:
+        vllm_moe_ms_aclgraph_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    vllm_eager_outputs_list = []
+    for output in vllm_eager_outputs:
+        vllm_eager_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    check_outputs_equal(
+        outputs_0_lst=vllm_eager_outputs_list,
+        outputs_1_lst=vllm_moe_ms_eager_outputs_list,
+        name_0="vllm_eager_outputs",
+        name_1="vllm_moe_ms_eager_outputs",
+    )
+
+    check_outputs_equal(
+        outputs_0_lst=vllm_eager_outputs_list,
+        outputs_1_lst=vllm_moe_ms_aclgraph_outputs_list,
+        name_0="vllm_eager_outputs",
+        name_1="vllm_moe_ms_aclgraph_outputs",
+    )

tests/ut/ops/test_fused_ops.py

@@ -94,7 +94,8 @@ def mock_dist_env(mocker: MockerFixture):
             return_value=mock_dp_and_tp_group(mocker)), \
         patch('vllm_ascend.ops.fused_moe.get_ascend_config',
             return_value=MagicMock(
-                torchair_graph_config=MagicMock(enabled=False, enable_multistream_moe=False),
+                torchair_graph_config=MagicMock(enabled=False),
+                enable_multistream_moe=False,
                 expert_map_path=None
             )), \
         patch('vllm_ascend.ops.fused_moe.determine_expert_map',

tests/ut/test_ascend_config.py

@@ -43,6 +43,7 @@ class TestAscendConfig(TestBase):
         # No additional config given, check the default value here.
         ascend_config = init_ascend_config(test_vllm_config)
         self.assertIsNone(ascend_config.expert_map_path)
+        self.assertFalse(ascend_config.multistream_overlap_shared_expert)
 
         torchair_graph_config = ascend_config.torchair_graph_config
         self.assertFalse(torchair_graph_config.enabled)
@@ -51,7 +52,6 @@ class TestAscendConfig(TestBase):
         self.assertEqual(torchair_graph_config.graph_batch_sizes, [])
         self.assertFalse(torchair_graph_config.graph_batch_sizes_init)
         self.assertFalse(torchair_graph_config.enable_multistream_mla)
-        self.assertFalse(torchair_graph_config.enable_multistream_moe)
         self.assertTrue(torchair_graph_config.enable_view_optimize)
         self.assertTrue(torchair_graph_config.enable_frozen_parameter)
         self.assertFalse(torchair_graph_config.enable_kv_nz)
@@ -69,11 +69,11 @@ class TestAscendConfig(TestBase):
                 "graph_batch_sizes": [1, 2, 4],
                 "graph_batch_sizes_init": False,
                 "enable_multistream_mla": True,
-                "enable_multistream_moe": True,
                 "enable_view_optimize": True,
                 "enable_frozen_parameter": True,
                 "enable_kv_nz": True
             },
+            "multistream_overlap_shared_expert": True,
             "ascend_scheduler_config": {
                 "enabled": True
             },
@@ -82,6 +82,7 @@ class TestAscendConfig(TestBase):
         }
         ascend_config = init_ascend_config(test_vllm_config)
         self.assertEqual(ascend_config.expert_map_path, "test_expert_map_path")
+        self.assertTrue(ascend_config.multistream_overlap_shared_expert)
 
         torchair_graph_config = ascend_config.torchair_graph_config
         self.assertTrue(torchair_graph_config.enabled)
@@ -89,7 +90,6 @@ class TestAscendConfig(TestBase):
         self.assertEqual(torchair_graph_config.graph_batch_sizes, [1, 2, 4])
         self.assertFalse(torchair_graph_config.graph_batch_sizes_init)
         self.assertTrue(torchair_graph_config.enable_multistream_mla)
-        self.assertTrue(torchair_graph_config.enable_multistream_moe)
         self.assertTrue(torchair_graph_config.enable_view_optimize)
         self.assertTrue(torchair_graph_config.enable_frozen_parameter)
         self.assertTrue(torchair_graph_config.enable_kv_nz)
@@ -306,17 +306,6 @@ class TestAscendConfig(TestBase):
             }
             init_ascend_config(test_vllm_config)
 
-        # enable_multistream_moe should not be enabled without torchair graph mode
-        with self.assertRaises(RuntimeError):
-            test_vllm_config.additional_config = {
-                "torchair_graph_config": {
-                    "enabled": False,
-                    "enable_multistream_moe": True,
-                },
-                "refresh": True
-            }
-            init_ascend_config(test_vllm_config)
-
         # mode should not be configured without torchair graph mode
         with self.assertRaises(RuntimeError):
             test_vllm_config.additional_config = {

tests/ut/torchair/ops/test_torchair_fused_moe.py

@@ -70,7 +70,8 @@ def mock_dist_env(mocker: MockerFixture):
                return_value=mock_dp_and_tp_group(mocker)), \
          patch('vllm_ascend.torchair.ops.torchair_fused_moe.get_ascend_config',
                return_value=MagicMock(
-                   torchair_graph_config=MagicMock(enabled=False, enable_multistream_moe=False),
+                   torchair_graph_config=MagicMock(enabled=False),
+                   enable_multistream_moe=False,
                    expert_map_path=None
                )), \
          patch('vllm_ascend.torchair.ops.torchair_fused_moe.determine_expert_map',