[Bugfix] Support Qwen3-MOE on aclgraph mode

tests/e2e/singlecard/test_aclgraph.py

@@ -29,7 +29,7 @@
 from tests.conftest import VllmRunner
 from tests.model_utils import check_outputs_equal
 
-MODELS = ["Qwen/Qwen2.5-0.5B-Instruct"]
+MODELS = ["Qwen/Qwen2.5-0.5B-Instruct", "vllm-ascend/Qwen3-30B-A3B-Puring"]
 
 
 @pytest.mark.skipif(os.getenv("VLLM_USE_V1") == "0",

vllm_ascend/ops/common_fused_moe.py

@@ -18,13 +18,23 @@
 from typing import Callable, Optional
 
 import torch
+from vllm.config import CompilationLevel, get_current_vllm_config
 from vllm.model_executor.layers.fused_moe.layer import \
     UnquantizedFusedMoEMethod
 
 from vllm_ascend.ops.fused_moe import (fused_experts, fused_experts_moge,
                                        select_experts)
 from vllm_ascend.utils import is_310p
 
+original_unquantized_fused_moe_init_func = UnquantizedFusedMoEMethod.__init__
+
+
+def unquantized_fused_moe_init_func(self, *args, **kwargs):
+    original_unquantized_fused_moe_init_func(self, *args, **kwargs)
+    vllm_config = get_current_vllm_config()
+    self.max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
+    self.use_aclgraph = vllm_config.compilation_config.level == CompilationLevel.PIECEWISE and not vllm_config.model_config.enforce_eager
+
 
 def forward_oot(
     self,
@@ -71,6 +81,10 @@
             expert_map=expert_map,
             apply_router_weight_on_input=apply_router_weight_on_input)
 
+    # If use aclgraph, we need to set max_num_tokens to make
+    # the input shape of `npu_moe_init_routing` fixed
+    max_num_tokens = self.max_num_batched_tokens if self.use_aclgraph else None
+
     return fused_experts(
         hidden_states=x,
         w1=layer.w13_weight,
@@ -79,7 +93,9 @@
         topk_ids=topk_ids,
         top_k=top_k,
         expert_map=expert_map,
-        apply_router_weight_on_input=apply_router_weight_on_input)
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        max_num_tokens=max_num_tokens)
 
 
+UnquantizedFusedMoEMethod.__init__ = unquantized_fused_moe_init_func
 UnquantizedFusedMoEMethod.forward_oot = forward_oot

vllm_ascend/ops/fused_moe.py

@@ -655,6 +655,7 @@
     top_k: int,
     expert_map: torch.Tensor = None,
     apply_router_weight_on_input: bool = False,
+    max_num_tokens: Optional[int] = None,
 ) -> torch.Tensor:
     """
     Fused experts with top-k routing.
@@ -748,11 +749,12 @@
                                 dtype=torch.int32,
                                 device=device).view(top_k, -1).permute(
                                     1, 0).contiguous())
+        active_num = max_num_tokens if max_num_tokens is not None else num_tokens
         sorted_hidden_states, expanded_row_idx, expanded_expert_idx = torch_npu.npu_moe_init_routing(
             hidden_states,
             row_idx=row_idx,
             expert_idx=topk_ids,
-            active_num=num_tokens)
+            active_num=active_num)
 
         expert_tokens = torch_npu.npu_moe_compute_expert_tokens(
             expanded_expert_idx, num_experts)