[main] adjust the position of warm_up_atb (#2823)

### What this PR does / why we need it?
Adjust the position of warm_up_atb.

### Does this PR introduce _any_ user-facing change?
not involved

### How was this patch tested?
CI passed with existing test.

- vLLM version: main
- vLLM main:
b23fb78623

Signed-off-by: huangxialu <huangxialu1@huawei.com>
This commit is contained in:
huangxialu
2025-09-10 14:06:38 +08:00
committed by GitHub
parent 22b425765a
commit 88d7af62be
3 changed files with 21 additions and 8 deletions

View File

@@ -1009,7 +1009,10 @@ class TestNPUWorker(TestBase):
@patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
@patch("vllm_ascend.worker.worker_v1.logger")
def test_compile_or_warm_up_model_with_eager_mode(self, mock_logger,
@patch("torch_npu._npu_matmul_add_fp32")
def test_compile_or_warm_up_model_with_eager_mode(self,
mock_npu_matmul_add,
mock_logger,
mock_seed_everything):
"""Test compile_or_warm_up_model method - eager mode"""
from vllm_ascend.worker.worker_v1 import NPUWorker
@@ -1051,10 +1054,14 @@ class TestNPUWorker(TestBase):
# Verify seed setting
mock_seed_everything.assert_called_once_with(12345)
# Verify calls
mock_npu_matmul_add.assert_called_once()
@patch("vllm_ascend.worker.worker_v1.NPUPlatform.seed_everything")
@patch("vllm_ascend.worker.worker_v1.logger")
@patch("torch_npu._npu_matmul_add_fp32")
def test_compile_or_warm_up_model_with_graph_capture(
self, mock_logger, mock_seed_everything):
self, mock_npu_matmul_add, mock_logger, mock_seed_everything):
"""Test compile_or_warm_up_model method - with graph capture enabled"""
from vllm_ascend.worker.worker_v1 import NPUWorker
@@ -1087,6 +1094,9 @@ class TestNPUWorker(TestBase):
# Verify seed setting
mock_seed_everything.assert_called_once_with(67890)
# Verify calls
mock_npu_matmul_add.assert_called_once()
@patch("vllm_ascend.worker.worker_v1.CaMemAllocator")
def test_initialize_from_config_with_sleep_mode(self,
mock_allocator_class):