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

[mem_cache] Move req_pool_idx into ReqKvInfo (#37094)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-30 14:46:21 -07:00
committed by GitHub
parent 8a87079dbb
commit 007ef5e23a
68 changed files with 372 additions and 357 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/chunked_prefill/test_scripted_abort.py
+14 -14
View File
@@ -23,7 +23,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle):
if all(
h.kv_pages == 0
and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None)
and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles
):
return
@@ -57,7 +57,7 @@ class TestAbortBasic(ScriptedTestCase):
r.kv_pages == 0
), f"abort must release KV; r.kv_pages={r.kv_pages} after abort"
assert (
r.req is None or r.req.req_pool_idx is None
r.req is None or r.req.kv.req_pool_idx is None
), f"abort must release row; r.req={r.req} after abort"
assert (
r.lock_refs == 0
@@ -78,7 +78,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_at_chunk_mid(self):
self.server.execute_script(self._script_abort_at_chunk_mid)
@@ -124,7 +124,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r)
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
def test_abort_at_admission_step(self):
@@ -137,7 +137,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r)
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_then_start_same_step_new_rid(self):
self.server.execute_script(self._script_abort_then_start_same_step_new_rid)
@@ -191,7 +191,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, *reqs)
for r in reqs:
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_unknown_rid_noop(self):
self.server.execute_script(self._script_abort_unknown_rid_noop)
@@ -264,7 +264,7 @@ class TestAbortBasic(ScriptedTestCase):
t.abort(r)
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
def test_double_abort_idempotent(self):
@@ -351,7 +351,7 @@ class TestAbortBasic(ScriptedTestCase):
f"aborted req revived and ran another chunk; "
f"chunks_done went {chunks_after_abort} -> {r.chunks_done}"
)
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
def test_abort_mid_chunk_no_extra_radix_node(self):
self.server.execute_script(self._script_abort_mid_chunk_no_extra_radix_node)
@@ -369,7 +369,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
chunks_after_release = r.chunks_done
for _ in range(4):
yield
@@ -406,7 +406,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from run_until_finished(r2)
assert r2.finished, "resubmit under same rid must complete independently"
assert r1.kv_pages == 0, "aborted r1 must release KV before resubmit"
assert r1.req is None or r1.req.req_pool_idx is None
assert r1.req is None or r1.req.kv.req_pool_idx is None
assert r1.lock_refs == 0
def test_abort_during_gap_inflight_middle_chunks_positive(self):
@@ -429,7 +429,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert not r.is_chunking, "aborted gap req must not re-enter chunking"
yield
@@ -511,7 +511,7 @@ class TestAbortBasic(ScriptedTestCase):
assert r1.kv_pages == 0, (
f"force_retract + abort same yield must release KV; got " f"{r1.kv_pages}"
)
assert r1.req is None or r1.req.req_pool_idx is None, (
assert r1.req is None or r1.req.kv.req_pool_idx is None, (
f"force_retract + abort same yield must release row; got " f"{r1.req}"
)
assert r1.lock_refs == 0, (
@@ -540,7 +540,7 @@ class TestAbortBasic(ScriptedTestCase):
yield from run_until(r2, lambda h: h.is_chunking)
assert r1.kv_pages == 0
assert r1.req is None or r1.req.req_pool_idx is None
assert r1.req is None or r1.req.kv.req_pool_idx is None
assert r1.lock_refs == 0
yield from run_until_finished(r2)
assert r2.finished, "baton handoff must let r2 complete"
@@ -570,7 +570,7 @@ class TestAbortPP(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
assert r.finished
test/manual/chunked_prefill/test_scripted_kv_pressure.py
+2 -2
View File
@@ -409,13 +409,13 @@ class TestKVPressureSmallPool(ScriptedTestCase):
if (
r_chunk.kv_pages == 0
and r_chunk.lock_refs == 0
and (r_chunk.req is None or r_chunk.req.req_pool_idx is None)
and (r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None)
):
break
yield
assert r_chunk.kv_pages == 0, f"kv_pages={r_chunk.kv_pages}"
assert r_chunk.lock_refs == 0, f"lock_refs={r_chunk.lock_refs}"
assert r_chunk.req is None or r_chunk.req.req_pool_idx is None
assert r_chunk.req is None or r_chunk.req.kv.req_pool_idx is None
t.abort(ballast)
for _ in range(200):
test/manual/chunked_prefill/test_scripted_lifecycle.py
+20 -12
View File
@@ -18,7 +18,7 @@ def _drain_until_released(t, *handles):
if all(
h.kv_pages == 0
and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None)
and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles
):
return
@@ -223,7 +223,7 @@ class TestLifecycleBasic(ScriptedTestCase):
r = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
yield from run_until_finished(r)
assert r.finished
assert r.req.req_pool_idx is None
assert r.req.kv.req_pool_idx is None
assert r.kv_pages == 0
assert r.lock_refs == 0
@@ -235,7 +235,7 @@ class TestLifecycleBasic(ScriptedTestCase):
r1 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
yield from run_until_finished(r1)
yield from _drain_until_released(t, r1)
assert r1.req.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0
assert r1.req.kv.req_pool_idx is None and r1.kv_pages == 0 and r1.lock_refs == 0
r1_output_len = len(r1.req.output_ids)
r2 = t.start_req(prompt_len=16, max_new_tokens=2, ignore_eos=True)
@@ -243,7 +243,7 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from _drain_until_released(t, r2)
assert r1.finished and r2.finished
assert r1_output_len == 2 and len(r2.req.output_ids) == 2
assert r2.req.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0
assert r2.req.kv.req_pool_idx is None and r2.kv_pages == 0 and r2.lock_refs == 0
def test_five_seq_clean(self):
self.server.execute_script(self._script_five_seq_clean)
@@ -256,7 +256,7 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None
assert r.req.kv.req_pool_idx is None
assert r.kv_pages == 0
assert r.lock_refs == 0
reqs.append(r)
@@ -300,7 +300,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
if prompt == VERY_LONG_PROMPT_LEN:
assert r.chunks_done == 8
else:
@@ -319,7 +321,7 @@ class TestLifecycleBasic(ScriptedTestCase):
assert r.finished
assert len(r.req.output_ids) == 1
yield from _drain_until_released(t, r)
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 and r.lock_refs == 0
if L > DEFAULT_CHUNK_SIZE:
assert (
@@ -341,7 +343,7 @@ class TestLifecycleBasic(ScriptedTestCase):
assert r.finished
assert len(r.req.output_ids) == 1
yield from _drain_until_released(t, r)
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.kv_pages == 0 and r.lock_refs == 0
if L > DEFAULT_CHUNK_SIZE:
assert (
@@ -360,7 +362,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
for _ in range(20):
yield
@@ -377,7 +381,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
if L == VERY_LONG_PROMPT_LEN:
assert r.chunks_done == 8
else:
@@ -394,7 +400,9 @@ class TestLifecycleBasic(ScriptedTestCase):
yield from run_until_finished(r)
assert r.finished
assert len(r.req.output_ids) == 2
assert r.req.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
assert (
r.req.kv.req_pool_idx is None and r.kv_pages == 0 and r.lock_refs == 0
)
for _ in range(5):
yield
t.flush_cache()
@@ -432,7 +440,7 @@ class TestLifecycleBasic(ScriptedTestCase):
)
assert r.finished or _error_message(r) is not None
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
test/manual/chunked_prefill/test_scripted_lora.py
+3 -1
View File
@@ -222,7 +222,9 @@ class TestLoRAAdapterEviction(ScriptedTestCase):
t.abort(r_a)
for _ in range(12):
if r_a.kv_pages == 0 and (r_a.req is None or r_a.req.req_pool_idx is None):
if r_a.kv_pages == 0 and (
r_a.req is None or r_a.req.kv.req_pool_idx is None
):
break
yield
test/manual/chunked_prefill/test_scripted_multi_req.py
+1 -1
View File
@@ -173,7 +173,7 @@ class TestMultiReqBasic(ScriptedTestCase):
for _ in range(5):
yield
assert r.kv_pages == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
def test_rid_reuse_after_finish(self):
self.server.execute_script(self._script_rid_reuse_after_finish)
test/manual/chunked_prefill/test_scripted_pp.py
+1 -1
View File
@@ -31,7 +31,7 @@ def _expected_chunks(prompt_len: int, chunk_size: int) -> int:
def _drain_until_released(t, *handles):
for _ in range(16):
if all(
h.kv_pages == 0 and (h.req is None or h.req.req_pool_idx is None)
h.kv_pages == 0 and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles
):
return
test/manual/chunked_prefill/test_scripted_priority.py
+2 -2
View File
@@ -168,13 +168,13 @@ class TestPriorityBasic(ScriptedTestCase):
if (
r.kv_pages == 0
and r.lock_refs == 0
and (r.req is None or r.req.req_pool_idx is None)
and (r.req is None or r.req.kv.req_pool_idx is None)
):
break
yield
assert r.kv_pages == 0
assert r.lock_refs == 0
assert r.req is None or r.req.req_pool_idx is None
assert r.req is None or r.req.kv.req_pool_idx is None
t.continue_generation()
yield
assert r.kv_pages == 0 and r.lock_refs == 0
test/manual/chunked_prefill/test_scripted_regression.py
+9 -9
View File
@@ -19,7 +19,7 @@ def _drain_until_released(t, *handles):
if all(
h.kv_pages == 0
and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None)
and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles
):
return
@@ -41,7 +41,7 @@ class TestRegressionBasic(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert r.kv_pages == 0
assert r.req.req_pool_idx is None
assert r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
assert not r.is_chunking
assert r.req.inflight_middle_chunks == 0
@@ -58,7 +58,7 @@ class TestRegressionBasic(ScriptedTestCase):
yield
assert r.kv_pages == 0
assert r.req.req_pool_idx is None
assert r.req.kv.req_pool_idx is None
assert r.lock_refs == 0
assert not r.is_chunking
@@ -240,7 +240,7 @@ class TestRegressionBasic(ScriptedTestCase):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
assert r.req.req_pool_idx is not None, "row must be held mid-chunk"
assert r.req.kv.req_pool_idx is not None, "row must be held mid-chunk"
assert r.kv_pages > 0, "committed KV must be held mid-chunk"
assert r.lock_refs >= 1, "radix lock_ref must be held mid-chunk"
@@ -248,8 +248,8 @@ class TestRegressionBasic(ScriptedTestCase):
yield from _drain_until_released(t, r)
assert (
r.req.req_pool_idx is None
), f"96d4749094: abort must release row; got row_idx={r.req.req_pool_idx!r}"
r.req.kv.req_pool_idx is None
), f"96d4749094: abort must release row; got row_idx={r.req.kv.req_pool_idx!r}"
assert (
r.kv_pages == 0
), f"96d4749094: abort must release KV; got kv_pages={r.kv_pages}"
@@ -269,14 +269,14 @@ class TestRegressionBasic(ScriptedTestCase):
def _script_pause_retract_releases_waiting_chunked_resume(t: ScriptedContext):
r = t.start_req(prompt_len=VERY_LONG_PROMPT_LEN, max_new_tokens=2)
yield from run_until(r, lambda h: h.is_chunking and h.chunks_done >= 1)
assert r.req.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1
assert r.req.kv.req_pool_idx is not None and r.kv_pages > 0 and r.lock_refs >= 1
t.pause_generation(mode="retract")
yield
assert r.req.req_pool_idx is None, (
assert r.req.kv.req_pool_idx is None, (
f"f38e69f87d: pause(retract) must release waiting "
f"chunked-resume row; got row_idx={r.req.req_pool_idx!r}"
f"chunked-resume row; got row_idx={r.req.kv.req_pool_idx!r}"
)
assert r.kv_pages == 0
assert r.lock_refs == 0
test/manual/scheduler/test_scripted_pp_abort.py
+1 -1
View File
@@ -14,7 +14,7 @@ def _drain_until_released(t: ScriptedContext, *handles: ScriptedReqHandle):
if all(
h.kv_pages == 0
and h.lock_refs == 0
and (h.req is None or h.req.req_pool_idx is None)
and (h.req is None or h.req.kv.req_pool_idx is None)
for h in handles
):
return
test/registered/unit/disaggregation/test_disaggregation_wire.py
+1 -1
View File
@@ -102,7 +102,7 @@ class TestDisaggregationWire(unittest.TestCase):
def test_prebuilt_skips_unused_prompt_tensor(self):
req = SimpleNamespace(
req_pool_idx=0,
kv=SimpleNamespace(req_pool_idx=0),
prefix_indices=[0, 1],
extend_range=SimpleNamespace(length=3),
origin_input_ids=[0, 1, 2, 3, 4],
test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py
+3 -2
View File
@@ -33,9 +33,10 @@ def _make_mock_req(
"""Create a mock Req with the KV cache state needed for testing."""
req = MagicMock()
req.rid = rid
req.req_pool_idx = req_pool_idx
req.kv = ReqKvInfo(
kv_committed_len=kv_committed_len, kv_allocated_len=kv_allocated_len
req_pool_idx=req_pool_idx,
kv_committed_len=kv_committed_len,
kv_allocated_len=kv_allocated_len,
)
req.prefix_indices = list(range(prefix_indices_len))
req.effective_kv_committed_len = lambda: req.kv.kv_committed_len
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+6 -6
View File
@@ -719,7 +719,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=[11, 12, 13],
prefix_slot_ids=[2, 3],
new_slot_ids=[4],
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
)
runner.eval_pending(pending)
runner.prefill_finalize(pending)
@@ -775,7 +775,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=token_ids,
prefix_slot_ids=[],
new_slot_ids=list(range(1, 71)),
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
req=req,
)
runner.eval_pending(pending)
@@ -837,7 +837,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
full_token_ids=token_ids,
prefix_slot_ids=list(range(1, 65)),
new_slot_ids=list(range(65, 258)),
req_pool_idx=req.req_pool_idx,
req_pool_idx=req.kv.req_pool_idx,
req=req,
)
runner.eval_pending(pending)
@@ -917,7 +917,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
req = FakeRequest()
req_indices = pool.alloc([req])
auxiliary_state_idx = pool.get_auxiliary_state_indices(req.req_pool_idx)
auxiliary_state_idx = pool.get_auxiliary_state_indices(req.kv.req_pool_idx)
pool.free(req)
# Which free slot a fresh alloc gets is not semantically meaningful
@@ -925,7 +925,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertEqual(len(req_indices), 1)
self.assertIn(req_indices[0], range(1, pool.size + 1))
self.assertIsNotNone(auxiliary_state_idx)
self.assertIsNone(req.req_pool_idx)
self.assertIsNone(req.kv.req_pool_idx)
self.assertIsNotNone(req.mamba_pool_idx)
self.assertIs(pool.mamba_allocator, pool.mamba_pool)
self.assertEqual(pool.auxiliary_state_pool.available_size(), 3)
@@ -1518,7 +1518,7 @@ if _HAS_MLX:
class FakeRequest:
def __init__(self):
self.req_pool_idx = None
self.kv = SimpleNamespace(req_pool_idx=None)
self.mamba_pool_idx = None
self.inflight_middle_chunks = 0
test/registered/unit/hardware_backend/mlx/test_max_running_requests.py
+1 -1
View File
@@ -119,8 +119,8 @@ def _hybrid_stub_for_initialize(
def _fake_req():
return SimpleNamespace(
req_pool_idx=None,
inflight_middle_chunks=0,
kv=SimpleNamespace(req_pool_idx=None),
mamba_pool_idx=None,
mamba_ping_pong_track_buffer=None,
)
test/registered/unit/hardware_backend/mlx/test_tp_worker_routing.py
+1 -1
View File
@@ -169,7 +169,7 @@ class _FakeReq:
self.rid = rid
self.prefix_indices = torch.empty(0, dtype=torch.long)
self.fill_ids = [0]
self.req_pool_idx = req_pool_idx
self.kv = SimpleNamespace(req_pool_idx=req_pool_idx)
# Mirrors Req's chunk-finality contract read by
# MlxTpModelWorker._chunk_needs_logits: extend_range=None means
# "not truncated" (final chunk / plain prefill).
test/registered/unit/layers/test_minicpm_sparse_cache.py
+4 -4
View File
@@ -7,6 +7,7 @@ import torch
from sglang.srt.layers.attention.minicpm.cache import (
attach_compressed_cache,
)
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.srt.managers.scheduler_components.invariant_checker import (
SchedulerInvariantChecker,
)
@@ -66,8 +67,8 @@ def make_pool_and_req(capacity: int = 64):
enable_memory_saver=False,
)
req = SimpleNamespace(
req_pool_idx=None,
inflight_middle_chunks=0,
kv=SimpleNamespace(req_pool_idx=None),
)
req_pool_idx = pool.alloc([req])[0]
return pool, req, req_pool_idx, allocator
@@ -224,8 +225,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, cache_protected_len=0),
kv=ReqKvInfo(req_pool_idx=req_pool_idx, kv_allocated_len=16),
)
session.release_session("session-a")
@@ -293,7 +293,7 @@ def test_partial_failure_rolls_back_and_free_releases_every_slot():
assert len(cache.free_slots) == 0
pool.free(req)
assert req.req_pool_idx is None
assert req.kv.req_pool_idx is None
assert allocator.available_size() == 12
assert len(cache.free_slots) == 8
test/registered/unit/managers/test_hisparse_unit.py
+26 -25
View File
@@ -49,8 +49,7 @@ def _make_req(rid="test-req-0", origin_input_ids=None, output_ids=None):
output_ids=output_ids,
fill_ids=origin_input_ids + output_ids,
seqlen=len(origin_input_ids) + len(output_ids),
req_pool_idx=None,
kv=SimpleNamespace(kv_allocated_len=0, kv_committed_len=0),
kv=SimpleNamespace(req_pool_idx=None, kv_allocated_len=0, kv_committed_len=0),
finished_reason=None,
hisparse_staging=False,
staging=False,
@@ -190,11 +189,11 @@ class TestHiSparseUnit(unittest.TestCase):
"""Allocate a req_pool_idx for the request."""
indices = self.req_to_token_pool.alloc([req])
self.assertIsNotNone(indices, "Failed to allocate req pool slot")
return req.req_pool_idx
return req.kv.req_pool_idx
def _free_req_slot(self, req):
"""Free the req_pool_idx."""
if req.req_pool_idx is not None:
if req.kv.req_pool_idx is not None:
self.req_to_token_pool.free(req)
def _alloc_kv(self, req, fill_len, *, logical_only=False):
@@ -216,7 +215,9 @@ class TestHiSparseUnit(unittest.TestCase):
extend_num_tokens=fill_len,
)
self.assertIsNotNone(kv_loc, "KV alloc failed")
self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
self.req_to_token_pool.write(
(req.kv.req_pool_idx, slice(0, len(kv_loc))), kv_loc
)
req.kv.kv_allocated_len = fill_len
req.kv.kv_committed_len = fill_len
req.full_untruncated_fill_ids = array("q", range(fill_len))
@@ -254,8 +255,8 @@ class TestHiSparseUnit(unittest.TestCase):
host_indices = host_pool.alloc(fill_len)
self.assertIsNotNone(host_indices, "Host alloc failed")
host_indices = host_indices.to(device="cuda")
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len] = host_indices
self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx] = fill_len
self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len] = host_indices
self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx] = fill_len
for lid in range(LAYER_NUM):
for i in range(fill_len):
host_pool.kv_buffer[lid][host_indices[i]] = self._kv_pattern(lid, i)
@@ -291,7 +292,7 @@ class TestHiSparseUnit(unittest.TestCase):
def _make_batch_tensors(self, reqs, fill_lens):
"""Build (req_pool_indices [int64], seq_lens [int32]) on cuda."""
rpi = torch.tensor(
[r.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda"
[r.kv.req_pool_idx for r in reqs], dtype=torch.int64, device="cuda"
)
sls = torch.tensor(fill_lens, dtype=torch.int32, device="cuda")
return rpi, sls
@@ -566,7 +567,7 @@ class TestHiSparseUnit(unittest.TestCase):
kv_loc = self._alloc_kv(req, fill_len)
self.coordinator.alloc_device_buffer(req)
self.coordinator._skip_first_backup[req.req_pool_idx] = True
self.coordinator._skip_first_backup[req.kv.req_pool_idx] = True
out_loc = self.allocator.alloc(1)
self.assertIsNotNone(out_loc)
@@ -576,7 +577,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertTrue(torch.all(stale_loc > 0), "Temporary mapping should exist")
seq_len = fill_len + 1
self.req_to_token_pool.write((req.req_pool_idx, fill_len), out_loc)
self.req_to_token_pool.write((req.kv.req_pool_idx, fill_len), out_loc)
req.kv.kv_allocated_len = seq_len
req.kv.kv_committed_len = seq_len
@@ -584,10 +585,10 @@ class TestHiSparseUnit(unittest.TestCase):
seq_lens=torch.tensor([seq_len], dtype=torch.int64, device=device),
out_cache_loc=out_loc,
req_pool_indices=torch.tensor(
[req.req_pool_idx], dtype=torch.int64, device=device
[req.kv.req_pool_idx], dtype=torch.int64, device=device
),
seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64),
req_pool_indices_cpu=torch.tensor([req.req_pool_idx], dtype=torch.int64),
req_pool_indices_cpu=torch.tensor([req.kv.req_pool_idx], dtype=torch.int64),
)
remapped_loc = self.allocator.full_to_hisparse_device_index_mapping[out_loc]
@@ -626,7 +627,7 @@ class TestHiSparseUnit(unittest.TestCase):
ready = self.coordinator.collect_ready_reqs()
self.assertEqual(len(ready), 1)
self.assertFalse(req.hisparse_staging)
self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx])
self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx])
tokens = self._build_topk_tokens(fill_len)
batch = tokens.unsqueeze(0)
@@ -662,11 +663,11 @@ class TestHiSparseUnit(unittest.TestCase):
ready = self.coordinator.collect_ready_reqs()
self.assertEqual(ready, [req])
host_row = self.coordinator.req_to_host_pool[req.req_pool_idx, :rounded_len]
host_row = self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :rounded_len]
self.assertTrue(torch.all(host_row >= 0))
self.assertEqual(torch.unique(host_row).numel(), rounded_len)
self.assertEqual(
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]),
int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
rounded_len,
)
@@ -674,7 +675,7 @@ class TestHiSparseUnit(unittest.TestCase):
next_host_index = self.coordinator.mem_pool_host.alloc_paged_token_slots(
self.coordinator.req_to_host_pool,
self.coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx,
req.kv.req_pool_idx,
fill_len,
1,
)
@@ -691,8 +692,8 @@ class TestHiSparseUnit(unittest.TestCase):
expected_total = rounded_len + expected_new_pages * self.page_size
allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices(
self.coordinator.req_to_host_pool,
req.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]),
req.kv.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
)
self.assertEqual(allocated_host_indices.numel(), expected_total)
@@ -714,9 +715,9 @@ class TestHiSparseUnit(unittest.TestCase):
self.coordinator.admit_request_direct(req)
self.assertFalse(req.staging)
self.assertTrue(self.coordinator._skip_first_backup[req.req_pool_idx])
self.assertTrue(self.coordinator._skip_first_backup[req.kv.req_pool_idx])
buf_tokens = self.coordinator.req_device_buffer_tokens[
:, req.req_pool_idx, :DEVICE_BUFFER_SIZE
:, req.kv.req_pool_idx, :DEVICE_BUFFER_SIZE
]
self.assertTrue(torch.all(buf_tokens == -1))
@@ -765,7 +766,7 @@ class TestHiSparseUnit(unittest.TestCase):
self.assertTrue(
torch.equal(
host_indices,
self.coordinator.req_to_host_pool[req.req_pool_idx, :fill_len],
self.coordinator.req_to_host_pool[req.kv.req_pool_idx, :fill_len],
)
)
self.assertEqual(req.kv.kv_allocated_len, fill_len)
@@ -774,18 +775,18 @@ class TestHiSparseUnit(unittest.TestCase):
rounded_len = (fill_len + self.page_size - 1) // self.page_size * self.page_size
self.assertEqual(
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]),
int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
rounded_len,
)
allocated_host_indices = self.coordinator.mem_pool_host.allocated_host_indices(
self.coordinator.req_to_host_pool,
req.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.req_pool_idx]),
req.kv.req_pool_idx,
int(self.coordinator.req_to_host_pool_allocated_len[req.kv.req_pool_idx]),
)
self.assertEqual(allocated_host_indices.numel(), rounded_len)
kv_loc = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : req.kv.kv_allocated_len
req.kv.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
+19 -20
View File
@@ -5,6 +5,7 @@ from types import SimpleNamespace
import torch
from sglang.srt.managers.schedule_batch import ReqKvInfo
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -43,23 +44,15 @@ def _make_checker(page_size=_PAGE_SIZE, row_width=4096, num_reqs=8, free_pages=N
return _FakeChecker(), rtt, tc, alloc
class _FakeReq:
def __init__(self, rid, rpi, committed, allocated):
class _FakeOwner:
# A req or a session slot; the checker reads only `kv` (and `rid` for reqs).
def __init__(self, rpi, committed, allocated, rid=None):
self.rid = rid
self.req_pool_idx = rpi
self.kv = SimpleNamespace(
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0
self.kv = ReqKvInfo(
req_pool_idx=rpi,
kv_committed_len=committed,
kv_allocated_len=allocated,
)
self.is_holding_kv = True
class _FakeSlot:
def __init__(self, rpi, committed, allocated):
self.req_pool_idx = rpi
self.kv = SimpleNamespace(
kv_committed_len=committed, kv_allocated_len=allocated, swa_evicted_seqlen=0
)
self.is_holding_kv = True
class TestKVPageInvariants(CustomTestCase):
@@ -70,21 +63,23 @@ class TestKVPageInvariants(CustomTestCase):
rtt[0, :256] = torch.arange(_PAGE_SIZE) # req 0 owns page 0
rtt[1, :256] = torch.arange(_PAGE_SIZE, 2 * _PAGE_SIZE) # req 1 owns page 1
chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeReq("a", 0, 256, 256), _FakeReq("b", 1, 200, 256)]
reqs=[_FakeOwner(0, 256, 256, rid="a"), _FakeOwner(1, 200, 256, rid="b")]
)
chk._check_kv_page_invariants()
self.assertEqual(chk.count_memory_leak_warnings, 0)
def test_committed_gt_allocated_raises(self):
chk, rtt, tc, alloc = _make_checker()
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 145, 144)])
chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 145, 144, rid="a")]
)
with self.assertRaises(AssertionError):
chk._check_kv_page_invariants()
def test_slot_committed_gt_allocated_raises(self):
chk, rtt, tc, alloc = _make_checker()
chk.get_last_batch = lambda: None
tc.slots = {"s1": _FakeSlot(0, 145, 144)}
tc.slots = {"s1": _FakeOwner(0, 145, 144)}
with self.assertRaises(AssertionError):
chk._check_kv_page_invariants()
@@ -94,14 +89,18 @@ class TestKVPageInvariants(CustomTestCase):
rtt[0, :3] = torch.tensor(
[5 * _PAGE_SIZE, 5 * _PAGE_SIZE + 1, 5 * _PAGE_SIZE + 2]
)
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 3, 3)])
chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 3, 3, rid="a")]
)
with self.assertRaises(ValueError):
chk._check_kv_page_invariants()
def test_free_pool_duplicate_raises(self):
chk, rtt, tc, alloc = _make_checker(free_pages=torch.tensor([3, 3, 4]))
rtt[0, :1] = torch.tensor([10 * _PAGE_SIZE]) # owner page 10, not in free
chk.get_last_batch = lambda: SimpleNamespace(reqs=[_FakeReq("a", 0, 1, 1)])
chk.get_last_batch = lambda: SimpleNamespace(
reqs=[_FakeOwner(0, 1, 1, rid="a")]
)
with self.assertRaises(ValueError):
chk._check_kv_page_invariants()
test/registered/unit/managers/test_priority_scheduling_disaggregation.py
+1 -2
View File
@@ -93,11 +93,10 @@ class TestDecodePreallocQueuePriority(unittest.TestCase):
priority=priority,
origin_input_ids=[1, 2, 3],
output_ids=[],
req_pool_idx=int(priority) % 8,
finished_reason=FINISH_ABORT("failed") if failed else None,
return_logprob=False,
sampling_params=SimpleNamespace(max_new_tokens=8),
kv=SimpleNamespace(cache_protected_len=0),
kv=SimpleNamespace(req_pool_idx=int(priority) % 8, cache_protected_len=0),
time_stats=MagicMock(),
)
return SimpleNamespace(
test/registered/unit/managers/test_schedule_batch_req_pool_indices.py
+1 -1
View File
@@ -18,7 +18,7 @@ register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _make_req(req_pool_idx, origin_input_ids, output_ids):
return types.SimpleNamespace(
req_pool_idx=req_pool_idx,
kv=types.SimpleNamespace(req_pool_idx=req_pool_idx),
origin_input_ids=origin_input_ids,
output_ids=output_ids,
return_logprob=False,
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+1 -2
View File
@@ -34,11 +34,10 @@ def _make_req(
req.output_ids = array("q")
req.full_untruncated_fill_ids = array("q", fill_ids)
req.prefix_indices = prefix_indices
req.req_pool_idx = req_pool_idx
req.extend_range = Range(fill_len - extend_input_len, fill_len)
req.inflight_middle_chunks = 0
req.host_hit_length = 0
req.kv = ReqKvInfo()
req.kv = ReqKvInfo(req_pool_idx=req_pool_idx)
req.skip_radix_cache_insert = False
req.last_node = None
req.swa_uuid_for_lock = None
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1 -1
View File
@@ -75,13 +75,13 @@ class MockReq:
"q", fill_ids[:-1] if len(fill_ids) > 1 else fill_ids
)
self.output_ids = array("q", [fill_ids[-1]] if len(fill_ids) > 1 else [])
self.req_pool_idx = req_pool_idx
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 = SimpleNamespace(
req_pool_idx=req_pool_idx,
kv_committed_len=len(fill_ids),
kv_allocated_len=len(fill_ids),
cache_protected_len=cache_protected_len,
test/registered/unit/mem_cache/test_decode_retraction_backup.py
+5 -3
View File
@@ -129,12 +129,14 @@ class TestDecodeRetractionBackup(unittest.TestCase):
)
def _admit_req(self, env, num_tokens: int):
req = SimpleNamespace(rid="request", req_pool_idx=None, seqlen=num_tokens + 1)
req = SimpleNamespace(
rid="request", kv=SimpleNamespace(req_pool_idx=None), seqlen=num_tokens + 1
)
self.assertIsNotNone(env.req_to_token_pool.alloc([req]))
source_indices = env.allocator.alloc(num_tokens)
self.assertIsNotNone(source_indices)
env.req_to_token_pool.write(
(req.req_pool_idx, slice(0, num_tokens)), source_indices
(req.kv.req_pool_idx, slice(0, num_tokens)), source_indices
)
return req, source_indices
@@ -200,7 +202,7 @@ class TestDecodeRetractionBackup(unittest.TestCase):
self.assertIsNotNone(destination_indices)
self.assertFalse(torch.equal(source_indices, destination_indices))
req_to_token_pool.write(
(req.req_pool_idx, slice(0, self.num_tokens)), destination_indices
(req.kv.req_pool_idx, slice(0, self.num_tokens)), destination_indices
)
cache.retraction_restore(req, backup)
test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py
+5 -5
View File
@@ -61,10 +61,10 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
return SimpleNamespace(
rid=rid,
prefix_indices=torch.tensor(prefix, dtype=torch.int32),
req_pool_idx=req_pool_idx,
dllm_incomplete_ids=array("q", range(block_size)) if reuse else array("q"),
inflight_middle_chunks=1 if req_pool_idx is not None else 0,
kv=SimpleNamespace(
req_pool_idx=req_pool_idx,
kv_committed_len=len(prefix) if req_pool_idx is not None else 0,
kv_allocated_len=(
len(prefix) + block_size if req_pool_idx is not None else 0
@@ -75,8 +75,8 @@ def _make_req(rid, prefix, block_size, *, req_pool_idx=None, reuse=False):
def _remove_allocated_req_slots(pool, *reqs):
for req in reqs:
if req.req_pool_idx in pool.free_slots:
pool.free_slots.remove(req.req_pool_idx)
if req.kv.req_pool_idx in pool.free_slots:
pool.free_slots.remove(req.kv.req_pool_idx)
def _make_batch(pool, allocator, reqs, extend_lens):
@@ -105,8 +105,8 @@ def _make_batch(pool, allocator, reqs, extend_lens):
def _seed_retained_block(pool, req, values):
prefix_len = len(req.prefix_indices)
pool.req_to_token[req.req_pool_idx, :prefix_len] = req.prefix_indices
pool.req_to_token[req.req_pool_idx, prefix_len : prefix_len + len(values)] = (
pool.req_to_token[req.kv.req_pool_idx, :prefix_len] = req.prefix_indices
pool.req_to_token[req.kv.req_pool_idx, prefix_len : prefix_len + len(values)] = (
torch.tensor(values, dtype=torch.int32)
)
test/registered/unit/mem_cache/test_hisparse_allocator.py
+2 -2
View File
@@ -106,7 +106,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
def alloc(self, reqs):
for item in reqs:
item.req_pool_idx = 0
item.kv.req_pool_idx = 0
return torch.tensor([0], dtype=torch.int64)
def write(self, indices, values):
@@ -160,7 +160,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
regular_host_alloc.assert_called_once_with(
coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx,
req.kv.req_pool_idx,
0,
len(host_indices),
)
test/registered/unit/mem_cache/test_paged_free_segment.py
+1 -1
View File
@@ -131,7 +131,7 @@ class TestFreeSegment(unittest.TestCase):
token_to_kv_pool_allocator=alloc,
req_to_token_pool=SimpleNamespace(req_to_token=row.unsqueeze(0)),
)
req = SimpleNamespace(req_pool_idx=0)
req = SimpleNamespace(kv=SimpleNamespace(req_pool_idx=0))
before = len(alloc.free_pages)
alloc.free_group_begin()
test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py
+4 -3
View File
@@ -21,11 +21,12 @@ class _FakeAllocator:
class _FakeReq:
req_pool_idx = 0
def __init__(self):
self.kv = SimpleNamespace(
swa_evicted_seqlen=6, swa_evict_floor=3, cache_protected_len=0
req_pool_idx=0,
swa_evicted_seqlen=6,
swa_evict_floor=3,
cache_protected_len=0,
)
def pop_committed_kv_cache(self):
test/registered/unit/mem_cache/test_radix_cache_unit.py
+1 -2
View File
@@ -518,8 +518,7 @@ class TestRadixCache(unittest.TestCase):
)
cache.req_to_token_pool = ReqToTokenPool(request_indices.clone())
req = unittest.mock.Mock(
req_pool_idx=0,
kv=SimpleNamespace(cache_protected_len=0),
kv=SimpleNamespace(req_pool_idx=0, cache_protected_len=0),
extra_key=None,
cache_salt=None,
priority=0,
test/registered/unit/mem_cache/test_retraction_mamba_backup.py
+2 -2
View File
@@ -2,7 +2,7 @@ import unittest
import torch
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_batch import Req, ReqKvInfo
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.test.ci.ci_register import register_cpu_ci
@@ -39,7 +39,7 @@ class _Allocator:
def _req_and_pool():
req = object.__new__(Req)
req.req_pool_idx = 0
req.kv = ReqKvInfo(req_pool_idx=0)
req.origin_input_ids = [1, 2]
req.output_ids = [3]
req.mamba_pool_idx = torch.tensor(1)
test/registered/unit/mem_cache/test_streaming_session_unit.py
+14 -14
View File
@@ -2,7 +2,7 @@ from types import SimpleNamespace
import torch
from sglang.srt.managers.schedule_batch import FINISH_ABORT
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
from sglang.srt.mem_cache.base_prefix_cache import MatchResult
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
from sglang.test.ci.ci_register import register_cpu_ci
@@ -24,8 +24,8 @@ class _FakeReqToTokenPool:
self.free_slots = []
def free(self, req):
self.free_slots.append(req.req_pool_idx)
req.req_pool_idx = None
self.free_slots.append(req.kv.req_pool_idx)
req.kv.req_pool_idx = None
class _FakeInnerCache:
@@ -67,8 +67,8 @@ class _FakeReq:
abort_req=lambda: None,
_inflight=False,
)
self.req_pool_idx = req_pool_idx
self.kv = SimpleNamespace(
self.kv = ReqKvInfo(
req_pool_idx=req_pool_idx,
kv_committed_len=committed,
kv_allocated_len=allocated,
swa_evicted_seqlen=0,
@@ -112,8 +112,8 @@ def test_preabort_detaches_session_and_preserves_slot():
)
tree_cache = StreamingSession(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv=SimpleNamespace(
kv=ReqKvInfo(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
swa_evicted_seqlen=0,
@@ -135,7 +135,7 @@ def test_preabort_detaches_session_and_preserves_slot():
assert req.session is None
# Slot untouched.
slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0
assert slot.kv.req_pool_idx == 0
assert slot.kv.kv_committed_len == 48
assert slot.kv.kv_allocated_len == 48
assert len(result.device_indices) == 0
@@ -160,7 +160,7 @@ def test_first_mid_abort_nukes_ephemeral_slot():
# Slot must NOT be created.
assert "session-a" not in tree_cache.slots
# Transient pool slot freed.
assert req.req_pool_idx is None
assert req.kv.req_pool_idx is None
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(20))
@@ -179,8 +179,8 @@ def test_nth_mid_abort_nukes_session_slot():
# Session already has a slot from a previous turn.
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv=SimpleNamespace(
kv=ReqKvInfo(
req_pool_idx=0,
kv_committed_len=50,
kv_allocated_len=50,
swa_evicted_seqlen=0,
@@ -202,7 +202,7 @@ def test_nth_mid_abort_nukes_session_slot():
assert allocator.freed[0].tolist() == list(range(65))
# Pool slot returned.
assert req_to_token_pool.free_slots == [0]
assert req.req_pool_idx is None
assert req.kv.req_pool_idx is None
def test_release_session_threads_mamba_skip_ids():
@@ -220,8 +220,8 @@ def test_release_session_threads_mamba_skip_ids():
lock_node = SimpleNamespace(id=42)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv=SimpleNamespace(
kv=ReqKvInfo(
req_pool_idx=0,
kv_committed_len=50,
kv_allocated_len=50,
swa_evicted_seqlen=0,
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+3 -3
View File
@@ -103,11 +103,11 @@ def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=51
def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
"""Mock Req with fields needed by _evict_swa and cache_finished_req."""
req = SimpleNamespace(
req_pool_idx=req_pool_idx,
is_holding_kv=True,
origin_input_ids=token_ids,
output_ids=[],
kv=ReqKvInfo(cache_protected_len=cache_protected_len),
kv=ReqKvInfo(
req_pool_idx=req_pool_idx, cache_protected_len=cache_protected_len
),
extra_key=None,
cache_salt=None,
last_node=tree.root_node,
test/registered/unit/mem_cache/test_swa_unittest.py
+5 -5
View File
@@ -669,13 +669,13 @@ class TestSWA(unittest.TestCase):
# Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = _DummyReq()
req.req_pool_idx = 0
req.kv.req_pool_idx = 0
req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6])
req.output_ids = array("q")
req._kv_committed_len = len(req.origin_input_ids)
kv_indices = allocator.alloc(req._kv_committed_len)
req_to_token_pool.write(
(req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
(req.kv.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
)
req.extra_key = None
req.cache_salt = None
@@ -707,13 +707,13 @@ class TestSWA(unittest.TestCase):
# Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
# even when len(prefix_indices) is intentionally larger.
req2 = _DummyReq()
req2.req_pool_idx = 1
req2.kv.req_pool_idx = 1
req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16])
req2.output_ids = array("q")
req2._kv_committed_len = len(req2.origin_input_ids)
kv_indices2 = allocator.alloc(req2._kv_committed_len)
req_to_token_pool.write(
(req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
(req2.kv.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
)
req2.extra_key = None
req2.cache_salt = None
@@ -911,7 +911,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
token_ids = array("q", range(1, num_tokens + 1))
req = _DummyReq()
req.req_pool_idx = 0
req.kv.req_pool_idx = 0
req.origin_input_ids = token_ids
req.output_ids = array("q")
req.get_fill_ids = lambda: token_ids
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+1 -1
View File
@@ -646,7 +646,7 @@ def bench_cache_finished(
req.kv.kv_committed_len = len(seq)
if hasattr(lr, "swa_uuid_for_lock"):
req.swa_uuid_for_lock = lr.swa_uuid_for_lock
env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices
env.rtp.req_to_token[req.kv.req_pool_idx, : len(kv_indices)] = kv_indices
req_items.append(req)
if not req_items:
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+16 -12
View File
@@ -1235,7 +1235,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q", output_ids)
kv_len = len(input_ids) + len(output_ids)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -1275,7 +1275,7 @@ class UnifiedRadixCacheSuite:
)
kv_len = req.extend_range.end
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len
req.kv.kv_allocated_len = kv_len
req.last_node = cache.root_node_handle()
@@ -1297,7 +1297,7 @@ class UnifiedRadixCacheSuite:
start_p = ((start_p + ps - 1) // ps) * ps
if start_p < end_p:
allocator.free(
req_to_token_pool.req_to_token[req.req_pool_idx][start_p:end_p]
req_to_token_pool.req_to_token[req.kv.req_pool_idx][start_p:end_p]
)
prompt_aligned = (len(prompt_ids) // ps) * ps
@@ -1320,7 +1320,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q")
kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -1355,7 +1355,7 @@ class UnifiedRadixCacheSuite:
)
kv_len = len(tokens)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -1392,7 +1392,9 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, len(tokens))
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
req_to_token_pool.write(
(req.kv.req_pool_idx, slice(0, len(tokens))), kv_indices
)
req.kv.kv_committed_len = len(tokens)
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -1494,7 +1496,7 @@ class UnifiedRadixCacheSuite:
req.output_ids = array("q")
kv_len = len(input_ids)
kv_indices = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), kv_indices)
req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -1619,7 +1621,7 @@ class UnifiedRadixCacheSuite:
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_len = len(tokens)
fresh_value = self._alloc(allocator, kv_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, kv_len)), fresh_value)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, kv_len)), fresh_value)
req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -2206,7 +2208,7 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -2295,7 +2297,7 @@ class UnifiedRadixCacheSuite:
req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc(allocator, pre_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, pre_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, pre_len)), kv_indices)
req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
@@ -6790,7 +6792,9 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
req = self._make_req(req_to_token_pool, tokens)
kv_indices = allocator.alloc(len(tokens))
self.assertIsNotNone(kv_indices)
req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
req_to_token_pool.write(
(req.kv.req_pool_idx, slice(0, len(tokens))), kv_indices
)
req.last_node = cache.root_node_handle()
cache.cache_finished_req(
@@ -8104,7 +8108,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
req.full_untruncated_fill_ids = array("q", tokens)
req.set_extend_range(0, len(req.full_untruncated_fill_ids))
kv_indices = self._alloc_paged(allocator, seq_len)
req_to_token_pool.write((req.req_pool_idx, slice(0, seq_len)), kv_indices)
req_to_token_pool.write((req.kv.req_pool_idx, slice(0, seq_len)), kv_indices)
req.kv.kv_committed_len = seq_len
req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0
test/registered/xpu/test_lmcache_radix_cache.py
+10 -3
View File
@@ -66,13 +66,16 @@ def _make_req(rid, req_pool_idx, token_ids, tree):
SimpleNamespace pattern in test_swa_eviction_boundary.py)."""
req = SimpleNamespace(
rid=rid,
req_pool_idx=req_pool_idx,
origin_input_ids=token_ids,
output_ids=[],
extra_key=None,
cache_salt=None,
last_node=tree.root_node,
kv=SimpleNamespace(cache_protected_len=0, kv_committed_len=len(token_ids)),
kv=SimpleNamespace(
req_pool_idx=req_pool_idx,
cache_protected_len=0,
kv_committed_len=len(token_ids),
),
priority=0,
kv_committed_freed=False,
)
@@ -162,7 +165,11 @@ class TestLMCRadixCacheXPU(unittest.TestCase):
# commit it as a finished request (inserts into radix + stores to
# LMCache on tree.store_stream).
req_pool_idx = req_to_token_pool.alloc(
[SimpleNamespace(req_pool_idx=None, inflight_middle_chunks=0)]
[
SimpleNamespace(
inflight_middle_chunks=0, kv=SimpleNamespace(req_pool_idx=None)
)
]
)[0]
kv_slots = allocator.alloc(self.INPUT_LEN)
self.assertIsNotNone(kv_slots)