fix bug

Signed-off-by: weijinqian_v1 <weijinqian@huawei.com>

Author: harygo22

vllm_ascend/ascend_forward_context.py

@@ -21,7 +21,7 @@ class FusedMoEState(Enum):
 def get_fused_moe_state(ep_size: int, with_prefill: bool):
     if ep_size == 1:
         return FusedMoEState.AllGather
-    elif with_prefill and envs_ascend.VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ:
+    elif envs_ascend.VLLM_ASCEND_ENABLE_MOE_ALL2ALL_SEQ:
         return FusedMoEState.All2AllSeq
     # NOTE: mc2 need ep_size >= 16 & all2all can't use in torchair graph.
     elif ep_size < 16 or with_prefill:

vllm_ascend/models/qwen3_dbo.py

@@ -6,7 +6,6 @@
 import torch_npu
 from torch import nn
 from transformers import PretrainedConfig
-from vllm.compilation.decorators import support_torch_compile
 
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
 from vllm.config import CacheConfig, VllmConfig
@@ -22,6 +21,7 @@
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.compilation.decorators import support_torch_compile
 
 from vllm_ascend.multistream.context import (
     advance_step_multistream_layer_context, get_multistream_comm_context,
@@ -35,6 +35,7 @@
 from vllm_ascend.ops.fused_moe import AscendFusedMoE, select_experts, apply_mlp
 from vllm_ascend.distributed.tensor_parallel import gather_from_sequence_parallel_region
 import vllm_ascend.envs as envs_ascend
+from vllm_ascend.models.qwen3_moe import CustomQwen3MoeForCausalLM
 
 VLLM_ASCEND_ENABLE_DBO: bool = envs_ascend.VLLM_ASCEND_ENABLE_DBO
 
@@ -197,7 +198,7 @@ def _forward_op_grouped_mlp(
             self, dispatched_input, tokens_per_expert
     ):
         return apply_mlp(
-            [dispatched_input],
+            dispatched_input,
             self.mlp.experts.w13_weight,
             self.mlp.experts.w2_weight,
             tokens_per_expert
@@ -207,8 +208,7 @@ def _forward_combine_comm(
             self, hidden_states, microbatch_id, num_tokens, chunked_hidden_states_sizes
     ):
         token_dispatcher = self.mlp.experts.token_dispatchers[microbatch_id]
-        token_dispatcher.combine_alltoall()
-        final_hidden_states = token_dispatcher.unpermute2()
+        final_hidden_states, _ = token_dispatcher.token_unpermutation(hidden_states)
         if hasattr(self.mlp, 'routed_scaling_factor'):
             final_hidden_states = final_hidden_states * self.mlp.routed_scaling_factor
 
@@ -267,13 +267,10 @@ def discard_tensor(tensor):
         # communication in the previous layer, and the attn computation of microbatch 2
         # can be overlapped with the attn communication of microbatch 1
         for i in range(num_micro_batchs):
-            # wait last layer moe finishing communication
-            ms_metadata.try_wait_event(layer_index - 1, i,
-                                       MSEventKey.MOE_AFTER_COMM)
-
             forward_context = get_forward_context()
             layer_index, ms_metadata, attn_metadata = get_multistream_layer_context(
             )
+            ms_metadata.try_wait_event(layer_index - 1, i, MSEventKey.FFN_AR_FINISH)
             forward_context.attn_metadata = attn_metadata[i]
 
             # input layernorm
@@ -309,36 +306,25 @@ def discard_tensor(tensor):
             with torch.npu.stream(dispatch_context.comm_stream):
                 dispatch_context.comm_stream.wait_event(dispatch_context.before_comm_event)
                 token_dispatchers[i].dispatch_alltoall()
+                dispatched_input[i], tokens_per_expert[i] = token_dispatchers[i].permute2()
                 dispatch_context.after_comm_event.record()
 
-                if has_shared_expert:
-                    token_dispatchers[i].cached_shared_expert_output = tensor_model_parallel_all_reduce(
-                        token_dispatchers[i].cached_shared_expert_output
-                    )
-                    ms_metadata.ms_events[layer_index][i][MSEventKey.MOE_SE_COMM_FINISH].record()
-
         # print_with_sync('begin experts...', torch.distributed.get_rank())
         # block 4 : Router Experts Computation
         # block 5 : Token Combine Communication
         for i in range(num_micro_batchs):
-
             ms_metadata.try_wait_event(layer_index, i, MSEventKey.MOE_AFTER_COMM)
             discard_tensor(hidden_states[i])
-
-            dispatched_input[i], tokens_per_expert[i] = token_dispatchers[i].permute2()
             router_expert_output[i] = self._forward_op_grouped_mlp(dispatched_input[i], tokens_per_expert[i])
             discard_tensor(dispatched_input[i])
-            token_dispatchers[i].unpermute1(router_expert_output[i])
-            if router_expert_output[i].shape[0] > 0 and token_dispatchers[i].num_local_experts > 1:
-                discard_tensor(router_expert_output[i])
 
             # Launch Combine Comm in a New Stream.
             combine_context = MultiStreamStepMetadata(
                 comm_stream=ms_metadata.communicate_stream,
                 before_comm_event=ms_metadata.ms_events[layer_index][i][
-                    MSEventKey.MOE_BEFORE_COMM],
+                    MSEventKey.FFN_COM_FINISH],
                 after_comm_event=ms_metadata.ms_events[layer_index][i][
-                    MSEventKey.MOE_AFTER_COMM],
+                    MSEventKey.FFN_AR_FINISH],
             )
             combine_context.before_comm_event.record()
             ms_metadata.try_wait_event(layer_index, i, MSEventKey.MOE_SE_COMM_FINISH)
