Write the cache of preempted req to storage (#7113)

Author: juncaipeng

fastdeploy/cache_manager/cache_transfer_manager.py

@@ -796,7 +796,7 @@ def read_storage_task(self, task: ReadStorageTask):
                 try:
                     valid_gpu_block_ids = self._run_read_storage(
                         task.task_id,
-                        task.token_ids[: match_block_num * self.block_size],
+                        task.token_ids[: match_block_num * self.block_size] if task.token_ids else None,
                         task.start_read_block_idx,
                         k_cache_keys,
                         v_cache_keys,

fastdeploy/cache_manager/prefix_cache_manager.py

@@ -872,7 +872,7 @@ def request_match_blocks(self, task: Request, block_size, *args):
                     read_storage_task = ReadStorageTask(
                         task_id=req_id,
                         keys=no_match_block_keys,
-                        token_ids=input_token_ids,
+                        token_ids=input_token_ids if self.kvcache_storage_backend == "attention_store" else None,
                         gpu_block_ids=gpu_recv_storage_block_ids,
                         start_read_block_idx=match_token_num // block_size,
                     )
@@ -1111,7 +1111,9 @@ def write_cache_to_storage(self, request: Request):
         if isinstance(token_ids, np.ndarray):
             token_ids = token_ids.tolist()
         if self.config.cache_config.enable_output_caching:
-            token_ids += request.output_token_ids
+            input_token_ids = token_ids + request.output_token_ids
+        else:
+            input_token_ids = token_ids
 
         req_id = request.request_id
         keys = []
@@ -1128,7 +1130,7 @@ def write_cache_to_storage(self, request: Request):
         write_storage_task = WriteStorageTask(
             task_id=req_id,
             keys=keys,
-            token_ids=token_ids,
+            token_ids=input_token_ids if self.kvcache_storage_backend == "attention_store" else None,
             gpu_block_ids=gpu_block_ids,
         )
         logger.debug(f"issue write storage task: {write_storage_task}")
@@ -2067,7 +2069,7 @@ def recv_data_transfer_result(self):
                 event_type = data[0]
 
                 if event_type.value == CacheStatus.STORAGE2GPU.value:
-                    logger.info(f"recv_data_transfer_result: {data}")
+                    logger.debug(f"recv_data_transfer_result: {data}")
                     task_id, hash_keys, block_ids = data[1:]
                     if task_id not in self.storage_prefetch_block_ids:
                         self.storage_prefetch_block_ids[task_id] = []
@@ -2078,7 +2080,7 @@ def recv_data_transfer_result(self):
                         if task_id in self.task_prefetch_event:
                             self.task_prefetch_event[task_id].set()
                 elif event_type.value == CacheStatus.GPU2STORAGE.value:
-                    logger.info(f"recv_data_transfer_result: {data}")
+                    logger.debug(f"recv_data_transfer_result: {data}")
                     task_id, hash_keys, block_ids = data[1:]
                     if task_id in self.task_write_back_event:
                         self.task_write_back_event[task_id].set()

fastdeploy/engine/common_engine.py

@@ -891,7 +891,7 @@ def _fetch_request():
                                 self.split_connector.send_splitwise_tasks([task], task.idx)
                                 status, msg = self.split_connector.check_decode_allocated(task)
                                 if not status:
-                                    self.llm_logger.error(
+                                    self.llm_logger.warning(
                                         f"D failed to allocate resource for request {task.request_id}, try again."
                                     )
                                     time.sleep(0.05)

fastdeploy/engine/sched/resource_manager_v1.py

@@ -367,6 +367,9 @@ def _trigger_preempt(self, request, num_new_blocks, preempted_reqs, scheduled_re
                     self._free_blocks(preempted_req)
                     llm_logger.info(f"Preemption is triggered! Preempted request id: {preempted_req.request_id}")
                 else:
+                    if envs.FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST:
+                        if self.config.cache_config.kvcache_storage_backend:
+                            self.cache_manager.write_cache_to_storage(preempted_req)
                     self._free_blocks(preempted_req)
                     preempted_req.num_cached_blocks = 0
                     self.to_be_rescheduled_request_id_set.add(preempted_req.request_id)

fastdeploy/envs.py

@@ -230,6 +230,10 @@
     # When v1 is enabled, the legacy /clear_load_weight and /update_model_weight
     # will adopt this new communication pattern.
     "FD_ENABLE_V1_UPDATE_WEIGHTS": lambda: bool(int(os.getenv("FD_ENABLE_V1_UPDATE_WEIGHTS", "0"))),
+    # Whether to save the cache of output token for preempted request to storage.
+    "FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST": lambda: bool(
+        int(os.getenv("FD_SAVE_OUTPUT_CACHE_FOR_PREEMPTED_REQUEST", "1"))
+    ),
 }