fixup

Signed-off-by: CodeNine-CJ <[email protected]>

Author: ForBetterCodeNine

tests/ut/torchair/test_torchair_model_runner.py

@@ -3,7 +3,7 @@
 import pytest
 import torch
 from pytest_mock import MockerFixture
-from vllm.config import CacheConfig, VllmConfig
+from vllm.config import VllmConfig
 
 from tests.ut.base import PytestBase
 from vllm_ascend.torchair.torchair_model_runner import NPUTorchairModelRunner
@@ -66,16 +66,18 @@ class TestNPUTorchairModelRunner(PytestBase):
     @pytest.fixture
     def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
         vllm_config = MagicMock(spec=VllmConfig)
+        vllm_config.device_config = MagicMock()
+        vllm_config.device_config.device = torch.device("cpu")
         vllm_config.model_config = MagicMock()
-        cache_config = CacheConfig(block_size=16)
-        vllm_config.cache_config = cache_config
         vllm_config.model_config.hf_config = MagicMock()
         vllm_config.model_config.hf_config.index_topk = 2
         vllm_config.model_config.max_model_len = 1024
         vllm_config.model_config.use_mla = False
         vllm_config.model_config.get_hidden_size.return_value = 512
         vllm_config.model_config.pooler_config = None
         vllm_config.model_config.logits_processors = []
+        vllm_config.model_config.dtype = torch.float16
+
         cache_config = MagicMock()
         cache_config.block_size = 16
         cache_config.cache_dtype = "auto"
@@ -85,25 +87,27 @@ def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
         speculative_config.num_speculative_tokens = 4
         vllm_config.speculative_config = speculative_config
 
-        vllm_config.compilation_config = MagicMock()
-        vllm_config.compilation_config.cudagraph_mode = Mock()
-        vllm_config.compilation_config.cudagraph_capture_sizes = [1, 2, 4]
+        compilation_config = MagicMock()
+        compilation_config.cudagraph_mode = Mock()
+        compilation_config.cudagraph_capture_sizes = [1, 2, 4]
+        vllm_config.compilation_config = compilation_config
 
-        vllm_config.lora_config = MagicMock()
-        vllm_config.parallel_config = MagicMock()
-        vllm_config.parallel_config.data_parallel_size = 1
-        vllm_config.parallel_config.data_parallel_rank = 0
-        vllm_config.parallel_config.cp_kv_cache_interleave_size = 1
+        parallel_config = MagicMock()
+        parallel_config.data_parallel_size = 1
+        parallel_config.data_parallel_rank = 0
+        parallel_config.cp_kv_cache_interleave_size = 1
+        vllm_config.parallel_config = parallel_config
 
         scheduler_config = MagicMock()
         scheduler_config.max_num_batched_tokens = 2048
         scheduler_config.max_num_seqs = 64
         scheduler_config.chunked_prefill_enabled = True
         scheduler_config.async_scheduling = False
+        scheduler_config.decode_max_num_seqs = 0
         vllm_config.scheduler_config = scheduler_config
 
         vllm_config.load_config = MagicMock()
-
+        vllm_config.lora_config = MagicMock()
         vllm_config.kv_transfer_config = None
 
         mocker.patch(
@@ -116,16 +120,14 @@ def setup_npu_torchair_model_runner(self, mocker: MockerFixture):
             return_value=False)
         mocker.patch("vllm_ascend.worker.model_runner_v1.get_dcp_group"
                      ).return_value.world_size = 1
-        mocker.patch(
-            "vllm_ascend.torchair.torchair_model_runner.get_attn_backend",
-            autospec=True)
-        mocker.patch(
-            "vllm_ascend.torchair.torchair_model_runner._set_up_drafter")
-        mocker.patch(
-            "vllm_ascend.torchair.torchair_model_runner._may_pad_kv_consumer_num_seq"
-        )
-
-        device = torch.device("npu:0")
+        mocker.patch("vllm.attention.get_attn_backend", autospec=True)
+        # mocker.patch(
+        #     "vllm_ascend.torchair.torchair_model_runner._set_up_drafter")
+        # mocker.patch(
+        #     "vllm_ascend.torchair.torchair_model_runner._may_pad_kv_consumer_num_seq"
+        # )
+
+        device = torch.device("npu")
         ascend_config = MagicMock()
 
         ascend_config.ascend_scheduler_config.enabled = False