[mem_cache] Move cache_protected_len and swa_evict_floor into ReqKvInfo (#36982)
This commit is contained in:
@@ -304,7 +304,7 @@ class TestRadixNoTailChunked(ScriptedTestCase):
|
||||
if req is not None and req.rid == r.rid:
|
||||
observed_mid_chunk = True
|
||||
prefix_len: int = len(req.prefix_indices)
|
||||
protected_len: int = req.cache_protected_len
|
||||
protected_len: int = req.kv.cache_protected_len
|
||||
assert prefix_len == protected_len, (
|
||||
f"page_size=1 must take the no-tail else branch: "
|
||||
f"len(prefix_indices)={prefix_len} != "
|
||||
@@ -460,7 +460,7 @@ class TestRadixPartialPage(ScriptedTestCase):
|
||||
req = s.chunked_req
|
||||
if req is not None and req.rid == r.rid:
|
||||
prefix_len: int = len(req.prefix_indices)
|
||||
protected_len: int = req.cache_protected_len
|
||||
protected_len: int = req.kv.cache_protected_len
|
||||
assert prefix_len >= protected_len, (
|
||||
f"len(prefix_indices)={prefix_len} dropped below "
|
||||
f"cache_protected_len={protected_len}: tail was freed "
|
||||
|
||||
@@ -226,7 +226,7 @@ def test_streaming_session_release_frees_compressed_slots():
|
||||
)
|
||||
session.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=req_pool_idx,
|
||||
kv=SimpleNamespace(kv_allocated_len=16),
|
||||
kv=SimpleNamespace(kv_allocated_len=16, cache_protected_len=0),
|
||||
)
|
||||
|
||||
session.release_session("session-a")
|
||||
|
||||
@@ -97,7 +97,7 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
|
||||
finished_reason=FINISH_ABORT("failed") if failed else None,
|
||||
return_logprob=False,
|
||||
sampling_params=SimpleNamespace(max_new_tokens=8),
|
||||
cache_protected_len=0,
|
||||
kv=SimpleNamespace(cache_protected_len=0),
|
||||
time_stats=MagicMock(),
|
||||
)
|
||||
return SimpleNamespace(
|
||||
|
||||
@@ -12,7 +12,7 @@ from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
|
||||
|
||||
maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.schedule_batch import NextBatchPlan, Req
|
||||
from sglang.srt.managers.schedule_batch import NextBatchPlan, Req, ReqKvInfo
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.mem_cache.chunk_cache import ChunkCache
|
||||
from sglang.srt.utils.common import Range
|
||||
@@ -38,7 +38,7 @@ def _make_req(
|
||||
req.extend_range = Range(fill_len - extend_input_len, fill_len)
|
||||
req.inflight_middle_chunks = 0
|
||||
req.host_hit_length = 0
|
||||
req.cache_protected_len = 0
|
||||
req.kv = ReqKvInfo()
|
||||
req.skip_radix_cache_insert = False
|
||||
req.last_node = None
|
||||
req.swa_uuid_for_lock = None
|
||||
|
||||
@@ -76,14 +76,15 @@ class MockReq:
|
||||
)
|
||||
self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else [])
|
||||
self.req_pool_idx = req_pool_idx
|
||||
self.cache_protected_len = cache_protected_len
|
||||
self.last_node = last_node
|
||||
self.extra_key = None
|
||||
self.cache_salt = None
|
||||
self.prefix_indices = torch.empty(0, dtype=torch.int64)
|
||||
self.priority = 0
|
||||
self.kv_committed_len = len(fill_ids)
|
||||
self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids))
|
||||
self.kv = SimpleNamespace(
|
||||
kv_allocated_len=len(fill_ids), cache_protected_len=cache_protected_len
|
||||
)
|
||||
|
||||
def get_fill_ids(self):
|
||||
return self.full_untruncated_fill_ids[: self.extend_range.end]
|
||||
@@ -368,7 +369,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
|
||||
req.output_ids = [99]
|
||||
req.last_node = object()
|
||||
req.finished_reason = None
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = 123
|
||||
req.swa_prefix_lock_released = False
|
||||
req.pd_rebootstrap_in_progress = False
|
||||
|
||||
@@ -22,9 +22,11 @@ class _FakeAllocator:
|
||||
|
||||
class _FakeReq:
|
||||
req_pool_idx = 0
|
||||
swa_evict_floor = 3
|
||||
cache_protected_len = 0
|
||||
kv = SimpleNamespace(swa_evicted_seqlen=6)
|
||||
|
||||
def __init__(self):
|
||||
self.kv = SimpleNamespace(
|
||||
swa_evicted_seqlen=6, swa_evict_floor=3, cache_protected_len=0
|
||||
)
|
||||
|
||||
def pop_committed_kv_cache(self):
|
||||
return 8
|
||||
@@ -51,7 +53,7 @@ class TestPureSWAChunkCache(CustomTestCase):
|
||||
def test_finished_req_skips_protected_prefix(self):
|
||||
cache = self._make_cache()
|
||||
req = _FakeReq()
|
||||
req.cache_protected_len = 2
|
||||
req.kv.cache_protected_len = 2
|
||||
|
||||
cache.cache_finished_req(req, kv_len_to_handle=8)
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ import random
|
||||
import unittest
|
||||
import unittest.mock
|
||||
from array import array
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
@@ -518,7 +519,7 @@ class TestRadixCache(unittest.TestCase):
|
||||
cache.req_to_token_pool = ReqToTokenPool(request_indices.clone())
|
||||
req = unittest.mock.Mock(
|
||||
req_pool_idx=0,
|
||||
cache_protected_len=0,
|
||||
kv=SimpleNamespace(cache_protected_len=0),
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
priority=0,
|
||||
|
||||
@@ -12,6 +12,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
import unittest
|
||||
import unittest.mock
|
||||
from array import array
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
@@ -39,7 +40,7 @@ class _StubReq:
|
||||
self.host_hit_length = None
|
||||
self.num_matched_prefix_tokens = 0
|
||||
self.mamba_branching_seqlen = None
|
||||
self.cache_protected_len = None
|
||||
self.kv = SimpleNamespace(cache_protected_len=None)
|
||||
|
||||
def _compute_max_prefix_len(self, input_len):
|
||||
return max(input_len - 1, 0)
|
||||
|
||||
@@ -72,13 +72,13 @@ class _FakeReq:
|
||||
self.kv = SimpleNamespace(
|
||||
kv_allocated_len=allocated,
|
||||
swa_evicted_seqlen=0,
|
||||
cache_protected_len=0,
|
||||
)
|
||||
self.origin_input_ids = list(range(committed))
|
||||
self.output_ids = []
|
||||
self.extra_key = None
|
||||
self.cache_salt = None
|
||||
self.last_node = None
|
||||
self.cache_protected_len = 0
|
||||
self.swa_uuid_for_lock = None
|
||||
self.skip_lock_node_ids = {}
|
||||
self.mamba_pool_idx = None
|
||||
@@ -114,8 +114,9 @@ def test_preabort_detaches_session_and_preserves_slot():
|
||||
tree_cache.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=0,
|
||||
kv_committed_len=48,
|
||||
kv=SimpleNamespace(kv_allocated_len=48, swa_evicted_seqlen=0),
|
||||
cache_protected_len=16,
|
||||
kv=SimpleNamespace(
|
||||
kv_allocated_len=48, swa_evicted_seqlen=0, cache_protected_len=16
|
||||
),
|
||||
)
|
||||
|
||||
req = _FakeReq("session-a", req_pool_idx=1, committed=1, allocated=1)
|
||||
@@ -178,9 +179,10 @@ def test_nth_mid_abort_nukes_session_slot():
|
||||
tree_cache.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=0,
|
||||
kv_committed_len=50,
|
||||
kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
|
||||
kv=SimpleNamespace(
|
||||
kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0
|
||||
),
|
||||
last_node=None,
|
||||
cache_protected_len=0,
|
||||
)
|
||||
|
||||
# Mid-processing abort: req has the SESSION slot's pool_idx (restore_to_req ran).
|
||||
@@ -216,9 +218,10 @@ def test_release_session_threads_mamba_skip_ids():
|
||||
tree_cache.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=0,
|
||||
kv_committed_len=50,
|
||||
kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
|
||||
kv=SimpleNamespace(
|
||||
kv_allocated_len=50, swa_evicted_seqlen=0, cache_protected_len=0
|
||||
),
|
||||
last_node=lock_node,
|
||||
cache_protected_len=0,
|
||||
skip_lock_node_ids={ComponentType.MAMBA: {42}},
|
||||
)
|
||||
|
||||
|
||||
@@ -18,7 +18,7 @@ from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import ScheduleBatch
|
||||
from sglang.srt.managers.schedule_batch import ReqKvInfo, ScheduleBatch
|
||||
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||
from sglang.srt.mem_cache.common import free_swa_out_of_window_slots
|
||||
@@ -107,10 +107,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
|
||||
is_holding_kv=True,
|
||||
origin_input_ids=token_ids,
|
||||
output_ids=[],
|
||||
cache_protected_len=cache_protected_len,
|
||||
kv=SimpleNamespace(
|
||||
swa_evicted_seqlen=0,
|
||||
),
|
||||
kv=ReqKvInfo(cache_protected_len=cache_protected_len),
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
last_node=tree.root_node,
|
||||
|
||||
@@ -36,7 +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)
|
||||
self.kv = SimpleNamespace(swa_evicted_seqlen=0, cache_protected_len=0)
|
||||
|
||||
|
||||
def _build_swa_tree(
|
||||
@@ -682,7 +682,7 @@ class TestSWA(unittest.TestCase):
|
||||
req.last_node = tree.root_node
|
||||
req.swa_uuid_for_lock = None
|
||||
req.kv.swa_evicted_seqlen = 0
|
||||
req.cache_protected_len = 1
|
||||
req.kv.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)
|
||||
|
||||
@@ -700,7 +700,7 @@ class TestSWA(unittest.TestCase):
|
||||
req, is_insert=True, kv_len_to_handle=req._kv_committed_len
|
||||
)
|
||||
|
||||
self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len)
|
||||
self.assertEqual(captured["prev_prefix_len"], req.kv.cache_protected_len)
|
||||
self.assertTrue(captured["is_bigram"])
|
||||
self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1)
|
||||
|
||||
@@ -720,7 +720,7 @@ class TestSWA(unittest.TestCase):
|
||||
req2.last_node = tree.root_node
|
||||
req2.swa_uuid_for_lock = None
|
||||
req2.kv.swa_evicted_seqlen = 0
|
||||
req2.cache_protected_len = 1
|
||||
req2.kv.cache_protected_len = 1
|
||||
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
|
||||
|
||||
freed_lens = []
|
||||
@@ -917,7 +917,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
|
||||
req.get_fill_ids = lambda: token_ids
|
||||
req.extra_key = None
|
||||
req.cache_salt = None
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.last_node = tree.root_node
|
||||
req.swa_uuid_for_lock = None
|
||||
req.prefix_indices = torch.empty(0, dtype=torch.int64, device=tree.device)
|
||||
@@ -928,7 +928,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
|
||||
# The insert itself frees nothing.
|
||||
self.assertEqual(allocator.swa_available_size(), swa_before)
|
||||
# The live leaf holds a full window, so the whole key stays matchable.
|
||||
self.assertEqual(req.cache_protected_len, num_tokens)
|
||||
self.assertEqual(req.kv.cache_protected_len, num_tokens)
|
||||
# [0, evicted) is a tombstone; only [evicted, num_tokens) counts as SWA.
|
||||
(first,) = tree.root_node.children.values()
|
||||
self.assertTrue(first.swa_tombstone)
|
||||
|
||||
@@ -642,7 +642,7 @@ def bench_cache_finished(
|
||||
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
|
||||
)
|
||||
req.last_node = node
|
||||
req.cache_protected_len = matched_len
|
||||
req.kv.cache_protected_len = matched_len
|
||||
req.kv_committed_len = len(seq)
|
||||
if hasattr(lr, "swa_uuid_for_lock"):
|
||||
req.swa_uuid_for_lock = lr.swa_uuid_for_lock
|
||||
|
||||
@@ -1228,7 +1228,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||
req.kv_committed_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
req.full_untruncated_fill_ids = array("q", input_ids + output_ids)
|
||||
@@ -1269,7 +1269,7 @@ class UnifiedRadixCacheSuite:
|
||||
req.kv_committed_len = kv_len
|
||||
req.kv.kv_allocated_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
if self.cfg.has_mamba:
|
||||
@@ -1313,7 +1313,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||
req.kv_committed_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.swa_prefix_lock_released = True
|
||||
req.extra_key = None
|
||||
@@ -1348,7 +1348,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||
req.kv_committed_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
if self.cfg.has_mamba:
|
||||
@@ -1357,7 +1357,7 @@ class UnifiedRadixCacheSuite:
|
||||
cache.cache_unfinished_req(req)
|
||||
|
||||
self.assertGreater(len(req.prefix_indices), 0)
|
||||
self.assertEqual(req.cache_protected_len, len(req.prefix_indices))
|
||||
self.assertEqual(req.kv.cache_protected_len, len(req.prefix_indices))
|
||||
self.assertIsNotNone(req.last_node)
|
||||
self.assertFalse(req.swa_prefix_lock_released)
|
||||
|
||||
@@ -1385,7 +1385,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
|
||||
req.kv_committed_len = len(tokens)
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
req.kv.swa_evicted_seqlen = evicted_len
|
||||
@@ -1487,7 +1487,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
|
||||
req.kv_committed_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
req.full_untruncated_fill_ids = array("q", input_ids)
|
||||
@@ -1612,7 +1612,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value)
|
||||
req.kv_committed_len = kv_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
req.kv.swa_evicted_seqlen = 0
|
||||
@@ -1639,7 +1639,7 @@ class UnifiedRadixCacheSuite:
|
||||
swa_value,
|
||||
)
|
||||
)
|
||||
self.assertEqual(req.cache_protected_len, len(tokens))
|
||||
self.assertEqual(req.kv.cache_protected_len, len(tokens))
|
||||
|
||||
cache.dec_lock_ref(
|
||||
req.last_node,
|
||||
@@ -2199,7 +2199,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
||||
req.kv_committed_len = pre_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
|
||||
@@ -2288,7 +2288,7 @@ class UnifiedRadixCacheSuite:
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
|
||||
req.kv_committed_len = pre_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
|
||||
@@ -6770,7 +6770,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
|
||||
req.output_ids = array("q")
|
||||
req.kv_committed_len = len(tokens)
|
||||
req.kv.kv_allocated_len = len(tokens)
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
req.mamba_last_track_seqlen = len(tokens)
|
||||
@@ -8097,7 +8097,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
|
||||
req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices)
|
||||
req.kv_committed_len = seq_len
|
||||
req.last_node = cache.root_node_handle()
|
||||
req.cache_protected_len = 0
|
||||
req.kv.cache_protected_len = 0
|
||||
req.swa_uuid_for_lock = None
|
||||
req.extra_key = None
|
||||
|
||||
@@ -8113,7 +8113,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
|
||||
f"{self.cfg.sliding_window_size} window",
|
||||
)
|
||||
self.assertEqual(
|
||||
req.cache_protected_len,
|
||||
req.kv.cache_protected_len,
|
||||
boundary,
|
||||
"the match after the insert must reach the leaf the insert created",
|
||||
)
|
||||
|
||||
@@ -72,7 +72,7 @@ def _make_req(rid, req_pool_idx, token_ids, tree):
|
||||
extra_key=None,
|
||||
cache_salt=None,
|
||||
last_node=tree.root_node,
|
||||
cache_protected_len=0,
|
||||
kv=SimpleNamespace(cache_protected_len=0),
|
||||
priority=0,
|
||||
kv_committed_freed=False,
|
||||
kv_committed_len=len(token_ids),
|
||||
|
||||
Reference in New Issue
Block a user