[Feat] shared expert dp for deepseek_mtp and e2e

Signed-off-by: chenmenglong <[email protected]>

Author: dragondream-chen

tests/e2e/multicard/test_shared_expert_dp.py

@@ -0,0 +1,93 @@
+import os
+
+import pytest
+from vllm import SamplingParams
+
+from tests.e2e.conftest import VllmRunner
+from tests.e2e.model_utils import check_outputs_equal
+
+MODELS = [
+    "vllm-ascend/DeepSeek-V2-Lite",
+]
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+
+@pytest.mark.parametrize("model", MODELS)
+def test_models_with_enable_shared_expert_dp(model: str) -> None:
+
+    if 'HCCL_OP_EXPANSION_MODE' in os.environ:
+        del os.environ['HCCL_OP_EXPANSION_MODE']
+
+    prompts = [
+        "Hello, my name is", "The capital of the United States is",
+        "The capital of France is", "The future of AI is"
+    ]
+    sampling_params = SamplingParams(max_tokens=32, temperature=0.0)
+
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            enforce_eager=True,
+            tensor_parallel_size=2,
+            enable_expert_parallel=True,
+    ) as runner:
+        vllm_eager_outputs = runner.model.generate(prompts, sampling_params)
+
+    os.environ["VLLM_ASCEND_ENABLE_FLASHCOMM1"] = "1"
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            enforce_eager=True,
+            tensor_parallel_size=2,
+            enable_expert_parallel=True,
+            additional_config={
+                "enable_shared_expert_dp": True,
+            },
+    ) as runner:
+        shared_expert_dp_eager_outputs = runner.model.generate(
+            prompts, sampling_params)
+
+    with VllmRunner(
+            model,
+            max_model_len=1024,
+            tensor_parallel_size=2,
+            enforce_eager=False,
+            compilation_config={
+                "cudagraph_capture_sizes": [1, 4, 8, 16],
+                "cudagraph_mode": "FULL_DECODE_ONLY",
+            },
+            additional_config={
+                "enable_shared_expert_dp": True,
+            },
+    ) as runner:
+        shared_expert_dp_aclgraph_outputs = runner.model.generate(
+            prompts, sampling_params)
+
+    vllm_eager_outputs_list = []
+    for output in vllm_eager_outputs:
+        vllm_eager_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    shared_expert_dp_eager_outputs_list = []
+    for output in shared_expert_dp_eager_outputs:
+        shared_expert_dp_eager_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    shared_expert_dp_aclgraph_outputs_list = []
+    for output in shared_expert_dp_aclgraph_outputs:
+        shared_expert_dp_aclgraph_outputs_list.append(
+            (output.outputs[0].index, output.outputs[0].text))
+
+    check_outputs_equal(
+        outputs_0_lst=vllm_eager_outputs_list,
+        outputs_1_lst=shared_expert_dp_eager_outputs_list,
+        name_0="vllm_eager_outputs",
+        name_1="shared_expert_dp_eager_outputs",
+    )
+
+    check_outputs_equal(
+        outputs_0_lst=vllm_eager_outputs_list,
+        outputs_1_lst=shared_expert_dp_aclgraph_outputs_list,
+        name_0="vllm_eager_outputs",
+        name_1="shared_expert_dp_aclgraph_outputs",
+    )

tests/ut/ops/test_layernorm.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import patch
 
 import pytest
 import torch
@@ -41,7 +42,9 @@ def context(self, mocker: MockerFixture):
     # Test case for the most common and basic scenario
     @pytest.mark.parametrize(
         "residual", [None, torch.randn(4, 8, dtype=torch.float16)])
-    def test_forward_oot_basic(self, residual):
+    @patch("torch.ops.vllm.maybe_chunk_residual")
+    def test_forward_oot_basic(self, mock_maybe_chunk_residual, residual):
+        mock_maybe_chunk_residual.side_effect = lambda x, residual: residual
         layer = RMSNorm(hidden_size=8, eps=1e-05)
         x = torch.randn(4, 8, dtype=torch.float16)
         if residual is not None:
@@ -105,6 +108,8 @@ def test_forward_oot_with_quant_fusion(self, mocker: MockerFixture):
         mock_forward_context.num_hidden_layers = num_hidden_layers
         mock_forward_context.fusion_linear = "gate_up_dense"
         mock_forward_context.weight_prefetch_method = None
+        mocker.patch("torch.ops.vllm.maybe_chunk_residual",
+                     lambda x, residual: residual)
 
         # Ensure fusion and layer_idx increment are handled correctly
         x = torch.randn(4, 8, dtype=torch.float16)

vllm_ascend/ascend_config.py

@@ -68,6 +68,10 @@ def __init__(self, vllm_config):
         self.enable_shared_expert_dp = additional_config.get(
             "enable_shared_expert_dp", False
         ) and not self.torchair_graph_config.enabled and vllm_config.parallel_config.enable_expert_parallel
