Merge pull request vllm-project#7 from wangxiyuan/upodate_model

Update qwen3 moe

Author: wangxiyuan

vllm_ascend/attention/attention_v1.py

@@ -17,7 +17,7 @@
 
 from dataclasses import dataclass
 from enum import Enum
-from typing import List, Optional, Tuple, Type, ClassVar
+from typing import ClassVar, List, Optional, Tuple, Type
 
 import torch
 import torch.nn as nn
@@ -32,13 +32,12 @@
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.utils import cdiv, direct_register_custom_op
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.kv_cache_interface import AttentionSpec
 
 from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, aligned_16, is_310p,
                                nd_to_nz_2d, nd_to_nz_spec)
-from vllm_ascend.worker.npu_input_batch import InputBatch
-from vllm.v1.kv_cache_interface import AttentionSpec
 
 
 def wait_for_kv_layer_from_connector(layer_name: str):
@@ -582,7 +581,7 @@ def unified_ascend_attention_with_output(
         attn_metadata = attn_metadata[layer_name]
     self = forward_context.no_compile_layers[layer_name]
     kv_cache = self.kv_cache[forward_context.virtual_engine]
-    print(100*"^", f"layer_name: {layer_name}")
+    print(100 * "^", f"layer_name: {layer_name}")
     self.impl.forward(self,
                       query,
                       key,

vllm_ascend/models/__init__.py

@@ -53,3 +53,6 @@ def register_model():
         "PanguProMoEForCausalLM",
         "vllm_ascend.torchair.models.torchair_pangu_moe:PanguProMoEForCausalLM"
     )
+    ModelRegistry.register_model(
+        "Qwen3NextForCausalLM",
+        "vllm_ascend.torchair.models.qwen3_next:Qwen3NextForCausalLM")