bugfix for torchair

Signed-off-by: panchao <panchao13@huawei.com>

Author: zhangdepeng

vllm_ascend/ascend_config.py

@@ -60,6 +60,7 @@ class TorchairGraphConfig:
 
     def __init__(self, torchair_graph_config):
         self.enabled = torchair_graph_config.get("enabled", False)
+        self.mode = torchair_graph_config.get("mode", 'max-autotune')
         self.use_cached_graph = torchair_graph_config.get(
             "use_cached_graph", False)
         self.graph_batch_sizes = torchair_graph_config.get(

vllm_ascend/torchair/torchair_attention.py

@@ -304,6 +304,7 @@ def __init__(
         self.num_queries_per_kv = self.num_heads // self.num_kv_heads
         self.key_cache = None
         self.value_cache = None
+        self.scale_tensor = torch.zeros((), device='npu' ,dtype=torch.int32)
 
     def forward(
         self,
@@ -366,7 +367,7 @@ def forward(
             key_cache, value_cache = kv_cache[0], kv_cache[1]
             slots = attn_metadata.slot_mapping
 
-            block_size = key_cache.shape[1]
+            block_size = self.scale_tensor + key_cache.shape[1]
             slots_indices = slots.reshape(-1, 1)
             block_indices = slots_indices // block_size
             slots_indices = slots_indices % block_size

vllm_ascend/torchair/torchair_model_runner.py

@@ -316,6 +316,7 @@ def _get_torchair_lazy_compiled_model(self, batch_size: int):
             communication_adaptation_310p()
 
         config = torchair.CompilerConfig()
+        config.mode = get_ascend_config().torchair_graph_config.mode
         config.experimental_config.frozen_parameter = True
         # enabling tiling_schedule_optimize on 300I Duo has some bugs, so we have to
         # disable it on 300I Duo platform now.

vllm_ascend/torchair/utils.py

@@ -168,7 +168,7 @@ def register_torchair_model():
         "vllm_ascend.torchair.models.qwen2:CustomQwen2ForCausalLM")
 
     ModelRegistry.register_model(
-        "Qwen3ForCausalLM",
+        "Qwen3MoeForCausalLM",
         "vllm_ascend.torchair.models.qwen3_moe:CustomQwen3MoeForCausalLM")