fix mla

Signed-off-by: wangli <wangli858794774@gmail.com>

Author: Potabk

tests/ut/torchair/test_torchair_mla.py

@@ -253,16 +253,20 @@ def test_reorder_batch_without_torchair_graph(self):
         ascend_config.torchair_graph_config = MagicMock()
         ascend_config.torchair_graph_config.enabled = False
 
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 1024
-        mock_vllm_config.get_head_size = lambda: 8
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.scheduler_config.max_num_seqs = 4
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_device = torch.device('cpu')
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
 
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
         mock_vllm_config.speculative_config = None
 
+        mock_device = torch.device('cpu')
+
         with patch("vllm_ascend.torchair.torchair_mla.get_ascend_config",
                    return_value=ascend_config):
             builder = AscendMLATorchairMetadataBuilder(None, None,
@@ -293,14 +297,21 @@ def test_get_graph_runner_block_tables_normal(self, mock_ascend_config):
         ascend_config = MagicMock()
         mock_ascend_config.return_value = ascend_config
         ascend_config.torchair_graph_config.enabled = False
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 1024
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_device = torch.device('cpu')
 
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
+
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
         mock_vllm_config.speculative_config = None
 
+        mock_device = torch.device('cpu')
+
         builder = AscendMLATorchairMetadataBuilder(None, None,
                                                    mock_vllm_config,
                                                    mock_device)
@@ -316,14 +327,21 @@ def test_get_graph_runner_block_tables_truncated(self, mock_ascend_config):
         ascend_config = MagicMock()
         mock_ascend_config.return_value = ascend_config
         ascend_config.torchair_graph_config.enabled = False
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 64
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_device = torch.device('cpu')
 
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
+
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
         mock_vllm_config.speculative_config = None
 
+        mock_device = torch.device('cpu')
+
         builder = AscendMLATorchairMetadataBuilder(None, None,
                                                    mock_vllm_config,
                                                    mock_device)
@@ -340,16 +358,21 @@ def test_get_graph_runner_block_tables_from_numpy(self,
         ascend_config = MagicMock()
         mock_ascend_config.return_value = ascend_config
         ascend_config.torchair_graph_config.enabled = False
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 1024
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.get_head_size = lambda: 28
-        mock_vllm_config.dtype = torch.bfloat16
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_device = torch.device('cpu')
 
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
+
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
         mock_vllm_config.speculative_config = None
 
+        mock_device = torch.device('cpu')
+
         builder = AscendMLATorchairMetadataBuilder(None, None,
                                                    mock_vllm_config,
                                                    mock_device)
@@ -368,16 +391,20 @@ def test_build_dummy(self, mock_ascend_config):
         mock_ascend_config.return_value = ascend_config
         ascend_config.torchair_graph_config.enabled = False
 
-        mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 1024
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_vllm_config.get_head_size.return_value = 64
-        mock_vllm_config.model_config.dtype = torch.float16
-        mock_device = torch.device('cpu')
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
 
+        mock_vllm_config = MagicMock()
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
         mock_vllm_config.speculative_config = None
 
+        mock_device = torch.device('cpu')
+
         builder = AscendMLATorchairMetadataBuilder(
             None,
             None,
@@ -435,18 +462,23 @@ def test_build_decode(self, mock_ascend_config):
         mock_ascend_config.return_value = ascend_config
         ascend_config.torchair_graph_config.enabled = False
 
+        mock_model_config = MagicMock()
+        mock_model_config.max_model_len = 1024
+        mock_model_config.get_head_size.return_value = 64
+        mock_model_config.dtype = torch.float16
+
         mock_vllm_config = MagicMock()
-        mock_vllm_config.model_config.max_model_len = 1024
-        mock_vllm_config.cache_config.block_size = 16
-        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
-        mock_vllm_config.get_head_size.return_value = 64
-        mock_vllm_config.model_config.dtype = torch.float16
+        mock_vllm_config.model_config = mock_model_config
+        mock_vllm_config.cache_config = MagicMock(block_size=16)
+        mock_vllm_config.scheduler_config = MagicMock(
+            max_num_seqs=4, enable_chunked_prefill=False)
+        mock_vllm_config.speculative_config = None
+
         mock_device = torch.device('cpu')
+
         model = MagicMock(spec=nn.Module)
         model.model = MagicMock(spec=nn.Module)
 
-        mock_vllm_config.speculative_config = None
-
         builder = AscendMLATorchairMetadataBuilder(
             None,
             None,