[Refactor] Add build_dummy_metadata in attention backend and refactor common metadata

Signed-off-by: Yizhou Liu <liu_yizhou@outlook.com>

Author: yiz-liu

vllm_ascend/attention/attention_v1.py

@@ -27,10 +27,11 @@
 from vllm.config import get_current_vllm_config
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.utils import direct_register_custom_op
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu_input_batch import InputBatch
 
+from vllm_ascend.attention.utils import \
+    AscendCommonAttentionMetadata as CommonAttentionMetadata
 from vllm_ascend.ops.attention import vanilla_chunked_prefill
 from vllm_ascend.utils import get_graph_params
 
@@ -163,13 +164,16 @@ def build(self,
         block_table[:num_reqs, :self.runner.max_num_blocks_per_req] = (
             block_table[:num_reqs])
 
-        query_lens = self.runner.query_lens
+        query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens
-        seq_lens_list = self.runner.seq_lens_list
+        # TODO: Refactor these two param to common metadata in runners,
+        # preparing for the hybrid KV groups feature
+        query_lens = common_attn_metadata.query_lens if common_attn_metadata.query_lens is not None else self.runner.query_lens
+        seq_lens_list = common_attn_metadata.seq_lens_list if common_attn_metadata.seq_lens_list is not None else self.runner.seq_lens_list
+
         slot_mapping = self.runner.slot_mapping[:num_actual_tokens]
         attn_mask = self.runner.attn_mask
         attn_state = self.runner.attn_state
-        query_start_loc = common_attn_metadata.query_start_loc
 
         attn_metadata = AscendMetadata(
             num_actual_tokens=num_actual_tokens,
@@ -185,6 +189,34 @@ def build(self,
             enable_dbo_across_dp=enable_dbo_across_dp)
         return attn_metadata
 
+    def build_dummy_metadata(self, num_actual_tokens, num_reqs,
+                             num_scheduled_tokens, attn_state):
+        if attn_state == AscendAttentionState.DecodeOnly:
+            # NOTE: We only need to pay attention to seq_lens_list and block_table here
+            common_attn_metadata = CommonAttentionMetadata(seq_lens_list=[2] *
+                                                           num_reqs)
+
+            block_table = self.runner.input_batch.block_table[0].block_table
+            block_table[:num_reqs, 0] = torch.arange(1,
+                                                     num_reqs + 1,
+                                                     device=block_table.device,
+                                                     dtype=block_table.dtype)
+
+            attn_metadata = self.build(
+                num_reqs=num_reqs,
+                num_actual_tokens=num_actual_tokens,
+                max_query_len=num_scheduled_tokens.max(),
+                common_prefix_len=0,
+                common_attn_metadata=common_attn_metadata,
+            )
+        else:
+            raise NotImplementedError(
+                "Currently we only support building dummy metadata for DecodeOnly state"
+            )
+
+        attn_metadata.attn_state = attn_state
+        return attn_metadata
+
 
 class AscendAttentionBackendImpl(AttentionImpl):
 

vllm_ascend/attention/mla_v1.py

@@ -16,6 +16,8 @@
 from vllm_ascend import envs
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
+from vllm_ascend.attention.utils import \
+    AscendCommonAttentionMetadata as CommonAttentionMetadata
 from vllm_ascend.multistream.base import MSAttentionMetadataSplitConfig
 from vllm_ascend.multistream.context import get_multistream_comm_context
 from vllm_ascend.multistream.ms_split import model_input_split_v1_mla_attn
@@ -28,20 +30,6 @@
     from vllm.v1.worker.gpu_input_batch import InputBatch
 
 
-@dataclass
-class CommonAttentionMetadata:
-    """
-    Attention metadata attributes that can be shared by layers in different KV
-    cache groups and thus having different block table.
-    """
-
-    query_start_loc: torch.Tensor
-    """(batch_size + 1,), the start location of each request in query Tensor"""
-    seq_lens: torch.Tensor
-    """(batch_size,), the length of each request including both computed tokens
-    and newly scheduled tokens"""
-
-
 class AscendMLABackend(AttentionBackend):
 
     accept_output_buffer: bool = True

vllm_ascend/attention/utils.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+
+@dataclass
+class AscendCommonAttentionMetadata:
+    """
+    Attention metadata attributes that can be shared by layers in different KV
+    cache groups and thus having different block table.
+    """
+
+    query_start_loc: Optional[torch.Tensor] = None
+    """(batch_size + 1,), the start location of each request in query Tensor"""
+    seq_lens: Optional[torch.Tensor] = None
+    """(batch_size,), the length of each request including both computed tokens
+    and newly scheduled tokens"""
+    query_lens: Optional[torch.Tensor] = None
+    """(batch_size,), the length of each request including only the newly
+    scheduled tokens"""
+    seq_lens_list: Optional[list] = None
+    """(num_input_tokens,), note that this is specifically for FIA kernel"""

vllm_ascend/worker/model_runner_v1.py

@@ -77,7 +77,8 @@
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
 from vllm_ascend.attention.attention import AttentionMaskBuilder
 from vllm_ascend.attention.attention_v1 import AscendAttentionState
-from vllm_ascend.attention.mla_v1 import CommonAttentionMetadata
+from vllm_ascend.attention.utils import \
+    AscendCommonAttentionMetadata as CommonAttentionMetadata
 from vllm_ascend.multistream.ms_split import compute_split_seq_index
 from vllm_ascend.platform import NPUPlatform
 from vllm_ascend.sample.rejection_sampler import AscendRejectionSampler
@@ -253,6 +254,9 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.slot_mapping = torch.zeros(self.max_num_tokens,
                                         dtype=torch.int32,
                                         device=self.device)
+        self.query_lens = torch.zeros(self.max_num_reqs,
+                                      dtype=torch.int32,
+                                      device=self.device)
         # None in the first PP rank. The rest are set after load_model.
         self.intermediate_tensors: Optional[IntermediateTensors] = None
 
@@ -1528,6 +1532,7 @@ def _dummy_run(
         skip_attn: bool = True,
         with_prefill: bool = False,
         is_torchair_compile: bool = False,
+        attn_state: AscendAttentionState = AscendAttentionState.DecodeOnly,
     ) -> torch.Tensor:
         if self.torchair_graph_enabled and not with_prefill:
             num_tokens = self.select_torchair_padded_batch_size(num_tokens)
@@ -1558,43 +1563,13 @@ def _dummy_run(
         elif skip_attn:
             attn_metadata = None
         else:
-            query_start_loc = self.query_start_loc[:num_reqs + 1]
-            query_start_loc[:] = torch.arange(
-                query_start_loc.numel(),
-                device=query_start_loc.device,
-                dtype=query_start_loc.dtype,
-            )
-            seq_lens = self.seq_lens_np[:num_reqs]
-            seq_lens[:] = seq_lens + 2
-            self.seq_lens_list = self.seq_lens_np.tolist()[:num_tokens]
-
-            common_attn_metadata = CommonAttentionMetadata(
-                query_start_loc=query_start_loc, seq_lens=seq_lens)
-
-            self.query_lens = torch.from_numpy(num_scheduled_tokens)
-
-            block_table = self.input_batch.block_table[0].block_table
-            block_table[:num_reqs, 0] = torch.arange(1,
-                                                     num_reqs + 1,
-                                                     device=block_table.device,
-                                                     dtype=block_table.dtype)
-
-            self.slot_mapping[:num_tokens] = torch.arange(
-                1,
-                num_tokens + 1,
-                device=self.slot_mapping.device,
-                dtype=self.slot_mapping.dtype) * self.block_size + 1
-
-            attn_metadata = self.attn_metadata_builder.build(
-                num_reqs=num_reqs,
+            attn_metadata = self.attn_metadata_builder.build_dummy_metadata(
                 num_actual_tokens=num_tokens,
-                max_query_len=num_tokens,
-                common_prefix_len=0,
-                common_attn_metadata=common_attn_metadata,
+                num_reqs=num_reqs,
+                num_scheduled_tokens=num_scheduled_tokens,
+                attn_state=attn_state,
             )
 
-            attn_metadata.attn_state = AscendAttentionState.DecodeOnly
-
         with self.maybe_dummy_run_with_lora(self.lora_config,
                                             num_scheduled_tokens):
             model = self.model
@@ -1977,6 +1952,7 @@ def capture_model(self) -> None:
             # TODO(zzzzwwjj): Check dummy_run with ACL Graph and full graph mode
             with graph_capture(device=self.device):
                 skip_attn = not self.vllm_config.compilation_config.full_cuda_graph
+                # TODO: Make sure passing attn_state to _dummy_run in the future
                 for num_tokens in reversed(self.aclgraph_batch_sizes):
                     for _ in range(self.vllm_config.compilation_config.
                                    cudagraph_num_of_warmups):

vllm_ascend/worker/mtp_proposer_v1.py

@@ -8,7 +8,8 @@
 from vllm.v1.sample.metadata import SamplingMetadata
 
 from vllm_ascend.ascend_forward_context import set_ascend_forward_context
-from vllm_ascend.attention.mla_v1 import CommonAttentionMetadata
+from vllm_ascend.attention.utils import \
+    AscendCommonAttentionMetadata as CommonAttentionMetadata
 from vllm_ascend.models.deepseek_mtp import CustomDeepSeekMTP