@@ -347,7 +333,7 @@ def discard_tensor(tensor):
                 hidden_states[i] = self._forward_combine_comm(
                     router_expert_output[i], i, num_tokens[i], chunked_hidden_states_sizes[i]
                 )
-                combine_context.after_comm_event.record()
+                ms_metadata.ms_events[layer_index][i][MSEventKey.FFN_AR_FINISH] = combine_context.comm_stream.record_event()
 
         return hidden_states, residual
 
@@ -443,11 +429,10 @@ def forward(
     def can_run_ms(self):
         attn_metadata = get_forward_context().attn_metadata
         # enable prefill overlap
-        with_prefill = getattr(attn_metadata, "with_prefill_across_dp", False)
+        with_prefill = get_forward_context().with_prefill
         if attn_metadata is None or not with_prefill or not attn_metadata.enable_dbo_across_dp:
             return False
-        # if torch.distributed.get_rank() == 0:
-        #     print(attn_metadata)
+
         return True
 
     def _forward_ms_layers(
@@ -465,9 +450,7 @@ def _forward_ms_layers(
         attn_metadata, [positions, hidden_states,
                         residual] = self.ms_pre_layer(
             [positions, hidden_states, residual], )
-        # if torch.distributed.get_rank() == 0:
-        #     print(attn_metadata[0], attn_metadata[1])
-        # exit()
+        num_micro_batch = len(attn_metadata)
         # the rest layers
         for i in range(moe_start_layer, self.end_layer):
             layer = self.layers[i]
@@ -481,6 +464,11 @@ def _forward_ms_layers(
             )
             advance_step_multistream_layer_context()
 
+        layer_index, ms_metadata, attn_metadata = get_multistream_layer_context()
+        for i in range(num_micro_batch):
+            ms_metadata.try_wait_event(layer_index - 1, i, MSEventKey.FFN_AR_FINISH)
+
+
         [hidden_states,
          residual] = self.ms_post_layer([hidden_states, residual], )
         return hidden_states, residual
@@ -517,17 +505,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.logits_processor = LogitsProcessor(config.vocab_size)
         self.make_empty_intermediate_tensors = (
             self.model.make_empty_intermediate_tensors)
+    
+    def forward(self, *args, **kwargs):
+        if "graph_enable" in kwargs:
+            kwargs.pop('graph_enable')
+        return super().forward(*args, **kwargs)
 
-    def forward(
-            self,
-            input_ids: torch.Tensor,
-            positions: torch.Tensor,
-            intermediate_tensors: Optional[IntermediateTensors] = None,
-            inputs_embeds: Optional[torch.Tensor] = None,
-            graph_enable: Optional[bool] = True
-    ) -> Union[torch.Tensor, IntermediateTensors]:
-        hidden_states = self.model(input_ids, positions, intermediate_tensors,
-                                   inputs_embeds)
-        return hidden_states
 
 

vllm_ascend/models/qwen3_moe.py

@@ -15,10 +15,26 @@
 # limitations under the License.
 # Adapted from vllm/model_executor/models/qwen3_moe.py
 # This file is a part of the vllm-ascend project.
+from typing import Optional
 
+import torch
+import vllm
+from torch import nn
+from transformers import PretrainedConfig
+from vllm.attention import AttentionMetadata
+from vllm.distributed import get_tensor_model_parallel_world_size, get_tp_group
+from vllm.distributed.parallel_state import get_dp_group
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM
+from vllm.distributed.parallel_state import get_ep_group
+from vllm.forward_context import get_forward_context
 
 
+from vllm_ascend.ascend_config import get_ascend_config
+from vllm_ascend.ops.fused_moe import AscendFusedMoE
+
 class CustomQwen3MoeForCausalLM(Qwen3MoeForCausalLM):
     packed_modules_mapping = {
         "qkv_proj": [
@@ -33,3 +49,86 @@ class CustomQwen3MoeForCausalLM(Qwen3MoeForCausalLM):
         "experts":
         ["experts.0.gate_proj", "experts.0.up_proj", "experts.0.down_proj"],
     }
+
+
+class AscendQwen3MoeSparseMoeBlock(nn.Module):
+    top_k: int
+
+    def __init__(
+            self,
+            config: PretrainedConfig,
+            quant_config: Optional[QuantizationConfig] = None,
+            prefix: str = "",
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        if self.tp_size > config.num_experts:
+            raise ValueError(
+                f"Tensor parallel size {self.tp_size} is greater than "
+                f"the number of experts {config.num_experts}.")
+
+        ascend_config = get_ascend_config()
+        self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
+        self.enable_multistream_moe = \
+            ascend_config.torchair_graph_config.enable_multistream_moe
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.num_experts,
+                                     bias=False,
+                                     quant_config=None,
+                                     prefix=f"{prefix}.gate")
+
+        self.experts = AscendFusedMoE(
+            num_experts=config.num_experts,
+            top_k=config.num_experts_per_tok,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=config.norm_topk_prob,
+            quant_config=quant_config,
+            prefix=f"{prefix}.experts")
+
+        self.top_k = config.num_experts_per_tok
+
+        self.dp_size = get_dp_group().world_size
+
+        self.tp_group = get_tp_group().device_group
+        self.tp_rank = get_tp_group().rank_in_group
+        self.ep_group = get_ep_group()
+
+        self.params_dtype = torch.get_default_dtype()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attn_metadata: Optional[AttentionMetadata] = None) -> torch.Tensor:
+        if attn_metadata is None:
+            attn_metadata = get_forward_context().attn_metadata
+        # when profile runs, force experts to load balanced tokens
+        # to avoid high memory consumption on a single rank.
+        # TODO: need a better flag to indicate whether in profile run or not.
+        if attn_metadata is None:
+            # for profile run
+            is_prefill = True
+            enable_force_load_balance = True
+        else:
+            is_prefill = get_forward_context().with_prefill
+            enable_force_load_balance = False
+            # if hasattr(attn_metadata, 'with_prefill_across_dp'):
+            #     is_prefill = attn_metadata.with_prefill_across_dp
+
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+
+        hidden_states = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits,
+            is_prefill=is_prefill,
+            top_k=self.top_k,
+            enable_force_load_balance=enable_force_load_balance,
+            shared_experts=None)
+
+        return hidden_states
+
+
+vllm.model_executor.models.qwen3_moe.Qwen3MoeSparseMoeBlock = AscendQwen3MoeSparseMoeBlock

vllm_ascend/multistream/ms_split.py

@@ -324,13 +324,13 @@ def model_input_split_v1_attn(
          query_start_loc=query_start_loc_pre,
          query_lens=query_lens_pre,
          seq_lens=seq_lens_pre,
+         seq_lens_list=seq_lens_pre.tolist(),
          max_query_len=max_query_len_pre,
          slot_mapping=slot_mapping_pre,
          is_only_prefill=is_only_prefill_pre,
          attn_state=attn_state_pre,
          attn_mask=attn_mask_pre,
          num_input_tokens=token_index,
-         with_prefill_across_dp=attn_metadata.with_prefill_across_dp,
          enable_dbo_across_dp=attn_metadata.enable_dbo_across_dp,
      )
 
@@ -340,13 +340,13 @@ def model_input_split_v1_attn(
          query_start_loc=query_start_loc_post,
          query_lens=query_lens_post,
          seq_lens=seq_lens_post,
+         seq_lens_list=seq_lens_post.tolist(),
          max_query_len=max_query_len_post,
          slot_mapping=slot_mapping_post,
          is_only_prefill=is_only_prefill_post,
          attn_state=attn_state_post,
          attn_mask=attn_mask_post,
          num_input_tokens=attn_metadata.num_input_tokens - token_index,
-         with_prefill_across_dp=attn_metadata.with_prefill_across_dp,
          enable_dbo_across_dp=attn_metadata.enable_dbo_across_dp,
      )
 

vllm_ascend/ops/fused_moe.py

@@ -989,7 +989,7 @@ def apply(
                 global_batch_size=self.global_batch_size,
                 expert_map=expert_map,
                 ep_group=get_ep_group())
-        elif fused_moe_state == FusedMoEState.All2AllSeq and is_prefill:
+        elif fused_moe_state == FusedMoEState.All2AllSeq:
             token_dispatcher = kwargs.get('token_dispatcher')
             return fused_experts_with_all2allv(token_dispatcher=token_dispatcher,
                                                probs=topk_weights,