+        if self.enable_shared_expert_dp:
+            from vllm_ascend.utils import enable_sp
+            assert enable_sp(
+                vllm_config), "shared_expert_dp requires enable_sp=True."
         self.multistream_overlap_shared_expert = additional_config.get(
             "multistream_overlap_shared_expert", False)
         self.recompute_scheduler_enable = additional_config.get(

vllm_ascend/attention/mla_v1.py

@@ -1456,6 +1456,8 @@ def forward(
         forward_context = get_forward_context()
         if (self.enable_mlapo and
             (attn_metadata is None or not forward_context.with_prefill)):
+            hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+                hidden_states.contiguous(), need_gather_q_kv)
             decode_preprocess_res, prefill_preprocess_res = self._mla_decode_preprocess(
                 hidden_states, kv_cache, attn_metadata)
         else:

vllm_ascend/ops/layernorm.py

@@ -109,6 +109,7 @@ def forward_oot(
         import torch_npu
 
         if residual is not None:
+            residual = torch.ops.vllm.maybe_chunk_residual(x, residual)
             assert x.size(0) == residual.size(0)
             x, residual = _addrmsnorm_forward_oot(
                 self, x, residual, self.next_need_quant_fusion_linear,

vllm_ascend/ops/register_custom_ops.py

@@ -2,6 +2,7 @@
 import torch.nn.functional as F
 import torch_npu
 from vllm.distributed import (get_dp_group, get_ep_group,
+                              get_tensor_model_parallel_rank,
                               get_tensor_model_parallel_world_size,
                               tensor_model_parallel_all_gather,
                               tensor_model_parallel_all_reduce,
@@ -20,6 +21,27 @@
     from vllm.utils.torch_utils import direct_register_custom_op
 
 
+def _maybe_chunk_residual_impl(x: torch.Tensor,
+                               residual: torch.Tensor) -> torch.Tensor:
+    try:
+        forward_context = get_forward_context()
+    except AssertionError:
+        return residual
+
+    if x.size(0) != residual.size(0):
+        sp_enabled = forward_context.sp_enabled
+        assert sp_enabled is True, ("Currently, this situation only occurs "
+                                    "when sp is enabled")
+        pad_size = forward_context.pad_size
+        if pad_size > 0:
+            residual = F.pad(residual, (0, 0, 0, pad_size))
+        tp_size = get_tensor_model_parallel_world_size()
+        tp_rank = get_tensor_model_parallel_rank()
+        residual = torch.chunk(residual, tp_size, dim=0)[tp_rank]
+
+    return residual
+
+
 def _maybe_all_gather_and_maybe_unpad_impl(
         x: torch.Tensor,
         label: bool,
@@ -264,6 +286,12 @@ def _matmul_and_reduce_impl_fake(input_parallel: torch.Tensor,
     return output
 
 
+direct_register_custom_op(op_name="maybe_chunk_residual",
+                          op_func=_maybe_chunk_residual_impl,
+                          fake_impl=lambda x, residual: x,
+                          mutates_args=[],
+                          dispatch_key="PrivateUse1")
+
 direct_register_custom_op(op_name="maybe_all_gather_and_maybe_unpad",
                           op_func=_maybe_all_gather_and_maybe_unpad_impl,
                           fake_impl=_maybe_all_gather_and_maybe_unpad_fake,

vllm_ascend/platform.py

@@ -364,7 +364,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         if parallel_config and parallel_config.worker_cls == "auto":
             # TODO: this is a tricky way to disable `use_sequence_parallel_moe` in vllm.
             os.environ["VLLM_ALL2ALL_BACKEND"] = "flashinfer_all2allv"
-            if ascend_config.torchair_graph_config.enabled or ascend_config.enable_shared_expert_dp:
+            if ascend_config.torchair_graph_config.enabled:
                 parallel_config.worker_cls = "vllm_ascend.torchair.torchair_worker.NPUTorchairWorker"
             else:
                 parallel_config.worker_cls = "vllm_ascend.worker.worker_v1.NPUWorker"
@@ -435,8 +435,6 @@ def get_attn_backend_cls(
         ascend_config = get_ascend_config()
 
         if use_mla and ascend_config.enable_shared_expert_dp:
-            if use_mla and not use_sparse:
-                return "vllm_ascend.torchair.torchair_mla.AscendMLATorchairBackend"
             if use_mla and use_sparse:
                 return "vllm_ascend.torchair.torchair_sfa.AscendSFATorchairBackend"
 

vllm_ascend/spec_decode/mtp_proposer.py

@@ -211,6 +211,11 @@ def dummy_run(self,
                     num_actual_tokens=0,
                     aclgraph_runtime_mode=aclgraph_runtime_mode,
                     batch_descriptor=batch_descriptor):
+                positions = positions.unsqueeze(-1)
+                positions = torch.ops.vllm.maybe_pad_and_reduce(positions)
+                positions = positions.squeeze(-1)
+                previous_hidden_states = torch.ops.vllm.maybe_pad_and_reduce(
+                    previous_hidden_states)
                 self.model(input_ids=input_ids,
                            positions=positions,
                            hidden_states=previous_hidden_states)
@@ -528,11 +533,20 @@ def _propose(
                 with ProfileExecuteDuration().capture_async('mtp_forward'):
                     model_kwargs = {}
                     model_kwargs["attn_metadata"] = attn_metadata
-
-                    hidden_states = self.model(
-                        input_ids=self.input_ids[:num_input_tokens],
-                        positions=self.positions[:num_input_tokens],
-                        hidden_states=self.hidden_states[:num_input_tokens])
+                    input_ids = self.input_ids[:num_input_tokens]
+                    positions = self.positions[:num_input_tokens]
+                    hidden_states = self.hidden_states[:num_input_tokens]
+
+                    # positions [N] -> [N, 1] for padding
+                    positions = positions.unsqueeze(-1)
+                    positions = torch.ops.vllm.maybe_pad_and_reduce(positions)
+                    positions = positions.squeeze(-1)
+
+                    hidden_states = self.model(input_ids=input_ids,
+                                               positions=positions,
+                                               hidden_states=hidden_states)
+                    hidden_states = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+                        hidden_states.contiguous(), True)
 
             num_indices = last_token_indices.shape[0]
             if lmhead_tp_enable():