|
4 | 4 | import pytest |
5 | 5 | import torch |
6 | 6 | from vllm.config.compilation import CUDAGraphMode |
| 7 | +from vllm.engine.arg_utils import EngineArgs |
7 | 8 | from vllm.platforms import PlatformEnum |
8 | 9 |
|
9 | 10 | from tests.ut.base import TestBase |
@@ -722,3 +723,32 @@ def test_get_static_graph_wrapper_cls_returns_correct_value(self): |
722 | 723 | self.platform.get_static_graph_wrapper_cls(), |
723 | 724 | "vllm_ascend.compilation.acl_graph.ACLGraphWrapper", |
724 | 725 | ) |
| 726 | + |
| 727 | + def test_aclgraph_enable(self): |
| 728 | + config = EngineArgs() |
| 729 | + VllmConfig = config.create_engine_config() |
| 730 | + self.assertEqual(VllmConfig.compilation_config.cudagraph_mode, |
| 731 | + CUDAGraphMode.PIECEWISE) |
| 732 | + |
| 733 | + with self.assertLogs(logger="vllm", level="INFO") as cm: |
| 734 | + from vllm_ascend import platform |
| 735 | + |
| 736 | + importlib.reload(platform) |
| 737 | + self.platform.check_and_update_config(VllmConfig) |
| 738 | + self.assertTrue( |
| 739 | + "PIECEWISE compilation enabled on NPU. use_inductor not supported - " |
| 740 | + "using only ACL Graph mode" in cm.output[1]) |
| 741 | + if vllm_version_is("0.11.0"): |
| 742 | + self.assertEqual( |
| 743 | + VllmConfig.compilation_config.level, |
| 744 | + CompilationLevel.PIECEWISE, |
| 745 | + ) |
| 746 | + else: |
| 747 | + self.assertEqual( |
| 748 | + VllmConfig.compilation_config.mode, |
| 749 | + CompilationMode.VLLM_COMPILE, |
| 750 | + ) |
| 751 | + self.assertEqual( |
| 752 | + VllmConfig.compilation_config.cudagraph_mode, |
| 753 | + CUDAGraphMode.PIECEWISE, |
| 754 | + ) |
0 commit comments