shared expert dp for deepseek r1 alltoall

vllm_ascend/ascend_forward_context.py

@@ -12,6 +12,7 @@
 
 import vllm_ascend.envs as envs_ascend
 from vllm_ascend.utils import enable_sp
+from vllm_ascend.ascend_config import get_ascend_config
 
 
 class FusedMoEState(Enum):
@@ -110,6 +111,8 @@ def set_ascend_forward_context(
             tp_world_size > 1 and \
             num_tokens is not None and num_tokens > 1000
 
+        if get_ascend_config().enable_shared_expert_dp:
+            sp_enabled = True
         if sp_enabled:
             pad_size = (tp_world_size -
                         (num_tokens % tp_world_size)) % tp_world_size

vllm_ascend/attention/mla_v1.py

@@ -887,9 +887,10 @@ def _mla_preprocess(self, layer_name, hidden_states, kv_cache,
 
         kv_no_split = self.kv_a_proj_with_mqa(hidden_states)[0]
         # Process for shared_expert_dp
-        if need_gather_q_kv:
-            q_c = get_tp_group().all_gather(q_c, 0)
-            kv_no_split = get_tp_group().all_gather(kv_no_split, 0)
+        q_c = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+            q_c, need_gather_q_kv)
+        kv_no_split = torch.ops.vllm.maybe_all_gather_and_maybe_unpad(
+            kv_no_split, need_gather_q_kv)
         decode_preprocess_res = None
         prefill_preprocess_res = None
         if has_prefill:
@@ -1011,19 +1012,15 @@ def forward(
                          enabled=self.enable_prefetch)
 
             output[...] = self.o_proj(
-                o_proj_input,
-                is_prefill=prefill_preprocess_res is not None,
-                is_force_scatter=self.enable_shared_expert_dp)[0]
+                o_proj_input)[0]
         else:
             with torch.npu.stream(current_ms_metadata.comm_stream):
                 npu_prefetch(self.o_proj.weight,
                              o_proj_input,
                              max_size=MAX_O_PROJ_PREFETCH_SIZE,
                              enabled=self.enable_prefetch)
                 output[...] = self.o_proj(
-                    o_proj_input,
-                    is_prefill=prefill_preprocess_res is not None,
-                    is_force_scatter=self.enable_shared_expert_dp)[0]
+                    o_proj_input)[0]
                 current_ms_metadata.after_comm_event.record()
         del o_proj_input
 

vllm_ascend/models/deepseek_v2.py

