minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix(hicache): isolate decode offload state per request (#37026)

Browse Source 
Co-authored-by: Guanyi Chen <939416532@qq.com>
This commit is contained in:
EchO
2026-08-29 18:15:05 -07:00
committed by GitHub
co-authored by Guanyi Chen
parent fbecd75c83
commit ca8ff035c3
2 changed files with 97 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
+29 -26
View File
@@ -93,25 +93,29 @@ class DecodeKVCacheOffloadManager:
self.ongoing_offload = {} self.ongoing_offload = {}
self.ongoing_backup = {} self.ongoing_backup = {}
self.offloaded_state = {} # A caller may reuse a rid as soon as the previous response finishes,
self.offload_inflight = {} # 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") logger.info("Enable offload kv cache for decode side")
def release_host_resources(self) -> None: def release_host_resources(self) -> None:
self.decode_host_mem_pool.destroy() self.decode_host_mem_pool.destroy()
def _mark_offload_started(self, rid): def _mark_offload_started(self, req: Req):
self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1 self.offload_inflight[req] = self.offload_inflight.get(req, 0) + 1
def _mark_offload_finished(self, rid): def _mark_offload_finished(self, req: Req):
count = self.offload_inflight.get(rid, 0) count = self.offload_inflight.get(req, 0)
if count <= 1: if count <= 1:
self.offload_inflight.pop(rid, None) self.offload_inflight.pop(req, None)
else: else:
self.offload_inflight[rid] = count - 1 self.offload_inflight[req] = count - 1
def _has_inflight_offload(self, rid): def _has_inflight_offload(self, req: Req):
return self.offload_inflight.get(rid, 0) > 0 return self.offload_inflight.get(req, 0) > 0
def offload_kv_cache(self, req) -> bool: def offload_kv_cache(self, req) -> bool:
"""Offload incremental KV cache for decode side.""" """Offload incremental KV cache for decode side."""
@@ -131,7 +135,7 @@ class DecodeKVCacheOffloadManager:
prefill_offloaded_len = ( prefill_offloaded_len = (
len(req.origin_input_ids) // self.page_size * self.page_size 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: if state is None:
prefill_hashes = self._compute_prefix_hash( prefill_hashes = self._compute_prefix_hash(
req.origin_input_ids[:prefill_offloaded_len] req.origin_input_ids[:prefill_offloaded_len]
@@ -144,7 +148,7 @@ class DecodeKVCacheOffloadManager:
inc_len=0, inc_len=0,
last_hash=last_prefill_hash, 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_total = len(all_tokens) - state.prefill_len
incremental_new = incremental_total - state.inc_len incremental_new = incremental_total - state.inc_len
incremental_aligned_len = ( incremental_aligned_len = (
@@ -177,7 +181,7 @@ class DecodeKVCacheOffloadManager:
logger.error(f"Not enough host memory for request {req.rid}") logger.error(f"Not enough host memory for request {req.rid}")
return False return False
self._mark_offload_started(req.rid) self._mark_offload_started(req)
self.ongoing_offload[ack_id] = ( self.ongoing_offload[ack_id] = (
req, req,
host_indices, host_indices,
@@ -224,20 +228,20 @@ class DecodeKVCacheOffloadManager:
end, end,
) = self.ongoing_offload.pop(ack_id) ) = self.ongoing_offload.pop(ack_id)
self._mark_offload_finished(req.rid) self._mark_offload_finished(req)
prior_hash = ( prior_hash = (
self.offloaded_state[req.rid].last_hash self.offloaded_state[req].last_hash
if req.rid in self.offloaded_state if req in self.offloaded_state
else None else None
) )
last_hash = self._trigger_backup( last_hash = self._trigger_backup(
req, host_indices, incremental_tokens, start_time, prior_hash req, host_indices, incremental_tokens, start_time, prior_hash
) )
if req.rid in self.offloaded_state: if req in self.offloaded_state:
self.offloaded_state[req.rid].last_hash = last_hash self.offloaded_state[req].last_hash = last_hash
if req.finished() and not self._has_inflight_offload(req.rid): if req.finished() and not self._has_inflight_offload(req):
state = self.offloaded_state.get(req.rid) state = self.offloaded_state.get(req)
start_offset = state.prefill_len if state is not None else start start_offset = state.prefill_len if state is not None else start
self._release_finished_req(req, start_offset) self._release_finished_req(req, start_offset)
finish_count -= 1 finish_count -= 1
@@ -257,7 +261,7 @@ class DecodeKVCacheOffloadManager:
# concurrent admission. Now consolidated here at request # concurrent admission. Now consolidated here at request
# finish, where the request is guaranteed to no longer attend # finish, where the request is guaranteed to no longer attend
# to those slots. # 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: if state is not None and state.prefill_len > 0:
prefill_indices = self.req_to_token_pool.req_to_token[ prefill_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : state.prefill_len req.req_pool_idx, : state.prefill_len
@@ -283,8 +287,7 @@ class DecodeKVCacheOffloadManager:
self.req_to_token_pool.free(req) self.req_to_token_pool.free(req)
req.kv.mark_released() req.kv.mark_released()
self.tree_cache.protected_size_ -= len(req.prefix_indices) self.tree_cache.protected_size_ -= len(req.prefix_indices)
if req.rid in self.offloaded_state: self.offloaded_state.pop(req, None)
del self.offloaded_state[req.rid]
def _check_backup_progress(self, finish_count): def _check_backup_progress(self, finish_count):
"""Check the progress of backup from host to storage.""" """Check the progress of backup from host to storage."""
@@ -328,7 +331,7 @@ class DecodeKVCacheOffloadManager:
# guard against both sentinels here. # guard against both sentinels here.
if req.req_pool_idx is None or req.req_pool_idx == -1: if req.req_pool_idx is None or req.req_pool_idx == -1:
return return
state = self.offloaded_state.get(req.rid) state = self.offloaded_state.get(req)
if state is None: if state is None:
prefill_len = len(req.origin_input_ids) // self.page_size * self.page_size prefill_len = len(req.origin_input_ids) // self.page_size * self.page_size
inc_len = 0 inc_len = 0
@@ -338,10 +341,10 @@ class DecodeKVCacheOffloadManager:
# Prefill-aligned slots are freed by _release_finished_req. Make # Prefill-aligned slots are freed by _release_finished_req. Make
# sure state exists so it can find prefill_len. # sure state exists so it can find prefill_len.
if state is None: 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 prefill_len=prefill_len, inc_len=0, last_hash=None
) )
if self._has_inflight_offload(req.rid): if self._has_inflight_offload(req):
return return
start_offset = prefill_len start_offset = prefill_len
self._release_finished_req(req, start_offset) self._release_finished_req(req, start_offset)
test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py
+68 -21
View File
@@ -175,7 +175,7 @@ class TestReleaseFinishedReq(unittest.TestCase):
def test_release_finished_req_frees_prefill_when_state_present(self): 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. free the prefill-aligned slots in addition to the committed range.
This is the consolidated free path that replaces the eager free that This is the consolidated free path that replaces the eager free that
@@ -190,7 +190,7 @@ class TestReleaseFinishedReq(unittest.TestCase):
kv_allocated_len=20, kv_allocated_len=20,
rid=rid, rid=rid,
) )
manager.offloaded_state[rid] = OffloadedState( manager.offloaded_state[req] = OffloadedState(
prefill_len=8, inc_len=0, last_hash=None 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[0], expected_prefill))
self.assertTrue(torch.equal(freed[1], expected_committed)) self.assertTrue(torch.equal(freed[1], expected_committed))
# State entry is removed at the end of _release_finished_req. # 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): 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, kv_allocated_len=10,
rid=rid, rid=rid,
) )
manager.offloaded_state[rid] = OffloadedState( manager.offloaded_state[req] = OffloadedState(
prefill_len=0, inc_len=0, last_hash=None 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[0], expected_prefill))
self.assertTrue(torch.equal(freed[1], expected_committed)) self.assertTrue(torch.equal(freed[1], expected_committed))
# State is deleted by _release_finished_req on the way out. # 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): def test_unfinished_offload_ack_does_not_free_incremental_slots(self):
manager, freed = _make_manager(pool_size=32) 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_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=1
) )
req.finished.return_value = False 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 prefill_len=4, inc_len=4, last_hash=None
) )
manager.offload_inflight[req.rid] = 1 manager.offload_inflight[req] = 1
manager.ongoing_offload[7] = ( manager.ongoing_offload[7] = (
req, req,
torch.arange(4, 8, dtype=torch.int64), torch.arange(4, 8, dtype=torch.int64),
@@ -289,7 +289,7 @@ class TestReleaseFinishedReq(unittest.TestCase):
self.assertEqual(freed, []) self.assertEqual(freed, [])
manager.req_to_token_pool.free.assert_not_called() 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): def test_offload_kv_cache_tracks_inflight_write_until_ack(self):
manager, freed = _make_manager(pool_size=32, page_size=4) 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) did_offload = manager.offload_kv_cache(req)
self.assertTrue(did_offload) self.assertTrue(did_offload)
self.assertEqual(manager.offload_inflight[req.rid], 1) self.assertEqual(manager.offload_inflight[req], 1)
self.assertEqual(manager.offloaded_state[req.rid].inc_len, 4) self.assertEqual(manager.offloaded_state[req].inc_len, 4)
manager.cache_controller.write.assert_called_once() manager.cache_controller.write.assert_called_once()
manager.cache_controller.ack_write_queue = [ manager.cache_controller.ack_write_queue = [
@@ -324,23 +324,70 @@ class TestReleaseFinishedReq(unittest.TestCase):
manager._check_offload_progress(1) manager._check_offload_progress(1)
self.assertEqual(freed, []) 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): def test_finalize_release_defers_while_offload_is_in_flight(self):
manager, freed = _make_manager(pool_size=32) manager, freed = _make_manager(pool_size=32)
req = _make_mock_req( req = _make_mock_req(
req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=2 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 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) manager.finalize_release_on_finish(req)
self.assertEqual(freed, []) self.assertEqual(freed, [])
manager.req_to_token_pool.free.assert_not_called() 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): def test_finished_offload_ack_waits_for_other_inflight_writes(self):
manager, freed = _make_manager(pool_size=32) 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_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=3
) )
req.finished.return_value = True 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 prefill_len=4, inc_len=8, last_hash=None
) )
manager.offload_inflight[req.rid] = 2 manager.offload_inflight[req] = 2
manager.ongoing_offload[8] = ( manager.ongoing_offload[8] = (
req, req,
torch.arange(4, 8, dtype=torch.int64), torch.arange(4, 8, dtype=torch.int64),
@@ -370,7 +417,7 @@ class TestReleaseFinishedReq(unittest.TestCase):
self.assertEqual(freed, []) self.assertEqual(freed, [])
manager.req_to_token_pool.free.assert_not_called() 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( def test_finished_request_releases_all_committed_slots_after_last_offload_ack(
self, self,
@@ -380,10 +427,10 @@ class TestReleaseFinishedReq(unittest.TestCase):
req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=4 req_pool_idx=0, kv_committed_len=20, kv_allocated_len=20, rid=4
) )
req.finished.return_value = True 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 prefill_len=4, inc_len=8, last_hash=None
) )
manager.offload_inflight[req.rid] = 1 manager.offload_inflight[req] = 1
manager.ongoing_offload[9] = ( manager.ongoing_offload[9] = (
req, req,
torch.arange(8, 12, dtype=torch.int64), 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[0], torch.arange(0, 4, dtype=torch.int64)))
self.assertTrue(torch.equal(freed[1], torch.arange(4, 20, 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) manager.req_to_token_pool.free.assert_called_once_with(req)
self.assertNotIn(req.rid, manager.offloaded_state) self.assertNotIn(req, manager.offloaded_state)
self.assertNotIn(req.rid, manager.offload_inflight) self.assertNotIn(req, manager.offload_inflight)
if __name__ == "__main__": if __name__ == "__main__":