From dcf462cfbaae7557f213ce2baceaebc5c61c2ccd Mon Sep 17 00:00:00 2001 From: Kangyan-Zhou Date: Sat, 28 Feb 2026 21:54:32 -0800 Subject: [PATCH] Revert "[HiCache] Enable spec v2 + decode KV cache offloading compatibility" (#19613) --- .../decode_kvcache_offload_manager.py | 13 -- python/sglang/srt/server_args.py | 7 + .../test_specv2_kvcache_offloading.py | 178 ------------------ 3 files changed, 7 insertions(+), 191 deletions(-) delete mode 100644 test/registered/disaggregation/test_specv2_kvcache_offloading.py diff --git a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py index 6abd49144..1cfe615a7 100644 --- a/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py +++ b/python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py @@ -20,7 +20,6 @@ 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 @@ -192,18 +191,6 @@ 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 cb71cba1d..a364d175c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2748,6 +2748,13 @@ 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 deleted file mode 100644 index 545531617..000000000 --- a/test/registered/disaggregation/test_specv2_kvcache_offloading.py +++ /dev/null @@ -1,178 +0,0 @@ -""" -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()