393 lines
14 KiB
Python
393 lines
14 KiB
Python
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
|
|
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
|
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams, MatchResult
|
|
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=12, suite="base-a-test-cpu")
|
|
|
|
|
|
class _FakeAllocator(BaseTokenToKVPoolAllocator):
|
|
"""Single-pool double. Subclassing the base routes free_full / free_segment /
|
|
free_segments into free(), so a new free API cannot slip past the recorder."""
|
|
|
|
def __init__(self, page_size: int = 1):
|
|
super().__init__(
|
|
size=1024,
|
|
page_size=page_size,
|
|
dtype=torch.bfloat16,
|
|
device="cpu",
|
|
kvcache=None,
|
|
need_sort=False,
|
|
)
|
|
self.freed = []
|
|
|
|
def clear(self):
|
|
self.freed = []
|
|
|
|
def alloc(self, need_size: int):
|
|
raise NotImplementedError
|
|
|
|
def free(self, free_index: torch.Tensor):
|
|
self.freed.append(free_index.clone())
|
|
|
|
|
|
class _FakeReqToTokenPool:
|
|
def __init__(self, req_to_token):
|
|
self.req_to_token = req_to_token
|
|
self.free_slots = []
|
|
|
|
def free(self, req):
|
|
self.free_slots.append(req.kv.req_pool_idx)
|
|
req.kv.req_pool_idx = None
|
|
|
|
|
|
class _FakeInnerCache:
|
|
def __init__(self, req_to_token_pool, allocator, page_size, match_results=None):
|
|
self.req_to_token_pool = req_to_token_pool
|
|
self.token_to_kv_pool_allocator = allocator
|
|
self.page_size = page_size
|
|
self.match_results = list(match_results or [])
|
|
self.dec_lock_ref_calls = []
|
|
self.dec_lock_ref_params = []
|
|
self.dec_lock_ref_skip_swa = []
|
|
|
|
def cache_finished_req(self, *args, **kwargs):
|
|
raise AssertionError("Streaming requests should not delegate to inner cache")
|
|
|
|
def match_prefix(self, *args, **kwargs):
|
|
if not self.match_results:
|
|
raise AssertionError("Unexpected match_prefix call")
|
|
return self.match_results.pop(0)
|
|
|
|
def dec_lock_ref(self, node, *args, **kwargs):
|
|
self.dec_lock_ref_calls.append(node)
|
|
self.dec_lock_ref_params.append(args[0] if args else kwargs.get("params"))
|
|
self.dec_lock_ref_skip_swa.append(kwargs.get("skip_swa", False))
|
|
|
|
def supports_mamba(self):
|
|
return False
|
|
|
|
def sanity_check(self):
|
|
return None
|
|
|
|
|
|
class _FakeReq:
|
|
def __init__(
|
|
self, session_id: str, req_pool_idx: int, committed: int, allocated: int
|
|
):
|
|
self.session = SimpleNamespace(
|
|
session_id=session_id,
|
|
streaming=True,
|
|
finish_req=lambda req: None,
|
|
abort_req=lambda: None,
|
|
_inflight=False,
|
|
)
|
|
self.kv = ReqKvInfo(
|
|
req_pool_idx=req_pool_idx,
|
|
kv_committed_len=committed,
|
|
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.swa_branching_seqlen = None
|
|
self.lock_receipt = DecLockRefParams()
|
|
self.swa_prefix_lock_released = False
|
|
self.to_finish = None
|
|
self.finished_reason = None
|
|
self.finished_len = None
|
|
|
|
def detach_kv(self):
|
|
kv, self.kv = self.kv, ReqKvInfo()
|
|
return kv
|
|
|
|
|
|
def test_session_slot_round_trip_preserves_mamba_state():
|
|
# The mamba state rides in the shared ReqKvInfo record. mamba_branching_seqlen
|
|
# is a per-turn match observation on the Req and is not preserved by the slot.
|
|
req = _FakeReq("session-a", req_pool_idx=0, committed=4, allocated=4)
|
|
req.kv.mamba_next_track_idx = 1
|
|
req.kv.mamba_last_track_idx = 0
|
|
req.kv.mamba_last_track_seqlen = 3
|
|
|
|
slot = SessionSlot()
|
|
slot.save_from_req(req, is_first=True)
|
|
|
|
next_req = _FakeReq("session-a", req_pool_idx=1, committed=0, allocated=0)
|
|
slot.restore_to_req(next_req)
|
|
|
|
assert next_req.kv.mamba_next_track_idx == 1
|
|
assert next_req.kv.mamba_last_track_idx == 0
|
|
assert next_req.kv.mamba_last_track_seqlen == 3
|
|
|
|
|
|
def test_preabort_detaches_session_and_preserves_slot():
|
|
"""Pre-aborted req (to_finish set before match_prefix) is detached from
|
|
the session: session=None, abort_req() called. Slot stays intact."""
|
|
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator(page_size=16)
|
|
inner = _FakeInnerCache(
|
|
req_to_token_pool,
|
|
allocator,
|
|
page_size=16,
|
|
match_results=[
|
|
MatchResult(
|
|
device_indices=torch.tensor([], dtype=torch.int64),
|
|
last_device_node=None,
|
|
last_host_node=None,
|
|
best_match_node=None,
|
|
)
|
|
],
|
|
)
|
|
tree_cache = StreamingSession(inner)
|
|
tree_cache.slots["session-a"] = SessionSlot(
|
|
kv=ReqKvInfo(
|
|
req_pool_idx=0,
|
|
kv_committed_len=48,
|
|
kv_allocated_len=48,
|
|
swa_evicted_seqlen=0,
|
|
cache_protected_len=16,
|
|
),
|
|
)
|
|
|
|
req = _FakeReq("session-a", req_pool_idx=1, committed=1, allocated=1)
|
|
req.to_finish = FINISH_ABORT("too long")
|
|
|
|
result = tree_cache.match_prefix(
|
|
SimpleNamespace(
|
|
req=req,
|
|
key=SimpleNamespace(token_ids=list(range(64))),
|
|
)
|
|
)
|
|
|
|
# Req detached from session.
|
|
assert req.session is None
|
|
# Slot untouched.
|
|
slot = tree_cache.slots["session-a"]
|
|
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
|
|
|
|
|
|
def test_first_mid_abort_nukes_ephemeral_slot():
|
|
"""First-request mid-processing abort: no slot exists yet, ephemeral
|
|
slot is created from req state and nuked via release_session."""
|
|
page_size = 1
|
|
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator()
|
|
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
|
|
tree_cache = StreamingSession(inner)
|
|
|
|
# No slot exists yet (first request).
|
|
req = _FakeReq("session-a", req_pool_idx=0, committed=0, allocated=20)
|
|
req.finished_reason = FINISH_ABORT("input too long")
|
|
|
|
tree_cache.cache_finished_req(req)
|
|
|
|
# Slot must NOT be created.
|
|
assert "session-a" not in tree_cache.slots
|
|
# Transient pool slot freed.
|
|
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))
|
|
|
|
|
|
def test_nth_mid_abort_nukes_session_slot():
|
|
"""Nth-request mid-processing abort: slot exists, restore_to_req ran.
|
|
ALL KV is wiped (release_session). Slot is deleted. Token IDs stay
|
|
in req_nodes for next turn's re-prefill."""
|
|
page_size = 1
|
|
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator()
|
|
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
|
|
tree_cache = StreamingSession(inner)
|
|
|
|
# Mid-processing abort: restore_to_req ran, so the req runs on the slot's
|
|
# record, which this turn has grown to committed=60 / allocated=65.
|
|
req = _FakeReq("session-a", req_pool_idx=0, committed=60, allocated=65)
|
|
req.finished_reason = FINISH_ABORT("client disconnected")
|
|
tree_cache.slots["session-a"] = SessionSlot(kv=req.kv, last_node=None)
|
|
|
|
tree_cache.cache_finished_req(req)
|
|
|
|
# Slot wiped — deleted from slots dict.
|
|
assert "session-a" not in tree_cache.slots
|
|
# All KV freed: [0, 65) from release_session.
|
|
assert len(allocator.freed) == 1
|
|
assert allocator.freed[0].tolist() == list(range(65))
|
|
# Pool slot returned.
|
|
assert req_to_token_pool.free_slots == [0]
|
|
assert req.kv.req_pool_idx is None
|
|
|
|
|
|
def test_release_session_threads_mamba_lock_receipt():
|
|
"""release_session must forward the slot's mamba lock receipt to
|
|
dec_lock_ref. The first req's last_node may be full-only-locked (mamba
|
|
not taken at inc), so without the receipt the release would drop a mamba
|
|
lock the session never took -- another request's, on a shared node."""
|
|
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator()
|
|
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size=1)
|
|
tree_cache = StreamingSession(inner)
|
|
|
|
lock_node = SimpleNamespace(id=42)
|
|
tree_cache.slots["session-a"] = SessionSlot(
|
|
kv=ReqKvInfo(
|
|
req_pool_idx=0,
|
|
kv_committed_len=50,
|
|
kv_allocated_len=50,
|
|
swa_evicted_seqlen=0,
|
|
cache_protected_len=0,
|
|
),
|
|
last_node=lock_node,
|
|
)
|
|
|
|
tree_cache.release_session("session-a")
|
|
|
|
assert inner.dec_lock_ref_calls == [lock_node]
|
|
params = inner.dec_lock_ref_params[0]
|
|
assert params is not None
|
|
assert params.skipped_lock_components == ()
|
|
assert inner.dec_lock_ref_skip_swa == [False]
|
|
|
|
|
|
def test_release_session_skips_swa_after_early_release():
|
|
"""A slot saved from a req that early-released its SWA lock
|
|
(swa_prefix_lock_released) must release with skip_swa, or the session
|
|
close double-releases the SWA segment."""
|
|
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator()
|
|
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size=1)
|
|
tree_cache = StreamingSession(inner)
|
|
|
|
lock_node = SimpleNamespace(id=42)
|
|
tree_cache.slots["session-a"] = SessionSlot(
|
|
kv=ReqKvInfo(
|
|
req_pool_idx=0,
|
|
kv_committed_len=50,
|
|
kv_allocated_len=50,
|
|
swa_evicted_seqlen=0,
|
|
cache_protected_len=0,
|
|
),
|
|
last_node=lock_node,
|
|
lock_receipt=DecLockRefParams(node_id=42, swa_uuid_for_lock=7),
|
|
swa_prefix_lock_released=True,
|
|
)
|
|
|
|
tree_cache.release_session("session-a")
|
|
|
|
assert inner.dec_lock_ref_calls == [lock_node]
|
|
assert inner.dec_lock_ref_params[0].swa_uuid_for_lock == 7
|
|
assert inner.dec_lock_ref_skip_swa == [True]
|
|
|
|
|
|
def test_session_slot_does_not_restore_swa_branching_seqlen():
|
|
req = _FakeReq("session-a", req_pool_idx=0, committed=4, allocated=4)
|
|
req.swa_branching_seqlen = 8
|
|
|
|
slot = SessionSlot()
|
|
slot.save_from_req(req, is_first=True)
|
|
|
|
next_req = _FakeReq("session-a", req_pool_idx=1, committed=0, allocated=0)
|
|
slot.restore_to_req(next_req)
|
|
|
|
assert req.swa_branching_seqlen is None
|
|
assert next_req.swa_branching_seqlen is None
|
|
|
|
|
|
# Shrink tests removed: streaming sessions are append-only after the
|
|
# rollback fix in session_controller (rollback_aborted_req). The shrink
|
|
# code path in cache_finished_req no longer exists.
|
|
|
|
|
|
def test_trim_overshoot_postcondition():
|
|
"""`_trim_overshoot` postcondition: every per-req KV field is capped at
|
|
target = origin+finished_len, output_ids is truncated, and the tail
|
|
KV slots are freed. Covers both non-SWA fields (kv_committed_len,
|
|
kv_allocated_len, output_ids) and SWA bookkeeping (swa_evicted_seqlen)
|
|
in one shot — same invariant `_free_tail` enforces on the match_prefix
|
|
path.
|
|
"""
|
|
page_size = 1
|
|
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator()
|
|
tree_cache = StreamingSession(
|
|
_FakeInnerCache(req_to_token_pool, allocator, page_size)
|
|
)
|
|
|
|
# Overshoot scenario: origin=26, finished_len=12 -> target=38.
|
|
# committed=40 (overshoot 2), allocated=44, swa_evicted=42 (> target),
|
|
# output_ids extended to 14 by the overshoot round.
|
|
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.kv.swa_evicted_seqlen = 42
|
|
|
|
tree_cache._trim_overshoot(req, finished_len=12)
|
|
|
|
target = 38
|
|
assert req.kv.kv_committed_len == 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, split at the pre-trim eviction
|
|
# floor 42: [38, 42) gave its SWA peers back already, so it goes back full-only.
|
|
assert [t.tolist() for t in allocator.freed] == [[38, 39, 40, 41], [42, 43]]
|
|
|
|
|
|
def test_trim_overshoot_keeps_cursor_page_aligned_on_paged():
|
|
"""A mid-page trim target must not become the SWA eviction cursor (the
|
|
dead/alive split there frees the shared page twice); rewind to the boundary."""
|
|
page_size = 16
|
|
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
|
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
|
allocator = _FakeAllocator(page_size=page_size)
|
|
tree_cache = StreamingSession(
|
|
_FakeInnerCache(req_to_token_pool, allocator, page_size)
|
|
)
|
|
|
|
# origin=26, finished=12 -> raw target 38 (mid-page); cursor 48 > target.
|
|
req = _FakeReq("session-a", req_pool_idx=0, committed=52, allocated=64)
|
|
req.origin_input_ids = list(range(26))
|
|
req.output_ids = list(range(14))
|
|
req.kv.swa_evicted_seqlen = 48
|
|
|
|
tree_cache._trim_overshoot(req, finished_len=12)
|
|
|
|
# Rewound to floor_align(38) = 32; every cursor lands page-aligned.
|
|
assert req.kv.kv_allocated_len == 32
|
|
assert req.kv.kv_committed_len == 32
|
|
assert req.kv.swa_evicted_seqlen == 32
|
|
assert len(req.output_ids) == 12
|
|
# Freed [32, 64): [32, 48) below the old cursor goes back full-only,
|
|
# [48, 64) both halves.
|
|
assert [t.tolist() for t in allocator.freed] == [
|
|
list(range(32, 48)),
|
|
list(range(48, 64)),
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main([__file__, "-v"]))
|