From f9792166c3570175dea4862c814f58b85c7952c3 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 15 Apr 2026 15:05:35 -0700 Subject: [PATCH] trim_overshoot: cap swa_evicted_seqlen + unit test (#22900) --- .../srt/mem_cache/session_aware_cache.py | 1 + .../mem_cache/test_streaming_session_unit.py | 36 +++++++++++++++++++ 2 files changed, 37 insertions(+) diff --git a/python/sglang/srt/mem_cache/session_aware_cache.py b/python/sglang/srt/mem_cache/session_aware_cache.py index be56c2b0c..9d6aabc16 100644 --- a/python/sglang/srt/mem_cache/session_aware_cache.py +++ b/python/sglang/srt/mem_cache/session_aware_cache.py @@ -299,6 +299,7 @@ class SessionAwareCache(BasePrefixCache): self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len) req.kv_allocated_len = min(req.kv_allocated_len, target) req.kv_committed_len = min(req.kv_committed_len, target) + req.swa_evicted_seqlen = min(req.swa_evicted_seqlen, target) req.output_ids = req.output_ids[:finished_len] def _free_kv_aligned(self, pool_idx: int, target: int, end: int): diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 2f32d5555..0a701dcb5 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -235,3 +235,39 @@ def test_nth_mid_abort_nukes_session_slot(): # 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 = SimpleNamespace(req_to_token=req_to_token, free_slots=[]) + allocator = _FakeAllocator() + tree_cache = SessionAwareCache( + _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.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 len(req.output_ids) == 12 + # Tail [38, 44) freed by _free_kv_aligned. + assert len(allocator.freed) == 1 + assert allocator.freed[0].tolist() == list(range(38, 44))