Merge branch 'release/2.1' into Jiang-Jia-Jun-patch-2

Author: Jiang-Jia-Jun

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -64,7 +64,10 @@ def __init__(
         self.speculative_config = config.speculative_config
         self.local_data_parallel_id = local_data_parallel_id
 
-        self.num_gpu_blocks = self.cache_config.prefill_kvcache_block_num
+        if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+            self.num_gpu_blocks = self.cache_config.total_block_num
+        else:
+            self.num_gpu_blocks = self.cache_config.prefill_kvcache_block_num
         self.num_cpu_blocks = self.cache_config.num_cpu_blocks
         self.gpu_free_block_list = list(range(self.num_gpu_blocks - 1, -1, -1))
         if self.num_cpu_blocks > 0:

fastdeploy/config.py

@@ -726,7 +726,10 @@ def __init__(self, args):
         self.block_size = 64
         self.gpu_memory_utilization = 0.9
         self.num_gpu_blocks_override = None
-        self.kv_cache_ratio = 0.75
+        if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+            self.kv_cache_ratio = 1.0
+        else:
+            self.kv_cache_ratio = 0.75
         self.enc_dec_block_num = 2
         self.prealloc_dec_block_slot_num_threshold = 5
         self.cache_dtype = "bfloat16"
@@ -811,7 +814,10 @@ def postprocess(self, num_total_tokens, number_of_tasks):
         self.dec_token_num = self.enc_dec_block_num * self.block_size
         if self.num_gpu_blocks_override is not None:
             self.total_block_num = self.num_gpu_blocks_override
-            self.prefill_kvcache_block_num = int(self.total_block_num * self.kv_cache_ratio)
+            if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+                self.prefill_kvcache_block_num = self.total_block_num
+            else:
+                self.prefill_kvcache_block_num = int(self.total_block_num * self.kv_cache_ratio)
         else:
             length = num_total_tokens // number_of_tasks
             block_num = (length + self.block_size - 1 + self.dec_token_num) // self.block_size
@@ -824,7 +830,10 @@ def reset(self, num_gpu_blocks):
         reset gpu block number
         """
         self.total_block_num = num_gpu_blocks
-        self.prefill_kvcache_block_num = int(self.total_block_num * self.kv_cache_ratio)
+        if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+            self.prefill_kvcache_block_num = self.total_block_num
+        else:
+            self.prefill_kvcache_block_num = int(self.total_block_num * self.kv_cache_ratio)
         logger.info(
             f"Reset block num, the total_block_num:{self.total_block_num},"
             f" prefill_kvcache_block_num:{self.prefill_kvcache_block_num}"

fastdeploy/engine/args_utils.py

@@ -18,6 +18,7 @@
 from dataclasses import asdict, dataclass
 from dataclasses import fields as dataclass_fields
 from typing import Any, Dict, List, Optional
+import os
 
 from fastdeploy.config import (
     CacheConfig,
@@ -865,7 +866,10 @@ def create_engine_config(self) -> Config:
             if self.enable_chunked_prefill:
                 self.max_num_batched_tokens = 2048
             else:
-                self.max_num_batched_tokens = self.max_model_len
+                if not int(os.getenv('ENABLE_V1_KVCACHE_SCHEDULER', '0')):
+                    self.max_num_batched_tokens = self.max_model_len
+                else:
+                    self.max_num_batched_tokens = 8192
 
         all_dict = asdict(self)
         all_dict["model_cfg"] = model_cfg

fastdeploy/engine/config.py

@@ -236,7 +236,10 @@ def postprocess(self):
             if self.cache_config.enable_chunked_prefill:
                 self.max_num_batched_tokens = 2048
             else:
-                self.max_num_batched_tokens = self.max_model_len
+                if not int(os.getenv('ENABLE_V1_KVCACHE_SCHEDULER', '0')):
+                    self.max_num_batched_tokens = self.max_model_len
+                else:
+                    self.max_num_batched_tokens = 8192
 
         if self.long_prefill_token_threshold == 0:
             self.long_prefill_token_threshold = int(self.max_model_len * 0.04)
@@ -284,10 +287,11 @@ def check(self):
         )
 
         if not self.cache_config.enable_chunked_prefill:
-            assert self.max_num_batched_tokens >= self.max_model_len, (
-                f"max_num_batched_tokens: {self.max_num_batched_tokens} "
-                f"should be larger than or equal to max_model_len: {self.max_model_len}"
-            )
+            if not int(os.getenv('ENABLE_V1_KVCACHE_SCHEDULER', '0')):
+                assert self.max_num_batched_tokens >= self.max_model_len, (
+                    f"max_num_batched_tokens: {self.max_num_batched_tokens} "
+                    f"should be larger than or equal to max_model_len: {self.max_model_len}"
+                )
         else:
             assert self.max_num_batched_tokens >= self.cache_config.block_size, (
                 f"max_num_batched_tokens: {self.max_num_batched_tokens} "

fastdeploy/engine/sched/resource_manager_v1.py

@@ -195,7 +195,6 @@ def _get_num_new_tokens(self, request, token_budget):
                 )
                 request.num_image_end = img_num_per_boundary[new_boundary_idx]
 
-            request.num_image_end = img_num_per_boundary[new_boundary_idx]
             request.image_type_ids_start = np.sum(grid_thw[: request.num_image_start, 0])
             request.image_type_ids_end = np.sum(grid_thw[: request.num_image_end, 0])
             request.image_start = np.sum(np.prod(grid_thw[: request.num_image_start], axis=1))

fastdeploy/output/token_processor.py

@@ -463,8 +463,9 @@ def _process_batch_output(self):
                 if recovery_stop:
                     llm_logger.info(f"recovery stop signal found at task {task_id}")
                 if not recovery_stop and token_id < 0:
-                    if task_id in self.resource_manager.to_be_rescheduled_request_id_set:
-                        self.resource_manager.reschedule_preempt_task(task_id)
+                    if envs.ENABLE_V1_KVCACHE_SCHEDULER:
+                        if task_id in self.resource_manager.to_be_rescheduled_request_id_set:
+                            self.resource_manager.reschedule_preempt_task(task_id)
                     continue
 
             if task.get("prefill_chunk_info", None) is not None:

fastdeploy/worker/worker_process.py

@@ -24,6 +24,7 @@
 import paddle.distributed as dist
 from paddle.distributed import fleet
 
+from fastdeploy import envs
 from fastdeploy.config import (
     CacheConfig,
     DecodingConfig,
@@ -289,8 +290,9 @@ def event_loop_normal(self) -> None:
             if self.local_rank % mp_num_per_node == 0:
                 if self.task_queue.num_tasks() > 0:
                     # VL only support 1 batch to prefill
-
-                    if not self.fd_config.model_config.enable_mm or not self.worker.exist_prefill():
+                    if envs.ENABLE_V1_KVCACHE_SCHEDULER or not (
+                        self.fd_config.model_config.enable_mm and self.worker.exist_prefill()
+                    ):
                         if self.nnode > 1 and self.parallel_config.tensor_parallel_size > self.max_chips_per_node:
                             self.task_queue.read_finish_flag.set(1)
                         else: