diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 412bb9b80..c441a88d5 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -93,25 +93,29 @@ class DecodeKVCacheOffloadManager: self.ongoing_offload = {} self.ongoing_backup = {} - self.offloaded_state = {} - self.offload_inflight = {} + # A caller may reuse a rid as soon as the previous response finishes, + # while that request's asynchronous D2H copy is still in flight. Key + # lifecycle state by the Req instance so a late ack cannot mutate the + # new request's state. + self.offloaded_state: dict[Req, OffloadedState] = {} + self.offload_inflight: dict[Req, int] = {} logger.info("Enable offload kv cache for decode side") def release_host_resources(self) -> None: self.decode_host_mem_pool.destroy() - def _mark_offload_started(self, rid): - self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1 + def _mark_offload_started(self, req: Req): + self.offload_inflight[req] = self.offload_inflight.get(req, 0) + 1 - def _mark_offload_finished(self, rid): - count = self.offload_inflight.get(rid, 0) + def _mark_offload_finished(self, req: Req): + count = self.offload_inflight.get(req, 0) if count <= 1: - self.offload_inflight.pop(rid, None) + self.offload_inflight.pop(req, None) else: - self.offload_inflight[rid] = count - 1 + self.offload_inflight[req] = count - 1 - def _has_inflight_offload(self, rid): - return self.offload_inflight.get(rid, 0) > 0 + def _has_inflight_offload(self, req: Req): + return self.offload_inflight.get(req, 0) > 0 def offload_kv_cache(self, req) -> bool: """Offload incremental KV cache for decode side.""" @@ -131,7 +135,7 @@ class DecodeKVCacheOffloadManager: prefill_offloaded_len = ( len(req.origin_input_ids) // self.page_size * self.page_size ) - state = self.offloaded_state.get(req.rid) + state = self.offloaded_state.get(req) if state is None: prefill_hashes = self._compute_prefix_hash( req.origin_input_ids[:prefill_offloaded_len] @@ -144,7 +148,7 @@ class DecodeKVCacheOffloadManager: inc_len=0, last_hash=last_prefill_hash, ) - self.offloaded_state[req.rid] = state + self.offloaded_state[req] = state incremental_total = len(all_tokens) - state.prefill_len incremental_new = incremental_total - state.inc_len incremental_aligned_len = ( @@ -177,7 +181,7 @@ class DecodeKVCacheOffloadManager: logger.error(f"Not enough host memory for request {req.rid}") return False - self._mark_offload_started(req.rid) + self._mark_offload_started(req) self.ongoing_offload[ack_id] = ( req, host_indices, @@ -224,20 +228,20 @@ class DecodeKVCacheOffloadManager: end, ) = self.ongoing_offload.pop(ack_id) - self._mark_offload_finished(req.rid) + self._mark_offload_finished(req) prior_hash = ( - self.offloaded_state[req.rid].last_hash - if req.rid in self.offloaded_state + self.offloaded_state[req].last_hash + if req in self.offloaded_state else None ) last_hash = self._trigger_backup( req, host_indices, incremental_tokens, start_time, prior_hash ) - if req.rid in self.offloaded_state: - self.offloaded_state[req.rid].last_hash = last_hash + if req in self.offloaded_state: + self.offloaded_state[req].last_hash = last_hash - if req.finished() and not self._has_inflight_offload(req.rid): - state = self.offloaded_state.get(req.rid) + if req.finished() and not self._has_inflight_offload(req): + state = self.offloaded_state.get(req) start_offset = state.prefill_len if state is not None else start self._release_finished_req(req, start_offset) finish_count -= 1 @@ -257,7 +261,7 @@ class DecodeKVCacheOffloadManager: # concurrent admission. Now consolidated here at request # finish, where the request is guaranteed to no longer attend # to those slots. - state = self.offloaded_state.get(req.rid) + state = self.offloaded_state.get(req) if state is not None and state.prefill_len > 0: prefill_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : state.prefill_len @@ -283,8 +287,7 @@ class DecodeKVCacheOffloadManager: self.req_to_token_pool.free(req) req.kv.mark_released() self.tree_cache.protected_size_ -= len(req.prefix_indices) - if req.rid in self.offloaded_state: - del self.offloaded_state[req.rid] + self.offloaded_state.pop(req, None) def _check_backup_progress(self, finish_count): """Check the progress of backup from host to storage.""" @@ -328,7 +331,7 @@ class DecodeKVCacheOffloadManager: # guard against both sentinels here. if req.req_pool_idx is None or req.req_pool_idx == -1: return - state = self.offloaded_state.get(req.rid) + state = self.offloaded_state.get(req) if state is None: prefill_len = len(req.origin_input_ids) // self.page_size * self.page_size inc_len = 0 @@ -338,10 +341,10 @@ class DecodeKVCacheOffloadManager: # Prefill-aligned slots are freed by _release_finished_req. Make # sure state exists so it can find prefill_len. if state is None: - self.offloaded_state[req.rid] = OffloadedState( + self.offloaded_state[req] = OffloadedState( prefill_len=prefill_len, inc_len=0, last_hash=None ) - if self._has_inflight_offload(req.rid): + if self._has_inflight_offload(req): return start_offset = prefill_len self._release_finished_req(req, start_offset) diff --git a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py index 67522f853..c4b40bd4b 100644 --- a/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py +++ b/test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py @@ -175,7 +175,7 @@ class TestReleaseFinishedReq(unittest.TestCase): def test_release_finished_req_frees_prefill_when_state_present(self): """ - When offloaded_state[rid].prefill_len > 0, _release_finished_req must + When offloaded_state[req].prefill_len > 0, _release_finished_req must free the prefill-aligned slots in addition to the committed range. This is the consolidated free path that replaces the eager free that @@ -190,7 +190,7 @@ class TestReleaseFinishedReq(unittest.TestCase): kv_allocated_len=20, rid=rid, ) - manager.offloaded_state[rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=8, inc_len=0, last_hash=None ) @@ -203,7 +203,7 @@ class TestReleaseFinishedReq(unittest.TestCase): self.assertTrue(torch.equal(freed[0], expected_prefill)) self.assertTrue(torch.equal(freed[1], expected_committed)) # State entry is removed at the end of _release_finished_req. - self.assertNotIn(rid, manager.offloaded_state) + self.assertNotIn(req, manager.offloaded_state) def test_release_finished_req_skips_prefill_free_when_prefill_len_zero(self): """ @@ -219,7 +219,7 @@ class TestReleaseFinishedReq(unittest.TestCase): kv_allocated_len=10, rid=rid, ) - manager.offloaded_state[rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=0, inc_len=0, last_hash=None ) @@ -259,7 +259,7 @@ class TestReleaseFinishedReq(unittest.TestCase): self.assertTrue(torch.equal(freed[0], expected_prefill)) self.assertTrue(torch.equal(freed[1], expected_committed)) # State is deleted by _release_finished_req on the way out. - self.assertNotIn(rid, manager.offloaded_state) + self.assertNotIn(req, manager.offloaded_state) def test_unfinished_offload_ack_does_not_free_incremental_slots(self): manager, freed = _make_manager(pool_size=32) @@ -267,10 +267,10 @@ class TestReleaseFinishedReq(unittest.TestCase): req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=1 ) req.finished.return_value = False - manager.offloaded_state[req.rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=4, inc_len=4, last_hash=None ) - manager.offload_inflight[req.rid] = 1 + manager.offload_inflight[req] = 1 manager.ongoing_offload[7] = ( req, torch.arange(4, 8, dtype=torch.int64), @@ -289,7 +289,7 @@ class TestReleaseFinishedReq(unittest.TestCase): self.assertEqual(freed, []) manager.req_to_token_pool.free.assert_not_called() - self.assertNotIn(req.rid, manager.offload_inflight) + self.assertNotIn(req, manager.offload_inflight) def test_offload_kv_cache_tracks_inflight_write_until_ack(self): manager, freed = _make_manager(pool_size=32, page_size=4) @@ -312,8 +312,8 @@ class TestReleaseFinishedReq(unittest.TestCase): did_offload = manager.offload_kv_cache(req) self.assertTrue(did_offload) - self.assertEqual(manager.offload_inflight[req.rid], 1) - self.assertEqual(manager.offloaded_state[req.rid].inc_len, 4) + self.assertEqual(manager.offload_inflight[req], 1) + self.assertEqual(manager.offloaded_state[req].inc_len, 4) manager.cache_controller.write.assert_called_once() manager.cache_controller.ack_write_queue = [ @@ -324,23 +324,70 @@ class TestReleaseFinishedReq(unittest.TestCase): manager._check_offload_progress(1) self.assertEqual(freed, []) - self.assertNotIn(req.rid, manager.offload_inflight) + self.assertNotIn(req, manager.offload_inflight) + + def test_reused_rid_does_not_share_offload_lifecycle(self): + manager, _ = _make_manager(pool_size=32, page_size=4) + manager.cache_controller = MagicMock() + manager.cache_controller.get_hash_str.return_value = "prefill_hash" + manager.cache_controller.write.return_value = torch.arange( + 4, 8, dtype=torch.int64 + ) + manager.decode_host_mem_pool = MagicMock() + manager.request_counter = 0 + manager.offload_stride = 4 + + old_req = _make_mock_req( + req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid="reused" + ) + new_req = _make_mock_req( + req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid="reused" + ) + for req in (old_req, new_req): + req.origin_input_ids = [0, 1, 2, 3] + req.output_ids = [4, 5, 6, 7, 8] + req.finished.return_value = False + + self.assertTrue(manager.offload_kv_cache(old_req)) + old_req.finished.return_value = True + + # A completed request leaves the API before its asynchronous D2H copy + # necessarily finishes, so a caller can reuse the same rid here. + self.assertTrue(manager.offload_kv_cache(new_req)) + self.assertIsNot(old_req, new_req) + self.assertIn(old_req, manager.offloaded_state) + self.assertIn(new_req, manager.offloaded_state) + self.assertEqual(manager.offload_inflight[old_req], 1) + self.assertEqual(manager.offload_inflight[new_req], 1) + + manager.cache_controller.ack_write_queue = [ + HiCacheAck(None, _FinishedEvent(), [1]) + ] + manager._trigger_backup = MagicMock(return_value="old_last_hash") + manager._check_offload_progress(1) + + self.assertNotIn(old_req, manager.offloaded_state) + self.assertNotIn(old_req, manager.offload_inflight) + self.assertIn(new_req, manager.offloaded_state) + self.assertEqual(manager.offloaded_state[new_req].inc_len, 4) + self.assertEqual(manager.offload_inflight[new_req], 1) + self.assertIn(2, manager.ongoing_offload) def test_finalize_release_defers_while_offload_is_in_flight(self): manager, freed = _make_manager(pool_size=32) req = _make_mock_req( req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=2 ) - manager.offloaded_state[req.rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=4, inc_len=8, last_hash=None ) - manager.offload_inflight[req.rid] = 1 + manager.offload_inflight[req] = 1 manager.finalize_release_on_finish(req) self.assertEqual(freed, []) manager.req_to_token_pool.free.assert_not_called() - self.assertIn(req.rid, manager.offloaded_state) + self.assertIn(req, manager.offloaded_state) def test_finished_offload_ack_waits_for_other_inflight_writes(self): manager, freed = _make_manager(pool_size=32) @@ -348,10 +395,10 @@ class TestReleaseFinishedReq(unittest.TestCase): req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=3 ) req.finished.return_value = True - manager.offloaded_state[req.rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=4, inc_len=8, last_hash=None ) - manager.offload_inflight[req.rid] = 2 + manager.offload_inflight[req] = 2 manager.ongoing_offload[8] = ( req, torch.arange(4, 8, dtype=torch.int64), @@ -370,7 +417,7 @@ class TestReleaseFinishedReq(unittest.TestCase): self.assertEqual(freed, []) manager.req_to_token_pool.free.assert_not_called() - self.assertEqual(manager.offload_inflight[req.rid], 1) + self.assertEqual(manager.offload_inflight[req], 1) def test_finished_request_releases_all_committed_slots_after_last_offload_ack( self, @@ -380,10 +427,10 @@ class TestReleaseFinishedReq(unittest.TestCase): req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=4 ) req.finished.return_value = True - manager.offloaded_state[req.rid] = OffloadedState( + manager.offloaded_state[req] = OffloadedState( prefill_len=4, inc_len=8, last_hash=None ) - manager.offload_inflight[req.rid] = 1 + manager.offload_inflight[req] = 1 manager.ongoing_offload[9] = ( req, torch.arange(8, 12, dtype=torch.int64), @@ -404,8 +451,8 @@ class TestReleaseFinishedReq(unittest.TestCase): self.assertTrue(torch.equal(freed[0], torch.arange(0, 4, dtype=torch.int64))) self.assertTrue(torch.equal(freed[1], torch.arange(4, 20, dtype=torch.int64))) manager.req_to_token_pool.free.assert_called_once_with(req) - self.assertNotIn(req.rid, manager.offloaded_state) - self.assertNotIn(req.rid, manager.offload_inflight) + self.assertNotIn(req, manager.offloaded_state) + self.assertNotIn(req, manager.offload_inflight) if __name__ == "__main__":