feat: add watermark schema

Signed-off-by: hw_whx <wanghexiang7@huawei.com>

Author: hw_whx

vllm_ascend/core/schedule_config.py

@@ -34,7 +34,7 @@ def __post_init__(self) -> None:
             raise NotImplementedError(f"currently AscendScheduler only supports fcfs policy, got {self.policy}")
         if self.is_multimodal_model:
             raise NotImplementedError(f"currently AscendScheduler only supports LLM modles.")
-        if self.num_scheduler_steps >t1:
+        if self.num_scheduler_steps > 1:
             raise NotImplementedError(f"currently AscendScheduler doesn't support multi-step.")
         if self.send_delta_data:
             raise NotImplementedError(f"currently AscendScheduler doesn't support send_delta_data.")

vllm_ascend/core/scheduler.py

@@ -1,8 +1,7 @@
 from collections import deque
-from dataclasses import dataclass
-from typing import List
 
 from vllm.logger import init_logger
+from vllm.utils import cdiv
 from vllm.v1.core.scheduler import Scheduler
 from vllm.v1.core.scheduler_output import (NewRequestData,
                                        SchedulerOutput)
@@ -88,6 +87,11 @@ def skip_cur_request():
                 skip_cur_request()
                 continue
 
+            if not self._check_watermark_for_prefill(num_new_tokens):
+                # Scheduling would exceed watermark, skip.
+                skip_cur_request()
+                continue
+
             assert num_new_tokens > 0
             new_blocks = self.kv_cache_manager.allocate_slots(
                     request, num_new_tokens, computed_blocks)
@@ -261,6 +265,14 @@ def skip_cur_request():
 
         self.finished_req_ids = set()
         return scheduler_output
+    
+
+    def _check_watermark_for_prefill(self, num_new_tokens, watermark = 0.01):
+        watermark_blocks = self.cache_config.num_gpu_blocks * watermark
+        num_required_blocks = cdiv(num_new_tokens, self.block_size)
+        if (self.kv_cache_manager.free_block_queue.num_free_blocks - num_required_blocks) < watermark_blocks:
+            return False
+        return True
 
 
     def _get_prompt_limit(self, request: Request) -> int:

vllm_ascend/worker/worker_v1.py

@@ -41,6 +41,7 @@
 
 from vllm_ascend.device_allocator.camem import CaMemAllocator
 from vllm_ascend.worker.model_runner_v1 import NPUModelRunner
+from vllm_ascend.utils import try_register_lib
 
 logger = init_logger(__name__)
 
@@ -74,6 +75,12 @@ def __init__(self,
         self.prompt_adapter_config = vllm_config.prompt_adapter_config
         self.observability_config = vllm_config.observability_config
 
+        # Try to import mindie_turbo to accelerate vLLM inference.
+        try_register_lib(
+            "mindie_turbo",
+            "MindIE Turbo is installed. vLLM inference will be accelerated with MindIE Turbo."
+        )
+
         if self.cache_config.cache_dtype == "auto":
             self.cache_dtype = self.model_config.dtype
         else: