[bugfix] fix torchair and mtp functionality

Co-authored-by: hust17yixuan <303660421@qq.com>

Signed-off-by: linfeng-yuan <1102311262@qq.com>

Author: linfeng-yuan

vllm_ascend/spec_decode/mtp_proposer.py

@@ -1,4 +1,5 @@
 import types
+from typing import Dict
 
 import torch
 import torch.nn as nn
@@ -186,6 +187,8 @@ def generate_token_ids(self,
                            hidden_states: torch.Tensor = None,
                            attn_metadata=None,
                            aux_hidden_states: torch.Tensor = None):
+        if attn_metadata is not None and isinstance(attn_metadata, Dict):
+            attn_metadata = attn_metadata['model.layers.0.self_attn.attn']
         next_token_ids: list[int] = []
         for i, token_ids in enumerate(valid_sampled_token_ids):
             if token_ids:
@@ -382,7 +385,7 @@ def _propose(
             num_computed_tokens_cpu=None,
             seq_lens=None)
         attn_metadata = self.runner.attn_metadata_builder.build(
-            common_attn_metadata, self.runner.get_model())
+            0, common_attn_metadata, self.runner.get_model())
 
         self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states

vllm_ascend/torchair/torchair_model_runner.py

@@ -19,7 +19,7 @@
 
 import math
 import types
-from typing import Optional
+from typing import Optional, Dict
 
 import torch
 import torch.distributed as dist
@@ -51,7 +51,7 @@ class NPUTorchairModelRunner(NPUModelRunner):
     def __init__(self, vllm_config: VllmConfig, device: torch.device):
         super().__init__(vllm_config, device)
         self.attn_metadata_builder = self.attn_backend.get_builder_cls()(
-            vllm_config, device)
+            None, None, vllm_config, device)
 
         ascend_config = get_ascend_config()
         self.new_kv_cache_bytes = -1
@@ -135,7 +135,8 @@ def _generate_dummy_run_hidden_states(self, with_prefill,
                                           is_torchair_compile, input_ids,
                                           positions, attn_metadata, num_tokens,
                                           intermediate_tensors, inputs_embeds):
-
+        if attn_metadata is not None and isinstance(attn_metadata, Dict):
+            attn_metadata = attn_metadata['model.layers.0.self_attn.attn']
         if not with_prefill:
             # Only mark static while compiling
             if is_torchair_compile:
@@ -281,6 +282,8 @@ def _generate_process_reqs_hidden_states(self, attn_metadata, with_prefill,
                                              input_ids, positions,
                                              intermediate_tensors,
                                              inputs_embeds):
+        if attn_metadata is not None and isinstance(attn_metadata, Dict):
+            attn_metadata = attn_metadata['model.layers.0.self_attn.attn']
         model_kwargs = {
             "kv_caches": self.kv_caches,
             "attn_metadata": attn_metadata