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

Introduce req.kv container for coupled owned kv field lifecycle (#29427)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:40:38 +08:00
committed by GitHub
parent 201ddeaba1
commit d8d76c4d12
24 changed files with 132 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/disaggregation/test_specv2_kvcache_offloading.py
+3 -2
View File
@@ -8,6 +8,7 @@ Requires: torch, sglang (run in an environment with sglang installed)
"""
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
@@ -34,7 +35,7 @@ def _make_mock_req(
req.rid = rid
req.req_pool_idx = req_pool_idx
req.kv_committed_len = kv_committed_len
req.kv_allocated_len = kv_allocated_len
req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len)
req.kv_committed_freed = False
req.kv_overallocated_freed = False
req.prefix_indices = list(range(prefix_indices_len))
@@ -47,7 +48,7 @@ def _make_mock_req(
def pop_overallocated():
assert not req.kv_overallocated_freed
req.kv_overallocated_freed = True
return req.kv_committed_len, req.kv_allocated_len
return req.kv_committed_len, req.kv.kv_allocated_len
req.pop_committed_kv_cache = pop_committed
req.pop_overallocated_kv_cache = pop_overallocated
test/registered/unit/managers/test_hisparse_unit.py
+5 -5
View File
@@ -50,7 +50,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
fill_ids=origin_input_ids + output_ids,
seqlen=len(origin_input_ids) + len(output_ids),
req_pool_idx=None,
kv_allocated_len=0,
kv=SimpleNamespace(kv_allocated_len=0),
kv_committed_len=0,
finished_reason=None,
hisparse_staging=False,
@@ -218,7 +218,7 @@ class TestHiSparseUnit(unittest.TestCase):
)
self.assertIsNotNone(kv_loc, "KV alloc failed")
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
req.kv_allocated_len = fill_len
req.kv.kv_allocated_len = fill_len
req.kv_committed_len = fill_len
req.full_untruncated_fill_ids = array("q", range(fill_len))
req.extend_range = Range(0, fill_len)
@@ -578,7 +578,7 @@ class TestHiSparseUnit(unittest.TestCase):
seq_len = fill_len + 1
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
req.kv_allocated_len = seq_len
req.kv.kv_allocated_len = seq_len
req.kv_committed_len = seq_len
self.coordinator.map_last_loc_to_buffer(
@@ -769,7 +769,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len],
)
)
self.assertEqual(req.kv_allocated_len, fill_len)
self.assertEqual(req.kv.kv_allocated_len, fill_len)
self.assertEqual(req.kv_committed_len, fill_len)
self.assertEqual(req.extend_range.length, fill_len)
@@ -786,7 +786,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertEqual(allocated_host_indices.numel(), rounded_len)
kv_loc = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.kv_allocated_len
req.req_pool_idx, : req.kv.kv_allocated_len
].clone()
self._cleanup_req(req, kv_loc, logical_only=True)
self._assert_sizes_restored(initial, "pd_decode_prealloc_hisparse")
test/registered/unit/managers/test_kv_page_invariants.py
+2 -2
View File
@@ -48,14 +48,14 @@ class _FakeReq:
self.rid = rid
self.req_pool_idx = rpi
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
class _FakeSlot:
def __init__(self, rpi, committed, allocated):
self.req_pool_idx = rpi
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(kv_allocated_len=allocated, swa_evicted_seqlen=0)
self.is_holding_kv = True
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+3 -2
View File
@@ -27,6 +27,7 @@ register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")
import unittest
from array import array
from types import SimpleNamespace
from unittest.mock import MagicMock
import torch
@@ -81,7 +82,7 @@ class MockReq:
self.prefix_indices = torch.empty(0, dtype=torch.int64)
self.priority = 0
self.kv_committed_len = len(fill_ids)
self.kv_allocated_len = len(fill_ids)
self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids))
self.kv_committed_freed = False
def get_fill_ids(self):
@@ -92,7 +93,7 @@ class MockReq:
return self.kv_committed_len
def pop_overallocated_kv_cache(self):
return (self.kv_committed_len, self.kv_allocated_len)
return (self.kv_committed_len, self.kv.kv_allocated_len)
def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None):
test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py
+1 -1
View File
@@ -23,7 +23,7 @@ class _FakeAllocator:
class _FakeReq:
req_pool_idx = 0
swa_evict_floor = 3
swa_evicted_seqlen = 6
kv = SimpleNamespace(swa_evicted_seqlen=6)
def pop_committed_kv_cache(self):
return 8
test/registered/unit/mem_cache/test_streaming_session_unit.py
+11 -9
View File
@@ -57,13 +57,15 @@ class _FakeReq:
)
self.req_pool_idx = req_pool_idx
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv = SimpleNamespace(
kv_allocated_len=allocated,
swa_evicted_seqlen=0,
)
self.kv_committed_freed = False
self.kv_overallocated_freed = False
self.origin_input_ids = list(range(committed))
self.output_ids = []
self.extra_key = None
self.swa_evicted_seqlen = 0
self.last_node = None
self.cache_protected_len = 0
self.swa_uuid_for_lock = None
@@ -86,7 +88,7 @@ class _FakeReq:
assert not self.kv_overallocated_freed
self.pop_overallocated_calls += 1
self.kv_overallocated_freed = True
return self.kv_committed_len, self.kv_allocated_len
return self.kv_committed_len, self.kv.kv_allocated_len
def test_preabort_detaches_session_and_preserves_slot():
@@ -112,7 +114,7 @@ def test_preabort_detaches_session_and_preserves_slot():
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0),
cache_protected_len=16,
)
@@ -132,7 +134,7 @@ def test_preabort_detaches_session_and_preserves_slot():
slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0
assert slot.kv_committed_len == 48
assert slot.kv_allocated_len == 48
assert slot.kv.kv_allocated_len == 48
assert len(result.device_indices) == 0
@@ -179,7 +181,7 @@ def test_nth_mid_abort_nukes_session_slot():
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=50,
kv_allocated_len=50,
kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
last_node=None,
cache_protected_len=0,
)
@@ -230,14 +232,14 @@ def test_trim_overshoot_postcondition():
req = _FakeReq("session-a", req_pool_idx=0, committed=40, allocated=44)
req.origin_input_ids = list(range(26))
req.output_ids = list(range(14))
req.swa_evicted_seqlen = 42
req.kv.swa_evicted_seqlen = 42
tree_cache._trim_overshoot(req, finished_len=12)
target = 38
assert req.kv_committed_len == target
assert req.kv_allocated_len == target
assert req.swa_evicted_seqlen == target
assert req.kv.kv_allocated_len == target
assert req.kv.swa_evicted_seqlen == target
assert len(req.output_ids) == 12
# Tail [38, 44) freed by _free_kv_aligned.
assert len(allocator.freed) == 1
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+15 -13
View File
@@ -106,7 +106,9 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
origin_input_ids=token_ids,
output_ids=[],
cache_protected_len=cache_protected_len,
swa_evicted_seqlen=0,
kv=SimpleNamespace(
swa_evicted_seqlen=0,
),
extra_key=None,
last_node=tree.root_node,
swa_uuid_for_lock=None,
@@ -162,7 +164,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
insert_len = seq_len // page_size * page_size
self.assertLess(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
insert_len,
f"page={page_size}, win={window}, seq={seq_len}",
)
@@ -187,7 +189,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertLess(req.swa_evicted_seqlen, insert_len)
self.assertLess(req.kv.swa_evicted_seqlen, insert_len)
tree.cache_finished_req(req, is_insert=True)
tree.sanity_check()
@@ -209,9 +211,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertLess(req.swa_evicted_seqlen, seq_len)
self.assertLess(req.kv.swa_evicted_seqlen, seq_len)
self.assertEqual(
req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
)
tree.cache_finished_req(req, is_insert=True)
@@ -235,7 +237,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
batch = _make_batch(tree, allocator, pool)
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertEqual(req.swa_evicted_seqlen, 0)
self.assertEqual(req.kv.swa_evicted_seqlen, 0)
# -- Insert case 1: swa_evicted <= total_prefix_length --
@@ -264,7 +266,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
# pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched.
ScheduleBatch._evict_swa(batch, req2, first_len - 1)
self.assertLessEqual(req2.swa_evicted_seqlen, first_len)
self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len)
swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req2, is_insert=True)
@@ -295,12 +297,12 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertGreater(req.swa_evicted_seqlen, 0, "Should have some eviction")
self.assertLess(req.swa_evicted_seqlen, insert_len, "Should be partial")
self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction")
self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial")
tree.cache_finished_req(req, is_insert=True)
non_tombstone = insert_len - req.swa_evicted_seqlen
non_tombstone = insert_len - req.kv.swa_evicted_seqlen
self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone)
self.assertGreater(tree.full_evictable_size_, 0)
tree.sanity_check()
@@ -333,7 +335,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
allocator.free_swa(pool.req_to_token[0, :old_evicted])
req = _make_req(0, list(range(seq_len)), 0, tree)
req.swa_evicted_seqlen = old_evicted
req.kv.swa_evicted_seqlen = old_evicted
swa_evictable_before = tree.swa_evictable_size_
tree.cache_finished_req(req, is_insert=True)
@@ -363,7 +365,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
insert_len = seq_len // page_size * page_size
self.assertLess(req.swa_evicted_seqlen, insert_len, f"turn {turn}")
self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}")
tree.cache_finished_req(req, is_insert=True)
tree.sanity_check()
@@ -386,7 +388,7 @@ class TestSWAEvictionBoundary(unittest.TestCase):
ScheduleBatch._evict_swa(batch, req, seq_len - 1)
self.assertEqual(
req.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
)
tree.cache_finished_req(req, is_insert=True)
test/registered/unit/mem_cache/test_swa_unittest.py
+4 -2
View File
@@ -1,5 +1,6 @@
import unittest
from array import array
from types import SimpleNamespace
import torch
@@ -35,6 +36,7 @@ class _DummyReq:
def __init__(self):
self._kv_committed_len = 0
self.swa_prefix_lock_released = False
self.kv = SimpleNamespace(swa_evicted_seqlen=0)
def pop_committed_kv_cache(self):
return self._kv_committed_len
@@ -575,7 +577,7 @@ class TestSWA(unittest.TestCase):
req.extra_key = None
req.last_node = tree.root_node
req.swa_uuid_for_lock = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
req.cache_protected_len = 1
# Intentionally mismatch to ensure code does not use len(prefix_indices).
req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device)
@@ -610,7 +612,7 @@ class TestSWA(unittest.TestCase):
req2.extra_key = None
req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None
req2.swa_evicted_seqlen = 0
req2.kv.swa_evicted_seqlen = 0
req2.cache_protected_len = 1
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+8 -8
View File
@@ -881,7 +881,7 @@ class UnifiedRadixCacheSuite:
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv_committed_len = kv_len
req.kv_allocated_len = kv_len
req.kv.kv_allocated_len = kv_len
req.last_node = cache.root_node
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
@@ -998,7 +998,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = evicted_len
req.kv.swa_evicted_seqlen = evicted_len
cache.cache_unfinished_req(req)
@@ -1220,7 +1220,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
full_available_before_insert = allocator.full_attn_allocator.available_size()
@@ -1761,7 +1761,7 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
swa_avail_before = allocator.swa_attn_allocator.available_size()
@@ -1771,10 +1771,10 @@ class UnifiedRadixCacheSuite:
cushion = max(self.cfg.sliding_window_size, self.cfg.page_size)
expected_evicted = (pre_len - 1) - cushion
self.assertEqual(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
expected_evicted,
f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = "
f"{expected_evicted}, got {req.swa_evicted_seqlen}",
f"{expected_evicted}, got {req.kv.swa_evicted_seqlen}",
)
swa_avail_after = allocator.swa_attn_allocator.available_size()
@@ -1813,13 +1813,13 @@ class UnifiedRadixCacheSuite:
req.cache_protected_len = 0
req.swa_uuid_for_lock = None
req.extra_key = None
req.swa_evicted_seqlen = 0
req.kv.swa_evicted_seqlen = 0
with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
cache.cache_unfinished_req(req)
self.assertEqual(
req.swa_evicted_seqlen,
req.kv.swa_evicted_seqlen,
0,
"Nothing should be evicted when prefill fits inside the cushion",
)
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
-2
View File
@@ -80,9 +80,7 @@ _OWNER_SITES = {
): 1,
# streaming session slot save/restore and tail trimming
(_SS, "SessionSlot.save_from_req", "kv_committed_len"): 1,
(_SS, "SessionSlot.save_from_req", "kv_allocated_len"): 1,
(_SS, "SessionSlot.restore_to_req", "kv_committed_len"): 1,
(_SS, "SessionSlot.restore_to_req", "kv_allocated_len"): 1,
(_SS, "StreamingSession._free_tail", "kv_committed_len"): 2,
(_SS, "StreamingSession._free_tail", "kv_allocated_len"): 2,
(_SS, "StreamingSession._trim_overshoot", "kv_committed_len"): 1,