Revert "drop ascend scheduler" (#4580)

Reverts vllm-project/vllm-ascend#4498
- vLLM version: v0.11.2
- vLLM main: https://github.com/vllm-project/vllm/commit/v0.11.2

tests/ut/core/test_schedule_config.py

@@ -0,0 +1,134 @@
+#
+# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from vllm.config import SchedulerConfig
+
+from tests.ut.base import TestBase
+from vllm_ascend.core.schedule_config import AscendSchedulerConfig
+
+
+class TestAscendSchedulerConfig(TestBase):
+
+    def setUp(self):
+        self.basic_scheduler_config = SchedulerConfig(
+            max_num_batched_tokens=8192,
+            max_model_len=8192,
+            is_multimodal_model=False,
+            send_delta_data=False,
+        )
+
+    def test_initialize_from_config_with_default(self):
+        # No additional config given, check the default value here.
+        ascend_config = AscendSchedulerConfig.initialize_from_config(
+            self.basic_scheduler_config, {})
+        self.assertEqual(ascend_config.enable_chunked_prefill, False)
+        self.assertEqual(ascend_config.policy, "fcfs")
+        self.assertEqual(ascend_config.scheduler_cls,
+                         "vllm_ascend.core.scheduler.AscendScheduler")
+        self.assertEqual(ascend_config.max_num_encoder_input_tokens, 8192)
+        self.assertEqual(ascend_config.encoder_cache_size, 8192)
+
+    def test_initialize_from_config_with_override(self):
+        # test override
+        ascend_config = AscendSchedulerConfig.initialize_from_config(
+            self.basic_scheduler_config,
+            AscendSchedulerConfig(
+                enable_chunked_prefill=False,
+                policy="fcfs",
+                scheduler_cls="vllm_ascend.core.scheduler.AscendScheduler",
+                max_num_batched_tokens=8192,
+                max_model_len=2048,
+                max_long_partial_prefills=1,
+                long_prefill_token_threshold=512,
+            ),
+        )
+        self.assertEqual(ascend_config.enable_chunked_prefill, False)
+        self.assertEqual(ascend_config.policy, "fcfs")
+        self.assertEqual(ascend_config.scheduler_cls,
+                         "vllm_ascend.core.scheduler.AscendScheduler")
+        self.assertEqual(ascend_config.max_num_batched_tokens, 8192)
+        self.assertEqual(ascend_config.encoder_cache_size, 8192)
+        self.assertEqual(ascend_config.max_long_partial_prefills, 1)
+        self.assertEqual(ascend_config.long_prefill_token_threshold, 512)
+
+    def test_not_implemented_policy(self):
+        with self.assertRaises(NotImplementedError) as context:
+            AscendSchedulerConfig.initialize_from_config(
+                self.basic_scheduler_config,
+                AscendSchedulerConfig(
+                    policy="custom_policy",
+                    max_num_batched_tokens=8192,
+                    max_model_len=2048,
+                ),
+            )
+        self.assertIn(
+            "currently AscendScheduler only supports fcfs policy",
+            str(context.exception),
+        )
+
+    def test_no_override(self):
+        ascend_config = AscendSchedulerConfig.initialize_from_config(
+            self.basic_scheduler_config, {})
+        self.assertEqual(ascend_config.max_num_encoder_input_tokens, 8192)
+        self.assertEqual(ascend_config.encoder_cache_size, 8192)
+
+    def test_valid_config_with_multimodal(self):
+        config = AscendSchedulerConfig.initialize_from_config(
+            SchedulerConfig(is_multimodal_model=True,
+                            max_num_batched_tokens=8192), {})
+        self.assertTrue(config.is_multimodal_model)
+
+    def test_valid_config_with_chunked_prefill(self):
+        ascend_config = AscendSchedulerConfig.initialize_from_config(
+            self.basic_scheduler_config,
+            AscendSchedulerConfig(
+                enable_chunked_prefill=True,
+                max_num_batched_tokens=8192,
+                max_model_len=8192,
+            ),
+        )
+        self.assertEqual(ascend_config.max_num_batched_tokens, 8192)
+        self.assertEqual(ascend_config.max_model_len, 8192)
+        self.assertTrue(ascend_config.enable_chunked_prefill)
+
+    def test_invalid_config_without_chunked_prefill(self):
+        with self.assertRaises(ValueError) as context:
+            AscendSchedulerConfig.initialize_from_config(
+                self.basic_scheduler_config,
+                AscendSchedulerConfig(
+                    enable_chunked_prefill=False,
+                    max_num_batched_tokens=2048,
+                    max_model_len=8192,
+                ),
+            )
+        self.assertIn(
+            "Ascend scheduler is enabled without chunked prefill feature",
+            str(context.exception),
+        )
+        self.assertIn("max_num_batched_tokens (2048)", str(context.exception))
+        self.assertIn("max_model_len (8192)", str(context.exception))
+
+    def test_initialize_from_config_with_pd_transfer(self):
+        ascend_config = AscendSchedulerConfig.initialize_from_config(
+            self.basic_scheduler_config,
+            AscendSchedulerConfig(
+                enable_pd_transfer=True,
+                decode_max_num_seqs=48,
+                max_num_batched_tokens=8192,
+                max_model_len=4096,
+            ),
+        )
+        self.assertEqual(ascend_config.enable_pd_transfer, True)
+        self.assertEqual(ascend_config.decode_max_num_seqs, 48)

tests/ut/ops/test_linear.py

@@ -99,6 +99,7 @@ class TestAscendRowParallelLinear(BaseLinearTest):
 
         ascend_config._ASCEND_CONFIG = MagicMock()
         ascend_config._ASCEND_CONFIG.oproj_tensor_parallel_size = 2
+        ascend_config._ASCEND_CONFIG.ascend_scheduler_config.enabled = False
 
         linear = AscendRowParallelLinear(
             input_size=16,

tests/ut/ops/test_vocab_parallel_embedding.py

@@ -209,7 +209,12 @@ class TestAscendLogitsProcessor(unittest.TestCase):
                 return_value=torch.randn(1, self.vocab_size)),
             patch(
                 "vllm_ascend.ops.vocab_parallel_embedding.get_lmhead_tp_group.all_gather",
-                return_value=torch.randn(1, self.vocab_size))
+                return_value=torch.randn(1, self.vocab_size)),
+            patch(
+                "vllm_ascend.core.schedule_config.AscendSchedulerConfig.initialize_from_config",
+                return_value=MagicMock(max_num_batched_tokens=1000,
+                                       max_model_len=512,
+                                       enable_chunked_prefill=False))
         ]
 
         for p in self.patches:

