refact attn metadata build

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

Author: weiguihua2

tests/ut/attention/test_attention_v1.py

@@ -72,7 +72,8 @@ def setUp(self):
         self.mock_vllm_config.model_config.max_model_len = 640
         self.mock_vllm_config.cache_config.block_size = 64
         self.mock_device = 'cpu:0'
-        self.builder = AscendAttentionMetadataBuilder(self.mock_vllm_config, self.mock_device)
+        self.builder = AscendAttentionMetadataBuilder(self.mock_vllm_config,
+                                                      self.mock_device)
 
     def test_reorder_batch(self):
         mock_input_batch = MagicMock()
@@ -104,18 +105,14 @@ def test_build_prefill_no_cache(self, mock_is_310p, mock_nd_to_nz_2d,
             positions=torch.tensor([10, 10]),
             attn_mask=torch.ones((10, 10)),
             spec_attn_mask=None,
-            attn_state=AscendAttentionState.PrefillNoCache
-        )
+            attn_state=AscendAttentionState.PrefillNoCache)
 
         mock_nz_tensor = MagicMock()
         mock_model = MagicMock()
         mock_nd_to_nz_2d.return_value = mock_nz_tensor
         mock_npu_format_cast.return_value = mock_nz_tensor
 
-        self.builder.build(
-            common_attn_metadata,
-            mock_model
-        )
+        self.builder.build(common_attn_metadata, mock_model)
 
     @patch('vllm_ascend.attention.attention_v1.AscendMetadata')
     @patch('torch_npu.npu_format_cast')
@@ -139,8 +136,7 @@ def test_build_chunked_prefill(self, mock_ascend_attention_state,
             positions=torch.tensor([10, 10]),
             attn_mask=torch.ones((15, 15)),
             spec_attn_mask=None,
-            attn_state=AscendAttentionState.ChunkedPrefill
-        )
+            attn_state=AscendAttentionState.ChunkedPrefill)
 
         mock_ascend_attention_state = MagicMock()
         mock_ascend_attention_state.PrefillNoCache = 0
@@ -169,8 +165,7 @@ def test_build_non_310p(self, mock_is_310p, mock_ascend_metadata):
             positions=torch.tensor([10, 10]),
             attn_mask=torch.ones((15, 15)),
             spec_attn_mask=None,
-            attn_state=AscendAttentionState.ChunkedPrefill
-        )
+            attn_state=AscendAttentionState.ChunkedPrefill)
         mock_model = MagicMock()
 
         self.builder.build(common_attn_metadata, mock_model)

tests/ut/attention/test_mla_v1.py

@@ -1,6 +1,5 @@
 from unittest.mock import MagicMock, patch
 
-import numpy as np
 import torch
 from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.model_executor.layers.linear import LinearBase
@@ -195,9 +194,11 @@ def test_ascend_mla_metadata_builder_default(self):
                    return_value=ascend_config):
             builder = AscendMLAMetadataBuilder(mock_vllm_config, mock_device)
 
-            self.assertEqual(builder.block_size, mock_vllm_config.cache_config.block_size)
-            self.assertEqual(builder.chunked_prefill_enabled,
-                             mock_vllm_config.scheduler_config.chunked_prefill_enabled)
+            self.assertEqual(builder.block_size,
+                             mock_vllm_config.cache_config.block_size)
+            self.assertEqual(
+                builder.chunked_prefill_enabled,
+                mock_vllm_config.scheduler_config.chunked_prefill_enabled)
             self.assertEqual(builder.torchair_graph_enabled, True)
 
     @patch("vllm_ascend.attention.mla_v1.get_ascend_config")
@@ -363,7 +364,7 @@ def test_build_dummy(self, mock_ascend_config):
                 num_reqs=3,
                 num_actual_tokens=3,
                 decode_token_per_req=1,
-                actual_seq_lengths_q=[0,1,2],
+                actual_seq_lengths_q=[0, 1, 2],
                 attn_mask=torch.zeros((1, 1), dtype=torch.bool),
                 spec_attn_mask=torch.zeros((1, 1), dtype=torch.bool),
             )

vllm_ascend/attention/mla_v1.py

@@ -282,8 +282,8 @@ def _get_graph_runner_block_tables(
         max_blocks = self.max_blocks
 
         graph_block_tables = torch.zeros((num_seqs, max_blocks),
-                                             dtype=block_tables.dtype,
-                                             device=block_tables.device)
+                                         dtype=block_tables.dtype,
+                                         device=block_tables.device)
 
         num_blocks = block_tables.size(1)
         if num_blocks <= max_blocks:

vllm_ascend/torchair/torchair_attention.py

@@ -117,8 +117,8 @@ def _get_graph_runner_block_tables(
         max_blocks = self.max_blocks
 
         graph_block_tables = torch.zeros((num_seqs, max_blocks),
-                                             dtype=block_tables.dtype,
-                                             device=block_tables.device)
+                                         dtype=block_tables.dtype,
+                                         device=block_tables.device)
 
         num_blocks = block_tables.size(1)
         if num_blocks <= max_blocks: