Refactor mla

Author: JC-ut0

vllm_ascend/attention/attention_v1.py

@@ -21,6 +21,8 @@
 
 import torch
 import torch_npu
+import torch.nn as nn
+from vllm.config import VllmConfig
 from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
                                               AttentionLayer, AttentionType)
 from vllm.attention.backends.utils import CommonAttentionState
@@ -35,6 +37,7 @@
 from vllm_ascend.multistream.base import MSAttentionMetadataSplitConfig
 from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import get_graph_params
+from vllm_ascend.attention.utils import AscendCommonAttentionMetadata
 
 
 class AscendAttentionBackend(AttentionBackend):
@@ -156,39 +159,48 @@ def split_metadata_for_multistream(
 
 class AscendAttentionMetadataBuilder:
 
-    def __init__(self, runner):
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        device: torch.device,
+        runner
+    ):
+        self.vllm_config = vllm_config
+        self.model_config = vllm_config.model_config
+        self.device = device
         self.runner = runner
 
     def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
         return False
 
-    def build(self,
-              num_reqs,
-              num_actual_tokens,
-              max_query_len,
-              common_attn_metadata: CommonAttentionMetadata,
-              enable_dbo_across_dp: bool = False,
-              is_only_prefill: bool = False,
-              *args,
-              **kwargs):
-
-        block_table = self.runner.input_batch.block_table[0].get_device_tensor(
-        )
-        block_table[:num_reqs, :self.runner.max_num_blocks_per_req] = (
-            block_table[:num_reqs])
-
-        query_start_loc = common_attn_metadata.query_start_loc
-        seq_lens = common_attn_metadata.seq_lens
+    def build(
+        self,
+        common_attn_metadata: AscendCommonAttentionMetadata,
+    ):
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[:
+                                                                       num_reqs
+                                                                       + 1]
+
+        block_table = common_attn_metadata.block_table_tensor
+        block_table[:num_reqs, :common_attn_metadata.
+                    max_num_blocks_per_req] = (block_table[:num_reqs])
+
+        seq_lens = common_attn_metadata.seq_lens_cpu[:num_reqs]
         # TODO: Refactor these two param to common metadata in runners,
         # preparing for the hybrid KV groups feature
-        query_lens = common_attn_metadata.query_lens or self.runner.query_lens
+        query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
         # Since FIA for GQA is not active now, we temporarily silence it
         seq_lens_list = common_attn_metadata.seq_lens_list
 
-        slot_mapping = self.runner.slot_mapping[:num_actual_tokens]
-        attn_mask = self.runner.attn_mask
-        attn_state = self.runner.attn_state
+        slot_mapping = common_attn_metadata.slot_mapping_cpu[:num_actual_tokens].to(
+            self.device, non_blocking=True)
+        attn_mask = common_attn_metadata.attn_mask
+        attn_state = common_attn_metadata.attn_state
+        query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[:num_reqs + 1]
+        query_start_loc = query_start_loc_cpu.to(self.device,non_blocking=True)
 
         attn_metadata = AscendMetadata(
             num_actual_tokens=num_actual_tokens,
@@ -197,12 +209,12 @@ def build(self,
             query_lens=query_lens,
             seq_lens=seq_lens,
             seq_lens_list=seq_lens_list,
-            max_query_len=max_query_len,
+            max_query_len=common_attn_metadata.max_query_len,
             slot_mapping=slot_mapping,
             attn_mask=attn_mask,
             attn_state=attn_state,
-            enable_dbo_across_dp=enable_dbo_across_dp,
-            is_only_prefill=is_only_prefill)
+            enable_dbo_across_dp=common_attn_metadata.enable_dbo_across_dp,
+            is_only_prefill=common_attn_metadata.is_only_prefill)
         return attn_metadata
 
     def build_dummy_metadata(self, num_actual_tokens, num_reqs,
@@ -217,14 +229,33 @@ def build_dummy_metadata(self, num_actual_tokens, num_reqs,
                                                      num_reqs + 1,
                                                      device=block_table.device,
                                                      dtype=block_table.dtype)
+            block_table = self.runner.input_batch.block_table[0].get_device_tensor(
+            )
+            block_table[:num_reqs, :self.runner.max_num_blocks_per_req] = (
+                block_table[:num_reqs])
+
+            query_start_loc = common_attn_metadata.query_start_loc
+            seq_lens = common_attn_metadata.seq_lens
+            query_lens = self.runner.query_lens
+            seq_lens_list = None
 
-            attn_metadata = self.build(
-                num_reqs=num_reqs,
+            slot_mapping = self.runner.slot_mapping[:num_actual_tokens]
+            attn_mask = self.runner.attn_mask
+            attn_state = self.runner.attn_state
+
+            attn_metadata = AscendMetadata(
                 num_actual_tokens=num_actual_tokens,
+                block_tables=block_table,
+                query_start_loc=query_start_loc,
+                query_lens=query_lens,
+                seq_lens=seq_lens,
+                seq_lens_list=seq_lens_list,
                 max_query_len=num_scheduled_tokens.max(),
-                common_prefix_len=0,
-                common_attn_metadata=common_attn_metadata,
-            )
+                slot_mapping=slot_mapping,
+                attn_mask=attn_mask,
+                attn_state=attn_state,
+                enable_dbo_across_dp=False,
+                is_only_prefill=False)
         else:
             raise NotImplementedError(
                 "Currently we only support building dummy metadata for DecodeOnly state"