diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 1cfe615a7..6abd49144 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -20,6 +20,7 @@ from sglang.srt.mem_cache.memory_pool_host import ( MLATokenToKVPoolHost, ) from sglang.srt.server_args import ServerArgs +from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -191,6 +192,18 @@ class DecodeKVCacheOffloadManager: # Free the incremental part of the request self.token_to_kv_pool_allocator.free(kv_indices) + + # Free over-allocated KV cache slots (e.g. from speculative decoding v2). + # Without spec v2, start_p == end_p so this is a no-op. + start_p, end_p = req.pop_overallocated_kv_cache() + if self.page_size > 1: + start_p = ceil_align(start_p, self.page_size) + if start_p < end_p: + overalloc_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, start_p:end_p + ] + self.token_to_kv_pool_allocator.free(overalloc_indices) + self.req_to_token_pool.free(req) self.tree_cache.protected_size_ -= len(req.prefix_indices) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a364d175c..cb71cba1d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2748,13 +2748,6 @@ class ServerArgs: raise ValueError( "The argument disaggregation-decode-enable-offload-kvcache is only supported for decode side." ) - if ( - self.disaggregation_mode == "decode" - and envs.SGLANG_ENABLE_SPEC_V2.get() - ): - raise ValueError( - "Spec v2 and decode offload kv cache are incompatible and cannot be enabled together." - ) if not (0 < self.swa_full_tokens_ratio <= 1.0): raise ValueError("--swa-full-tokens-ratio should be in range (0, 1.0].") diff --git a/test/registered/disaggregation/test_specv2_kvcache_offloading.py b/test/registered/disaggregation/test_specv2_kvcache_offloading.py new file mode 100644 index 000000000..545531617 --- /dev/null +++ b/test/registered/disaggregation/test_specv2_kvcache_offloading.py @@ -0,0 +1,178 @@ +""" +Unit tests for _release_finished_req in DecodeKVCacheOffloadManager. + +Verifies that over-allocated KV cache slots (from speculative decoding v2) +are correctly freed when a request finishes, preventing GPU memory leaks. + +Requires: torch, sglang (run in an environment with sglang installed) +""" + +import unittest +from unittest.mock import MagicMock + +import torch + +from sglang.srt.disaggregation.decode_kvcache_offload_manager import ( + DecodeKVCacheOffloadManager, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="stage-a-cpu-only") + + +def _make_mock_req( + req_pool_idx: int, + kv_committed_len: int, + kv_allocated_len: int, + prefix_indices_len: int = 0, +): + """Create a mock Req with the KV cache state needed for testing.""" + req = MagicMock() + req.req_pool_idx = req_pool_idx + req.kv_committed_len = kv_committed_len + req.kv_allocated_len = kv_allocated_len + req.kv_committed_freed = False + req.kv_overallocated_freed = False + req.prefix_indices = list(range(prefix_indices_len)) + + def pop_committed(): + assert not req.kv_committed_freed + req.kv_committed_freed = True + return req.kv_committed_len + + def pop_overallocated(): + assert not req.kv_overallocated_freed + req.kv_overallocated_freed = True + return req.kv_committed_len, req.kv_allocated_len + + req.pop_committed_kv_cache = pop_committed + req.pop_overallocated_kv_cache = pop_overallocated + return req + + +def _make_manager(pool_size: int, page_size: int = 1): + """Create a DecodeKVCacheOffloadManager with mock pools for testing.""" + # Build a real req_to_token tensor so indexing works + req_to_token = torch.arange(pool_size, dtype=torch.int64).unsqueeze(0) + + req_to_token_pool = MagicMock() + req_to_token_pool.req_to_token = req_to_token + + freed_indices = [] + + allocator = MagicMock() + allocator.free = MagicMock( + side_effect=lambda idx: freed_indices.append(idx.clone()) + ) + + tree_cache = MagicMock() + tree_cache.protected_size_ = 0 + + # Bypass __init__ entirely and set attributes directly + manager = object.__new__(DecodeKVCacheOffloadManager) + manager.req_to_token_pool = req_to_token_pool + manager.token_to_kv_pool_allocator = allocator + manager.page_size = page_size + manager.tree_cache = tree_cache + + return manager, freed_indices + + +class TestReleaseFinishedReq(unittest.TestCase): + """Tests for _release_finished_req overallocation cleanup.""" + + def test_no_overallocation(self): + """Without spec v2, kv_committed == kv_allocated; no extra free.""" + manager, freed = _make_manager(pool_size=32) + req = _make_mock_req( + req_pool_idx=0, + kv_committed_len=20, + kv_allocated_len=20, # no overallocation + ) + prefill_offloaded_len = 8 + + manager._release_finished_req(req, prefill_offloaded_len) + + # Only one free call: the committed range [8:20] + self.assertEqual(len(freed), 1) + expected = torch.arange(8, 20, dtype=torch.int64) + self.assertTrue(torch.equal(freed[0], expected)) + manager.req_to_token_pool.free.assert_called_once_with(req) + + def test_with_overallocation(self): + """With spec v2, overallocated slots [committed:allocated] must be freed.""" + manager, freed = _make_manager(pool_size=32) + req = _make_mock_req( + req_pool_idx=0, + kv_committed_len=20, + kv_allocated_len=28, # 8 over-allocated slots + ) + prefill_offloaded_len = 8 + + manager._release_finished_req(req, prefill_offloaded_len) + + # Two free calls: committed [8:20] and overallocated [20:28] + self.assertEqual(len(freed), 2) + expected_committed = torch.arange(8, 20, dtype=torch.int64) + expected_overalloc = torch.arange(20, 28, dtype=torch.int64) + self.assertTrue(torch.equal(freed[0], expected_committed)) + self.assertTrue(torch.equal(freed[1], expected_overalloc)) + manager.req_to_token_pool.free.assert_called_once_with(req) + + def test_overallocation_with_page_alignment(self): + """With page_size > 1, start of overallocated range is ceil-aligned.""" + page_size = 4 + manager, freed = _make_manager(pool_size=32, page_size=page_size) + req = _make_mock_req( + req_pool_idx=0, + kv_committed_len=10, # not page-aligned + kv_allocated_len=28, + ) + prefill_offloaded_len = 4 + + manager._release_finished_req(req, prefill_offloaded_len) + + # Committed range [4:10] + # Overallocated: start_p = ceil_align(10, 4) = 12, end_p = 28 => [12:28] + self.assertEqual(len(freed), 2) + expected_committed = torch.arange(4, 10, dtype=torch.int64) + expected_overalloc = torch.arange(12, 28, dtype=torch.int64) + self.assertTrue(torch.equal(freed[0], expected_committed)) + self.assertTrue(torch.equal(freed[1], expected_overalloc)) + + def test_overallocation_page_aligned_noop(self): + """When ceil_align(committed, page_size) >= allocated, no overalloc free.""" + page_size = 4 + manager, freed = _make_manager(pool_size=32, page_size=page_size) + req = _make_mock_req( + req_pool_idx=0, + kv_committed_len=10, # ceil_align(10, 4) = 12 + kv_allocated_len=12, # same as aligned start + ) + prefill_offloaded_len = 4 + + manager._release_finished_req(req, prefill_offloaded_len) + + # Only committed [4:10], no overalloc because start_p == end_p + self.assertEqual(len(freed), 1) + expected_committed = torch.arange(4, 10, dtype=torch.int64) + self.assertTrue(torch.equal(freed[0], expected_committed)) + + def test_prefix_indices_decremented(self): + """protected_size_ is decremented by len(req.prefix_indices).""" + manager, _ = _make_manager(pool_size=32) + manager.tree_cache.protected_size_ = 10 + req = _make_mock_req( + req_pool_idx=0, + kv_committed_len=20, + kv_allocated_len=20, + prefix_indices_len=5, + ) + + manager._release_finished_req(req, prefill_offloaded_len=0) + + self.assertEqual(manager.tree_cache.protected_size_, 5) + + +if __name__ == "__main__": + unittest.main()