refact attn metadata build

Signed-off-by: weiguihua2 <[email protected]>

Author: weiguihua2

vllm_ascend/attention/attention_v1_torchair.py

@@ -33,7 +33,7 @@
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_NZ, aligned_16, is_310p,
                                nd_to_nz_2d)
 from vllm_ascend.worker.npu_input_batch import InputBatch
-from vllm_ascend.attention.utils import AscendCommonAttentionMetadata, get_decode_token_per_req
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 
 
 class AscendAttentionTorchairBackend(AttentionBackend):
@@ -157,7 +157,7 @@ def __init__(self,
         self.device = device
         self.max_num_blocks_per_req = cdiv(self.model_config.max_model_len,
                                            vllm_config.cache_config.block_size)
-        self.decode_token_per_req = get_decode_token_per_req(vllm_config.speculative_config)
+        self.max_blocks = (self.model_config.max_model_len + vllm_config.cache_config.block_size - 1) // vllm_config.cache_config.block_size
 
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
@@ -175,7 +175,7 @@ def build_torchair_graph_dummy(
             self, common_attn_metadata: AscendCommonAttentionMetadata) -> AscendTorchairMetadata:
         device = self.device
         num_reqs = common_attn_metadata.num_reqs
-        _, max_blocks = self.runner.graph_block_tables.shape
+        _, max_blocks = self.max_blocks
         block_table = torch.zeros((num_reqs, max_blocks),
                                   dtype=torch.int32,
                                   device=device)
@@ -257,7 +257,7 @@ def build(self,
                     pad_value = 0
                     num_token_pad_size = graph_pad_size - num_actual_tokens
                     num_reqs_pad_size = (
-                        graph_pad_size // self.decode_token_per_req -
+                        graph_pad_size // common_attn_metadata.decode_token_per_req -
                         num_reqs)
                 pad_value = 1
                 padded_seq_lens = seq_lens.tolist() + [pad_value

vllm_ascend/attention/mla_v1.py

@@ -25,7 +25,7 @@
 from vllm_ascend.torchair.utils import npu_stream_switch, npu_wait_tensor
 from vllm_ascend.utils import npu_prefetch
 from vllm_ascend.worker.npu_input_batch import InputBatch
-from vllm_ascend.attention.utils import (AscendCommonAttentionMetadata,split_decodes_and_prefills, get_decode_token_per_req)
+from vllm_ascend.attention.utils import (AscendCommonAttentionMetadata,split_decodes_and_prefills)
 
 
 if TYPE_CHECKING:
@@ -186,7 +186,6 @@ def __init__(self,
         scheduler_config = vllm_config.scheduler_config
         self.block_size = vllm_config.cache_config.block_size
         self.max_blocks = (vllm_config.model_config.max_model_len + self.block_size - 1) // self.block_size
-        self.decode_token_per_req = get_decode_token_per_req(vllm_config.speculative_config)
         self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
         if self.chunked_prefill_enabled:
             self.chunked_prefill_workspace_size = min(
@@ -288,13 +287,13 @@ def build_torchair_graph_dummy(
             self, common_attn_metadata: AscendCommonAttentionMetadata,) -> AscendMLAMetadata:
         device = self.device
         num_reqs = common_attn_metadata.num_reqs
-        _, max_blocks = self.runner.graph_block_tables.shape
+        _, max_blocks = self.max_blocks
         block_table = torch.zeros((num_reqs, max_blocks),
                                   dtype=torch.int32,
                                   device=device)
         block_table = self._get_graph_runner_block_tables(
             num_reqs, block_table)
-        num_tokens = num_reqs * self.decode_token_per_req
+        num_tokens = num_reqs * common_attn_metadata.decode_token_per_req
         seq_lens = torch.zeros(num_reqs, dtype=torch.int32, device=device)
         seq_lens_list = [0] * num_reqs
         input_positions = torch.zeros(num_tokens,
@@ -382,8 +381,8 @@ def build(
         input_positions = common_attn_metadata.positions[:num_actual_tokens].long()
 
         if self.cos_cache is None:
-            self.cos_cache = model.layers[0].self_attn.rotary_emb.cos_cached
-            self.sin_cache = model.layers[0].self_attn.rotary_emb.sin_cached
+            self.cos_cache = model.model.layers[0].self_attn.rotary_emb.cos_cached
+            self.sin_cache = model.model.layers[0].self_attn.rotary_emb.sin_cached
         if self.cos_cache.dtype != self.model_config.dtype:  # type: ignore
             self.cos_cache = self.cos_cache.to(  # type: ignore
                 self.model_config.dtype)  # type: ignore
@@ -392,10 +391,9 @@ def build(
 
         query_seq_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         query_lens = query_seq_lens_cpu[:num_reqs]
-        num_computed_tokens_cpu = (common_attn_metadata.seq_lens_cpu -
-                                   query_seq_lens_cpu)
-        
         seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs]
+        num_computed_tokens_cpu = (seq_lens - query_lens)
+        
         prefill_metadata = None
         chunked_context_metadata = None
         if num_prefills > 0:
@@ -418,12 +416,12 @@ def build(
                 assert max_context_chunk > 0
                 num_chunks = cdiv(max_context_len_cpu, max_context_chunk)
                 chunk_starts = torch.arange(num_chunks, dtype=torch.int32) \
-                    .unsqueeze(1).expand(-1, self._num_prefills) * max_context_chunk
+                    .unsqueeze(1).expand(-1, num_prefills) * max_context_chunk
                 chunk_ends = torch.min(context_lens_cpu.unsqueeze(0),
                                        chunk_starts + max_context_chunk)
                 chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
                 cu_seq_lens_cpu = torch.zeros(num_chunks,
-                                              self._num_prefills + 1,
+                                              num_prefills + 1,
                                               dtype=torch.int32,
                                               pin_memory=True)
                 torch.cumsum(chunk_seq_lens,

vllm_ascend/attention/utils.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
+from enum import Enum
 
 from vllm.config import SpeculativeConfig
-from vllm_ascend.attention.attention_v1 import AscendAttentionState
 
 import torch
 
@@ -28,18 +28,20 @@ class AscendCommonAttentionMetadata:
     num_actual_tokens: int
     """Total number of tokens in batch"""
 
-    actual_seq_lengths_q: list[int] = None
+    max_query_len: int
+
+    decode_token_per_req: int
 
     block_table_tensor: torch.Tensor
     slot_mapping_cpu: torch.Tensor
 
+    actual_seq_lengths_q: list[int] = None
+
     positions: torch.Tensor = None
 
     attn_mask: torch.Tensor = None
     spec_attn_mask: torch.Tensor = None
-    attn_state: AscendAttentionState = None
-
-    max_query_len: int
+    attn_state: Enum = None
 
     enable_dbo_across_dp: bool = False
 
@@ -61,6 +63,8 @@ class TorchairCommonAttentionMetadata:
     num_actual_tokens: int
     """Total number of tokens in batch"""
 
+    decode_token_per_req: int
+
     actual_seq_lengths_q: list[int] = None
 
     attn_mask: torch.Tensor = None
@@ -110,11 +114,3 @@ def split_decodes_and_prefills(
     num_prefill_tokens = num_tokens - num_decode_tokens
     return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens)
 
-
-def get_decode_token_per_req(speculative_config: SpeculativeConfig):
-    decode_token_per_req = 1
-    if not speculative_config:
-        return decode_token_per_req
-    spec_token_num = speculative_config.num_speculative_tokens
-    assert spec_token_num > 0
-    return decode_token_per_req + spec_token_num

vllm_ascend/worker/eagle_proposer_v1.py

@@ -140,9 +140,10 @@ def propose(
             attn_mask=self.runner.attn_mask,
             spec_attn_mask=self.runner.spec_attn_mask,
             attn_state=self.runner.attn_state,
+            decode_token_per_req=self.runner.decode_token_per_req,
         )
         # FIXME(woosuk): The below two ops cause synchronization. Optimize.
-        attn_metadata = self.runner.attn_metadata_builder.build(common_attn_metadata)
+        attn_metadata = self.runner.attn_metadata_builder.build(common_attn_metadata, self.runner.model)
         if self.use_cuda_graph and \
             num_tokens <= self.cudagraph_batch_sizes[-1]:
             num_input_tokens = self.vllm_config.pad_for_cudagraph(num_tokens)

vllm_ascend/worker/model_runner_v1.py

@@ -92,7 +92,7 @@
 from vllm_ascend.worker.eagle_proposer_v1 import EagleProposer
 from vllm_ascend.worker.mtp_proposer_v1 import MtpProposer
 from vllm_ascend.worker.npu_input_batch import CachedRequestState, InputBatch
-from vllm_ascend.attention.utils import AscendCommonAttentionMetadata, get_decode_token_per_req
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 
 if TYPE_CHECKING:
     import xgrammar as xgr  # type: ignore[import-untyped]
@@ -229,9 +229,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.drafter: Optional[Union[NgramProposer, EagleProposer,
                                      MtpProposer]] = None
         self.actual_seq_lengths_q = []
-        self.decode_token_per_req = get_decode_token_per_req(self.speculative_config)
+        self.decode_token_per_req = 1
         if self.speculative_config:
             self.use_spec_decode = True
+            spec_token_num = self.speculative_config.num_speculative_tokens
+            assert spec_token_num > 0
+            self.decode_token_per_req = 1 + spec_token_num
             self.actual_seq_lengths_q = [
                 len for len in
                 range(self.decode_token_per_req, self.max_num_tokens +
@@ -810,8 +813,9 @@ def get_eagle_atten_dict(
                 spec_attn_mask=self.spec_attn_mask,
                 attn_state=self.attn_state,
                 max_num_blocks_per_req=self.max_num_blocks_per_req,
+                decode_token_per_req=self.decode_token_per_req,
             )
-            attn_metadata_i = self.attn_metadata_builder.build(common_attn_metadata)
+            attn_metadata_i = self.attn_metadata_builder.build(common_attn_metadata, self.model)
             for layer_name in kv_cache_group_spec.layer_names:
                 attn_metadata[layer_name] = attn_metadata_i
 
@@ -1228,9 +1232,11 @@ def _process_reqs(
             attn_state=self.attn_state,
             enable_dbo_across_dp=enable_dbo,
             is_only_prefill=is_only_prefill,
-            graph_pad_size=self.graph_pad_size
+            max_query_len=max_num_scheduled_tokens,
+            graph_pad_size=self.graph_pad_size,
+            decode_token_per_req=self.decode_token_per_req,
         )
-        attn_metadata = self.attn_metadata_builder.build(common_attn_metadata)
+        attn_metadata = self.attn_metadata_builder.build(common_attn_metadata, self.model)
         if self.vllm_config.model_config.use_mla:
             attn_metadata.num_input_tokens = num_input_tokens
 

vllm_ascend/worker/mtp_proposer_v1.py

@@ -181,9 +181,10 @@ def propose(
             attn_mask=self.runner.attn_mask,
             spec_attn_mask=self.runner.spec_attn_mask,
             attn_state=self.runner.attn_state,
-            graph_pad_size=extra_builder_kwargs['graph_pad_size']
+            graph_pad_size=extra_builder_kwargs['graph_pad_size'],
+            decode_token_per_req=self.runner.decode_token_per_req,
         )
-        attn_metadata = self.runner.attn_metadata_builder.build(common_attn_metadata)
+        attn_metadata = self.runner.attn_metadata_builder.build(common_attn_metadata, self.runner.model)
 
         self.positions[:num_tokens] = target_positions
         self.hidden_states[:num_tokens] = target_hidden_states
@@ -294,12 +295,13 @@ def dummy_run(self,
             attn_metadata = None
         else:
             common_attn_metadata = TorchairCommonAttentionMetadata(
-            num_reqs=num_reqs,
-            num_actual_tokens=1,
-            actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
-            attn_mask=self.runner.attn_mask,
-            spec_attn_mask=self.runner.spec_attn_mask,
-        )
+                num_reqs=num_reqs,
+                num_actual_tokens=1,
+                actual_seq_lengths_q=self.runner.actual_seq_lengths_q,
+                attn_mask=self.runner.attn_mask,
+                spec_attn_mask=self.runner.spec_attn_mask,
+                decode_token_per_req=self.runner.decode_token_per_req,
+            )
             attn_metadata = self.runner.attn_metadata_builder.build_torchair_graph_dummy(common_attn_metadata)
 
         input_ids = self.input_ids[:num_tokens]