fix kv cache blcok shape

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

Author: wangxiyuan

tests/ut/attention/test_attention_v1.py

@@ -115,7 +115,7 @@ def test_build_prefill_no_cache(self, mock_is_310p, mock_nd_to_nz_2d,
         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(1, common_attn_metadata, mock_model)
 
     @patch('vllm_ascend.attention.attention_v1.AscendMetadata')
     @patch('torch_npu.npu_format_cast')
@@ -151,7 +151,7 @@ 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(common_attn_metadata, mock_model)
+        self.builder.build(1, common_attn_metadata, mock_model)
 
     @patch('vllm_ascend.attention.attention_v1.AscendMetadata')
     @patch('vllm_ascend.attention.attention_v1.is_310p', return_value=False)
@@ -175,7 +175,7 @@ def test_build_non_310p(self, mock_is_310p, mock_ascend_metadata):
             seq_lens=None)
         mock_model = MagicMock()
 
-        self.builder.build(common_attn_metadata, mock_model)
+        self.builder.build(1, common_attn_metadata, mock_model)
 
 
 class TestAscendAttentionBackendImpl(TestBase):

tests/ut/torchair/test_torchair_mla.py

@@ -456,7 +456,7 @@ def test_build_decode(self, mock_ascend_config):
                 num_computed_tokens_cpu=None,
                 seq_lens=None)
 
-            metadata = builder.build(common_attn_metadata, model)
+            metadata = builder.build(1, common_attn_metadata, model)
 
         self.assertIsInstance(metadata, AscendMLATorchairMetadata)
         self.assertEqual(metadata.num_input_tokens, 0)

vllm_ascend/worker/model_runner_v1.py

@@ -2459,16 +2459,17 @@ def align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor:
                             kv_cache_spec.num_kv_heads,
                             kv_cache_spec.head_size)
                     elif hasattr(attn_backend, "get_supported_block_size"):
-                        # kv_cache_shape = attn_backend.get_kv_cache_shape(
-                        #     num_blocks, kv_cache_spec.block_size,
-                        #     kv_cache_spec.num_kv_heads,
-                        #     kv_cache_spec.head_size)
                         block_size = attn_backend.get_supported_block_size()[0]
                         block_size_chunk = kv_cache_spec.block_size // block_size
                         kv_cache_shape = attn_backend.get_kv_cache_shape(
                             num_blocks * block_size_chunk, block_size,
                             kv_cache_spec.num_kv_heads,
                             kv_cache_spec.head_size)
+                    else:
+                        kv_cache_shape = self.attn_backend.get_kv_cache_shape(
+                            num_blocks, kv_cache_spec.block_size,
+                            kv_cache_spec.num_kv_heads,
+                            kv_cache_spec.head_size)
                     dtype = kv_cache_spec.dtype
                     if self.model_config.is_deepseek_mla:
                         num_blocks, block_size, num_kv_heads, head_size = kv_cache_shape