diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 5e16b4352..11add894e 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -60,6 +60,7 @@ class DecodeKVCacheOffloadManager: self.tp_group = tp_group self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group) + self.cache_controller = HiCacheController( token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, mem_pool_host=self.decode_host_mem_pool, @@ -77,41 +78,59 @@ class DecodeKVCacheOffloadManager: logger.info("Enable offload kv cache for decode side") def offload_kv_cache(self, req) -> bool: - """Offload a finished request's KV cache to storage.""" + """Offload incremental KV cache for decode side.""" if self.cache_controller is None or self.decode_host_mem_pool is None: return False - if req.req_pool_idx == -1: + if req.req_pool_idx == -1 or len(req.output_ids) == 0: return False token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx] if token_indices.dim() == 0 or token_indices.numel() == 0: - logger.debug( - f"Request {req.rid} has invalid token_indices: {token_indices}" - ) return False - tokens = req.origin_input_ids + req.output_ids - aligned_len = (len(tokens) // self.page_size) * self.page_size - if aligned_len == 0: + # Prefill side offloads page-aligned origin_input_ids, decode side offloads the incremental part + all_tokens = req.origin_input_ids + req.output_ids[:-1] + prefill_offloaded_len = ( + len(req.origin_input_ids) // self.page_size * self.page_size + ) + incremental_len = len(all_tokens) - prefill_offloaded_len + incremental_aligned_len = incremental_len // self.page_size * self.page_size + + if incremental_aligned_len == 0: return False - token_indices = token_indices[:aligned_len] - tokens = tokens[:aligned_len] + # Extract incremental tokens and indices + start, end = ( + prefill_offloaded_len, + prefill_offloaded_len + incremental_aligned_len, + ) + incremental_tokens = all_tokens[start:end] + incremental_indices = token_indices[start:end] - # Asynchronously offload KV cache from device to host by cache controller + # Early free prefill-offloaded GPU memory + if prefill_offloaded_len > 0: + self.token_to_kv_pool_allocator.free(token_indices[:prefill_offloaded_len]) + + # Asynchronously offload incremental KV cache from device to host self.request_counter += 1 ack_id = self.request_counter host_indices = self.cache_controller.write( - device_indices=token_indices.long(), + device_indices=incremental_indices.long(), node_id=ack_id, ) if host_indices is None: logger.error(f"Not enough host memory for request {req.rid}") return False - self.ongoing_offload[ack_id] = (req, host_indices, tokens, time.time()) + self.ongoing_offload[ack_id] = ( + req, + host_indices, + incremental_tokens, + time.time(), + prefill_offloaded_len, + ) return True def check_offload_progress(self): @@ -140,15 +159,34 @@ class DecodeKVCacheOffloadManager: _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0) finish_event.synchronize() for ack_id in ack_list: - req, host_indices, tokens, start_time = self.ongoing_offload.pop(ack_id) + ( + req, + host_indices, + incremental_tokens, + start_time, + prefill_offloaded_len, + ) = self.ongoing_offload.pop(ack_id) - # Release device - self.tree_cache.cache_finished_req(req) - - # Trigger async backup from host to storage by cache controller - self._trigger_backup(req.rid, host_indices, tokens, start_time) + self._release_finished_req(req, prefill_offloaded_len) + self._trigger_backup( + req, + host_indices, + incremental_tokens, + start_time, + prefill_offloaded_len, + ) finish_count -= 1 + def _release_finished_req(self, req, prefill_offloaded_len): + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, + : len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0), + ] + + # Free the incremental part of the request + self.token_to_kv_pool_allocator.free(kv_indices[prefill_offloaded_len:]) + self.req_to_token_pool.free(req.req_pool_idx) + def _check_backup_progress(self, finish_count): """Check the progress of backup from host to storage.""" for _ in range(finish_count): @@ -159,25 +197,30 @@ class DecodeKVCacheOffloadManager: # Release host memory self.decode_host_mem_pool.free(host_indices) - logger.debug( + logger.info( f"Finished backup request {req_id}, free host memory, len:{len(host_indices)}, cost time:{time.time() - start_time:.2f} seconds." ) - def _trigger_backup(self, req_id, host_indices, tokens, start_time): - """Trigger async backup from host to storage by cache controller.""" + def _trigger_backup( + self, req, host_indices, incremental_tokens, start_time, prefill_offloaded_len + ): + """Trigger async backup from host to storage.""" + prefill_hashes = self._compute_prefix_hash( + req.origin_input_ids[:prefill_offloaded_len] + ) + last_prefill_hash = prefill_hashes[-1] if prefill_offloaded_len > 0 else "" - # Generate page hashes and write to storage - page_hashes = self._compute_prefix_hash(tokens) + page_hashes = self._compute_prefix_hash(incremental_tokens, last_prefill_hash) ack_id = self.cache_controller.write_storage( host_indices, - tokens, + incremental_tokens, hash_value=page_hashes, ) - self.ongoing_backup[ack_id] = (req_id, host_indices, start_time) + self.ongoing_backup[ack_id] = (req.rid, host_indices, start_time) - def _compute_prefix_hash(self, tokens): - last_hash = "" + def _compute_prefix_hash(self, tokens, prior_hash=""): page_hashes = [] + last_hash = prior_hash for offset in range(0, len(tokens), self.page_size): page_tokens = tokens[offset : offset + self.page_size] last_hash = self.cache_controller.get_hash_str(page_tokens, last_hash)