@@ -181,7 +181,7 @@ def __init__(
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.kv_b_proj")
-        self.o_proj = CustomDeepseekV2RowParallelLinear(
+        self.o_proj = RowParallelLinear(
             self.num_heads * self.v_head_dim,
             self.hidden_size,
             bias=False,

vllm_ascend/models/layers/mla.py

@@ -120,18 +120,17 @@ def forward(
             hidden_states: torch.Tensor,
             kv_cache: Optional[torch.Tensor] = None,
             attn_metadata: Optional[AttentionMetadata] = None) -> torch.Tensor:
-        num_tokens = hidden_states.shape[0]
+        forward_context = get_forward_context()
+        sp_enabled = forward_context.sp_enabled
         need_gather_q_kv = False
-        if self.enable_shared_expert_dp and self.debug_layer_idx > self.first_k_dense_replace and self.debug_layer_idx < self.layers:
-            # Simulate all gather to calculate output shape
-            num_tokens = num_tokens * self.tp_size
+        if sp_enabled and self.debug_layer_idx > 0:
             need_gather_q_kv = True
-        if not self.enable_shared_expert_dp or self.debug_layer_idx < self.first_k_dense_replace:
+        if not self.enable_shared_expert_dp or self.debug_layer_idx > 0:
             output_shape = hidden_states.shape
         else:
-            rows = num_tokens // self.tp_size
-            if num_tokens % self.tp_size:
-                rows += 1
+            num_tokens = hidden_states.shape[0]
+            pad_size = forward_context.pad_size
+            rows = (num_tokens + pad_size) // self.tp_size
             output_shape = (rows, hidden_states.shape[1])
         # FIXME: This does not seem right, should make sure the buffer is fixed
         output = torch.empty(output_shape,

vllm_ascend/ops/common_fused_moe.py

@@ -346,7 +346,8 @@ def forward(
             forward_context = get_forward_context()
             moe_comm_type = forward_context.moe_comm_type
             if moe_comm_type in {MoECommType.ALLTOALL, MoECommType.MC2}:
-                shared_out = tensor_model_parallel_all_reduce(shared_out)
+                if not get_forward_context().sp_enabled:
+                    shared_out = tensor_model_parallel_all_reduce(shared_out)
 
         _, fused_out = AscendFusedMoE.forward(
             self,

vllm_ascend/ops/layernorm.py

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

vllm_ascend/ops/linear_op.py

@@ -46,6 +46,7 @@
 from vllm.distributed import split_tensor_along_last_dim
 from vllm.distributed.parallel_state import get_tp_group
 
+from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.distributed.parallel_state import (get_mlp_tp_group,
                                                     get_otp_group)
 from vllm_ascend.utils import (dense_optim_enable, enable_sp,
@@ -418,11 +419,13 @@ def get_column_parallel_op(
         SequenceMergedColumnParallelOp,
         SequenceQKVParallelOp,
     ]] = None
-    if "gate_up_proj" in prefix and mlp_tp_enable():
+    if "shared_experts.gate_up_proj" in prefix and enable_sp():
+        return None, 0, 1
+    elif "gate_up_proj" in prefix and mlp_tp_enable():
         custom_op = MLPColumnParallelOp(layer)
     elif "gate_up_proj" in prefix and enable_sp():
         custom_op = SequenceMergedColumnParallelOp(layer)
-    elif enable_sp():
+    elif "qkv_proj" in prefix and enable_sp():
         custom_op = SequenceQKVParallelOp(layer, prefix)
 
     if custom_op is not None:
@@ -442,7 +445,9 @@ def get_row_parallel_op(
     custom_op: Optional[Union[MLPRowParallelOp, OProjRowParallelOp,
                               MatmulAllreduceRowParallelOp,
                               SequenceRowParallelOp]] = None
-    if "down_proj" in prefix and mlp_tp_enable():
+    if "shared_experts.down_proj" in prefix and enable_sp():
+        return None, 0, 1
+    elif "down_proj" in prefix and mlp_tp_enable():
         custom_op = MLPRowParallelOp(layer)
     elif "o_proj" in prefix and oproj_tp_enable():
         custom_op = OProjRowParallelOp(layer)

vllm_ascend/platform.py

@@ -260,7 +260,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
             compilation_config.level = CompilationLevel.NO_COMPILATION
 
         if parallel_config and parallel_config.worker_cls == "auto":
-            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"
@@ -303,9 +303,6 @@ def get_attn_backend_cls(cls,
 
         ascend_config = get_ascend_config()
 
-        if use_mla and ascend_config.enable_shared_expert_dp:
-            return "vllm_ascend.torchair.torchair_mla.AscendMLATorchairBackend"
-
         use_torchair = ascend_config.torchair_graph_config.enabled
         # choose attention backend based on use_mla and use_torchair
         backend_map = {

vllm_ascend/utils.py

@@ -595,7 +595,8 @@ def enable_sp() -> bool:
 
     return (
         get_cached_compilation_config().pass_config.enable_sequence_parallelism
-        or envs_ascend.VLLM_ASCEND_ENABLE_FLASHCOMM)
+        or envs_ascend.VLLM_ASCEND_ENABLE_FLASHCOMM
+        or get_ascend_config().enable_shared_expert_dp)
 
 
 def is_moe_model(vllm_config: VllmConfig):