tests/ut/quantization/test_w8a8_dynamic.py

@@ -33,6 +33,13 @@ class TestAscendW8A8FusedMoEMethod(TestBase):
             mock_get_ep_group.return_value = mock_ep_group
             mock_ascend_config = Mock()
 
+            # 创建一个具有具体属性的 Mock 对象来表示 ascend_scheduler_config
+            mock_ascend_scheduler_config = Mock()
+            mock_ascend_scheduler_config.enabled = False
+            mock_ascend_scheduler_config.max_num_batched_tokens = 1024
+            mock_ascend_scheduler_config.max_model_len = 2048
+            mock_ascend_config.ascend_scheduler_config = mock_ascend_scheduler_config
+
             mock_ascend_config.torchair_graph_config = Mock(enabled=False)
             mock_ascend_config.enable_chunked_prefill = False
             mock_get_ascend_config.return_value = mock_ascend_config

tests/ut/test_ascend_config.py

@@ -56,6 +56,9 @@ class TestAscendConfig(TestBase):
         self.assertTrue(torchair_graph_config.enable_frozen_parameter)
         self.assertFalse(torchair_graph_config.enable_kv_nz)
 
+        ascend_scheduler_config = ascend_config.ascend_scheduler_config
+        self.assertFalse(ascend_scheduler_config.enabled)
+
     @_clean_up_ascend_config
     def test_init_ascend_config_with_additional_config(self):
         test_vllm_config = VllmConfig()
@@ -71,6 +74,9 @@ class TestAscendConfig(TestBase):
                 "enable_kv_nz": True
             },
             "multistream_overlap_shared_expert": True,
