fxi: the unit tests are updated by mocking the common metadata object instead of runner attributes. Additionally, the forward pass test now uses a forward context to manage graph capturing state.

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

Author: yiz-liu

tests/ut/attention/test_attention_v1.py

@@ -1,6 +1,7 @@
 from unittest.mock import MagicMock, patch
 
 import torch
+from vllm.forward_context import get_forward_context, set_forward_context
 
 from tests.ut.base import TestBase
 from vllm_ascend.attention.attention_v1 import (AscendAttentionBackend,
@@ -95,12 +96,15 @@ def test_build_prefill_no_cache(self, mock_is_310p, mock_nd_to_nz_2d,
             0].get_device_tensor.return_value = torch.zeros((10, 10))
         self.mock_runner.max_num_blocks_per_req = 10
         self.mock_runner.query_lens = torch.tensor([3, 4])
-        self.mock_runner.seq_lens_cpu = torch.tensor([5, 6])
         self.mock_runner.slot_mapping_cpu = torch.tensor(range(20))
         self.mock_runner.device = 'cpu:0'
         self.mock_runner.attn_mask = torch.ones((10, 10))
         self.mock_runner.attn_state = AscendAttentionState.PrefillNoCache
-        self.mock_runner.query_start_loc_cpu = torch.tensor([0, 3, 7])
+
+        mock_common_attn_metadata = MagicMock()
+        mock_common_attn_metadata.seq_lens = torch.tensor([5, 6])
+        mock_common_attn_metadata.query_start_loc = torch.tensor([0, 3, 7])
+        mock_common_attn_metadata.seq_lens_list = [5, 6]
 
         mock_nz_tensor = MagicMock()
         mock_nd_to_nz_2d.return_value = mock_nz_tensor
@@ -110,6 +114,7 @@ def test_build_prefill_no_cache(self, mock_is_310p, mock_nd_to_nz_2d,
             num_reqs,
             num_actual_tokens,
             max_query_len,
+            mock_common_attn_metadata,
         )
 
     @patch('vllm_ascend.attention.attention_v1.AscendMetadata')
@@ -129,12 +134,15 @@ def test_build_chunked_prefill(self, mock_ascend_attention_state,
             0].get_device_tensor.return_value = torch.zeros((10, 10))
         self.mock_runner.max_num_blocks_per_req = 10
         self.mock_runner.query_lens = torch.tensor([2, 3, 4])
-        self.mock_runner.seq_lens_cpu = torch.tensor([4, 5, 6])
         self.mock_runner.slot_mapping_cpu = torch.tensor(range(20))
         self.mock_runner.device = 'cpu:0'
         self.mock_runner.attn_mask = torch.ones((15, 15))
         self.mock_runner.attn_state = AscendAttentionState.ChunkedPrefill
-        self.mock_runner.query_start_loc_cpu = torch.tensor([0, 2, 5, 9])
+
+        mock_common_attn_metadata = MagicMock()
+        mock_common_attn_metadata.seq_lens = torch.tensor([4, 5, 6])
+        mock_common_attn_metadata.query_start_loc = torch.tensor([0, 2, 5, 9])
+        mock_common_attn_metadata.seq_lens_list = [4, 5, 6]
 
         mock_ascend_attention_state = MagicMock()
         mock_ascend_attention_state.PrefillNoCache = 0
@@ -143,7 +151,8 @@ def test_build_chunked_prefill(self, mock_ascend_attention_state,
         mock_nd_to_nz_spec.return_value = mock_nz_tensor
         mock_npu_format_cast.return_value = mock_nz_tensor
 
-        self.builder.build(num_reqs, num_actual_tokens, max_query_len)
+        self.builder.build(num_reqs, num_actual_tokens, max_query_len,
+                           mock_common_attn_metadata)
 
     @patch('vllm_ascend.attention.attention_v1.AscendMetadata')
     @patch('vllm_ascend.attention.attention_v1.is_310p', return_value=False)
@@ -157,14 +166,18 @@ def test_build_non_310p(self, mock_is_310p, mock_ascend_metadata):
             0].get_device_tensor.return_value = torch.zeros((10, 10))
         self.mock_runner.max_num_blocks_per_req = 10
         self.mock_runner.query_lens = torch.tensor([2, 3, 4])
-        self.mock_runner.seq_lens_cpu = torch.tensor([4, 5, 6])
         self.mock_runner.slot_mapping_cpu = torch.tensor(range(20))
         self.mock_runner.device = 'cpu:0'
         self.mock_runner.attn_mask = torch.ones((15, 15))
         self.mock_runner.attn_state = AscendAttentionState.ChunkedPrefill
-        self.mock_runner.query_start_loc_cpu = torch.tensor([0, 2, 5, 9])
 
-        self.builder.build(num_reqs, num_actual_tokens, max_query_len)
+        mock_common_attn_metadata = MagicMock()
+        mock_common_attn_metadata.seq_lens = torch.tensor([4, 5, 6])
+        mock_common_attn_metadata.query_start_loc = torch.tensor([0, 2, 5, 9])
+        mock_common_attn_metadata.seq_lens_list = [4, 5, 6]
+
+        self.builder.build(num_reqs, num_actual_tokens, max_query_len,
+                           mock_common_attn_metadata)
 
 
 class TestAscendAttentionBackendImpl(TestBase):
@@ -372,13 +385,19 @@ def test_forward_decode_only(self, mock_paged_attention,
         metadata.slot_mapping = torch.zeros(10, dtype=torch.long)
         layer = self.layer_no_quant
 
-        output = self.impl.forward(layer,
-                                   query,
-                                   key,
-                                   value,
-                                   kv_cache,
-                                   metadata,
-                                   trace_flag=False)
+        vllm_config = MagicMock()
+        vllm_config.parallel_config.data_parallel_size = 1
+        with set_forward_context(attn_metadata=metadata,
+                                 vllm_config=vllm_config):
+            forward_context = get_forward_context()
+            forward_context.capturing = False
+            output = self.impl.forward(layer,
+                                       query,
+                                       key,
+                                       value,
+                                       kv_cache,
+                                       metadata,
+                                       trace_flag=False)
 
         mock_paged_attention.assert_called_once()
         assert output.shape == (10, 8 * 64)

vllm_ascend/attention/attention_v1.py

@@ -167,7 +167,10 @@ def reorder_batch(self, input_batch: "InputBatch",
                       scheduler_output: "SchedulerOutput") -> bool:
         return False
 
-    def build(self, num_reqs, num_actual_tokens, max_query_len,
+    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):
@@ -179,9 +182,9 @@ def build(self, num_reqs, num_actual_tokens, max_query_len,
 
         query_start_loc = common_attn_metadata.query_start_loc
         seq_lens = common_attn_metadata.seq_lens
-        # TODO: Refactor these two param to common metadata in runners,
+        # TODO: Refactor this 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 = self.runner.query_lens
         # Since FIA for GQA is not active now, we temporarily silence it
         seq_lens_list = common_attn_metadata.seq_lens_list
 
@@ -407,6 +410,8 @@ def forward(
             elif attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
                 graph_params = get_graph_params()
 
+                # TODO(Yizhou): Find another way to handle the
+                # graph capturing mode. Say, GraphCaptureContext?
                 forward_context = get_forward_context()
                 if not forward_context.capturing:
                     if is_310p():

vllm_ascend/worker/model_runner_v1.py

@@ -79,9 +79,9 @@
                                                 AscendMetadata)
 from vllm_ascend.attention.attention_v1_torchair import AscendTorchairMetadata
 from vllm_ascend.attention.mla_v1 import AscendMLAMetadata
-from vllm_ascend.multistream.ms_split import compute_split_seq_index
 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
 from vllm_ascend.utils import (ACL_FORMAT_FRACTAL_ND, ACL_FORMAT_FRACTAL_NZ,