diff --git a/tests/ut/test_platform.py b/tests/ut/test_platform.py index d5d43098..9e0c0b32 100644 --- a/tests/ut/test_platform.py +++ b/tests/ut/test_platform.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch import pytest import torch from vllm.config.compilation import CUDAGraphMode +from vllm.engine.arg_utils import EngineArgs from vllm.platforms import PlatformEnum from tests.ut.base import TestBase @@ -722,3 +723,32 @@ class TestNPUPlatform(TestBase): self.platform.get_static_graph_wrapper_cls(), "vllm_ascend.compilation.acl_graph.ACLGraphWrapper", ) + + def test_aclgraph_enable(self): + config = EngineArgs() + VllmConfig = config.create_engine_config() + self.assertEqual(VllmConfig.compilation_config.cudagraph_mode, + CUDAGraphMode.PIECEWISE) + + with self.assertLogs(logger="vllm", level="INFO") as cm: + from vllm_ascend import platform + + importlib.reload(platform) + self.platform.check_and_update_config(VllmConfig) + self.assertTrue( + "PIECEWISE compilation enabled on NPU. use_inductor not supported - " + "using only ACL Graph mode" in cm.output[1]) + if vllm_version_is("0.11.0"): + self.assertEqual( + VllmConfig.compilation_config.level, + CompilationLevel.PIECEWISE, + ) + else: + self.assertEqual( + VllmConfig.compilation_config.mode, + CompilationMode.VLLM_COMPILE, + ) + self.assertEqual( + VllmConfig.compilation_config.cudagraph_mode, + CUDAGraphMode.PIECEWISE, + )