+            "ascend_scheduler_config": {
+                "enabled": True
+            },
             "expert_map_path": "test_expert_map_path",
             "refresh": True,
         }
@@ -88,6 +94,9 @@ class TestAscendConfig(TestBase):
         self.assertTrue(torchair_graph_config.enable_frozen_parameter)
         self.assertTrue(torchair_graph_config.enable_kv_nz)
 
+        ascend_scheduler_config = ascend_config.ascend_scheduler_config
+        self.assertTrue(ascend_scheduler_config.enabled)
+
     @_clean_up_ascend_config
     def test_init_ascend_config_with_refresh(self):
         test_vllm_config = VllmConfig()

tests/ut/test_platform.py

@@ -32,6 +32,7 @@ class TestNPUPlatform(TestBase):
     def mock_vllm_ascend_config():
         mock_ascend_config = MagicMock()
         mock_ascend_config.torchair_graph_config.enabled = False
+        mock_ascend_config.ascend_scheduler_config.enabled = False
         mock_ascend_config.enable_shared_expert_dp = False
         return mock_ascend_config
 
@@ -521,6 +522,31 @@ class TestNPUPlatform(TestBase):
         self.platform.check_and_update_config(vllm_config)
         self.assertEqual(vllm_config.compilation_config.custom_ops, [])
 
+    @patch('vllm_ascend.utils.get_ascend_device_type',
+           return_value=AscendDeviceType._910_93)
+    @patch("vllm_ascend.ascend_config.check_ascend_config")
+    @patch("vllm_ascend.ascend_config.init_ascend_config")
+    @patch(
+        "vllm_ascend.core.recompute_schedule_config.RecomputeSchedulerConfig.initialize_from_config"
+    )
+    def test_check_and_update_config_ascend_scheduler_config(
+            self, mock_init_recompute, mock_init_ascend, mock_check_ascend,
+            mock_soc_version):
+        mock_ascend_config = TestNPUPlatform.mock_vllm_ascend_config()
+        mock_ascend_config.ascend_scheduler_config.enabled = True
+        mock_init_ascend.return_value = mock_ascend_config
+        vllm_config = TestNPUPlatform.mock_vllm_config()
+        vllm_config.parallel_config.tensor_parallel_size = 1
+        mock_init_recompute.return_value = MagicMock()
+
+        with patch("vllm_ascend.core.schedule_config.AscendSchedulerConfig"
+                   ) as mock_scheduler:
+            from vllm_ascend import platform
+
+            importlib.reload(platform)
+            self.platform.check_and_update_config(vllm_config)
+            mock_scheduler.initialize_from_config.assert_called_once()
+
     @patch('vllm_ascend.platform.get_ascend_config')
     def test_get_attn_backend_cls_use_v1_and_mla(self, mock_get_ascend_config):
         mock_config = MagicMock()

tests/ut/test_utils.py

@@ -253,10 +253,12 @@ class TestUtils(TestBase):
         model_path = os.path.join(os.path.dirname(__file__), "fake_weight")
         test_model_config = ModelConfig(model=model_path, enforce_eager=True)
         test_parallel_config = ParallelConfig()
+        ascend_config = {"ascend_scheduler_config": {"enabled": False}}
         test_vllm_config = VllmConfig(
             model_config=test_model_config,
             compilation_config=test_compilation_config,
-            parallel_config=test_parallel_config)
+            parallel_config=test_parallel_config,
+            additional_config=ascend_config)
         utils.update_aclgraph_sizes(test_vllm_config)
         os.environ['HCCL_OP_EXPANSION_MODE'] = 'AIV'
         utils.update_aclgraph_sizes(test_vllm_config)

tests/ut/torchair/models/test_torchair_deepseek_v2.py

@@ -235,6 +235,8 @@ def test_torchair_deepseek_v2_mlp(mock_distributed, base_config):
                                 hidden_act="silu",
                                 quant_config=None)
     assert isinstance(mlp.act_fn, TorchairDeepseekV2SiluAndMul)
+    ascend_config = MagicMock()
+    ascend_config._ASCEND_CONFIG.ascend_scheduler_config.enabled = False
     with patch(
             "vllm_ascend.torchair.models.torchair_deepseek_v2.QuantizationConfig"
     ) as mock_quant_config: