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

Revert "[HiCache] Enable spec v2 + decode KV cache offloading compatibility" (#19613)

Browse Source
This commit is contained in:
Kangyan-Zhou
2026-02-28 21:54:32 -08:00
committed by GitHub
parent 8167346609
commit dcf462cfba
3 changed files with 7 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py
View File

@@ -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)

7
python/sglang/srt/server_args.py
View File

@@ -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].")

178
test/registered/disaggregation/test_specv2_kvcache_offloading.py
View File

@@ -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()