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_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)
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):
"""
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__":