fix ut

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

Author: Potabk

tests/ut/attention/test_mla_v1.py

@@ -1,7 +1,7 @@
 from unittest.mock import MagicMock, patch
 
 import torch
-from vllm.config import VllmConfig
+from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig
 from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.model_executor.layers.linear import LinearBase
 
@@ -184,19 +184,15 @@ class TestAscendMLAMetadataBuilder(TestBase):
            return_value=1)
     def test_ascend_mla_metadata_builder_default(self, mock_get_dcp_size,
                                                  mock_dcp, mock_get_dcp_group):
-        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')
+        mock_vllm_config.model_config.max_model_len = 1024
+        mock_vllm_config.model_config.get_head_size.return_value = 64
+        mock_vllm_config.model_config.dtype = torch.float16
+        mock_vllm_config.cache_config.block_size = 16
+        mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
+        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
+        mock_device = 'cpu'
 
         mock_dcp.world_size = 1
         dcp_group = MagicMock(spec=GroupCoordinator)
@@ -205,6 +201,8 @@ def test_ascend_mla_metadata_builder_default(self, mock_get_dcp_size,
         dcp_group.device_group = MagicMock()
         mock_get_dcp_group.return_value = dcp_group
 
+        mock_vllm_config.speculative_config = None
+
         ascend_config = MagicMock()
         with patch("vllm_ascend.attention.mla_v1.get_ascend_config",
                    return_value=ascend_config):
@@ -225,19 +223,15 @@ def test_ascend_mla_metadata_builder_default(self, mock_get_dcp_size,
     def test_ascend_mla_metadata_builder_spec_decode(self, mock_get_dcp_size,
                                                      mock_dcp,
                                                      mock_get_dcp_group):
-        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')
+        mock_vllm_config.model_config.max_model_len = 1024
+        mock_vllm_config.model_config.get_head_size.return_value = 64
+        mock_vllm_config.model_config.dtype = torch.float16
+        mock_vllm_config.cache_config.block_size = 16
+        mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
+        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
+        mock_device = 'cpu'
 
         mock_dcp.world_size = 1
         dcp_group = MagicMock(spec=GroupCoordinator)
@@ -260,7 +254,7 @@ def test_ascend_mla_metadata_builder_spec_decode(self, mock_get_dcp_size,
                              mock_vllm_config.cache_config.block_size)
             self.assertEqual(
                 builder.chunked_prefill_enabled,
-                mock_vllm_config.scheduler_config.chunked_prefill_enabled)
+                mock_vllm_config.scheduler_config.enable_chunked_prefill)
 
     @patch('vllm.distributed.parallel_state.get_dcp_group')
     @patch('vllm.distributed.parallel_state._DCP',
@@ -322,19 +316,13 @@ def test_reorder_batch(self, mock_get_dcp_size, mock_dcp,
                            mock_get_dcp_group):
         ascend_config = MagicMock()
 
-        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')
+        mock_vllm_config.model_config.max_model_len = 1024
+        mock_vllm_config.cache_config.block_size = 16
+        mock_vllm_config.scheduler_config.max_num_seqs = 4
+        mock_vllm_config.scheduler_config.decode_max_num_seqs = 4
+        mock_vllm_config.scheduler_config.enable_chunked_prefill = False
+        mock_device = 'cpu'
 
         mock_dcp.world_size = 1
         dcp_group = MagicMock(spec=GroupCoordinator)
@@ -343,6 +331,8 @@ def test_reorder_batch(self, mock_get_dcp_size, mock_dcp,
         dcp_group.device_group = MagicMock()
         mock_get_dcp_group.return_value = dcp_group
 
+        mock_vllm_config.speculative_config = None
+
         with patch("vllm_ascend.attention.mla_v1.get_ascend_config",
                    return_value=ascend_config):
             builder = AscendMLAMetadataBuilder(None, None, mock_vllm_config,
@@ -447,21 +437,15 @@ class TestAscendMLAMetadataBuilderBuild(TestBase):
 
     def setUp(self):
         self.mock_vllm_config = MagicMock(spec=VllmConfig)
-        # NOTE: Do not init the ModelConfig from constructor
-        # Which will try to download a model
-        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
-
-        from vllm.config.scheduler import SchedulerConfig
-        self.mock_vllm_config.scheduler_config = SchedulerConfig()
-
-        self.mock_vllm_config.model_config = mock_model_config
-        self.mock_vllm_config.cache_config = MagicMock(block_size=16)
+        self.mock_vllm_config.model_config = ModelConfig(max_model_len=2048)
+        self.mock_vllm_config.model_config.hf_text_config.qk_rope_head_dim = 32
+        self.mock_vllm_config.cache_config = CacheConfig(block_size=32)
+        mock_scheduler_config = MagicMock(spec=SchedulerConfig)
+        mock_scheduler_config.max_num_seqs = 8  # 设置为整数，不是 MagicMock
+        mock_scheduler_config.chunked_prefill_enabled = True
+        self.mock_vllm_config.scheduler_config = mock_scheduler_config
         self.mock_vllm_config.speculative_config = None
-
-        self.mock_device = torch.device('cpu')
+        self.mock_device = torch.device("cpu")
 
         self.kv_cache_spec = MagicMock()
         self.kv_cache_spec.num_layers = 32