trim_overshoot: cap swa_evicted_seqlen + unit test (#22900)
This commit is contained in:
@@ -299,6 +299,7 @@ class SessionAwareCache(BasePrefixCache):
|
|||||||
self._free_kv_aligned(req.req_pool_idx, target, req.kv_allocated_len)
|
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_allocated_len = min(req.kv_allocated_len, target)
|
||||||
req.kv_committed_len = min(req.kv_committed_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]
|
req.output_ids = req.output_ids[:finished_len]
|
||||||
|
|
||||||
def _free_kv_aligned(self, pool_idx: int, target: int, end: int):
|
def _free_kv_aligned(self, pool_idx: int, target: int, end: int):
|
||||||
|
|||||||
@@ -235,3 +235,39 @@ def test_nth_mid_abort_nukes_session_slot():
|
|||||||
# Shrink tests removed: streaming sessions are append-only after the
|
# Shrink tests removed: streaming sessions are append-only after the
|
||||||
# rollback fix in session_controller (rollback_aborted_req). The shrink
|
# rollback fix in session_controller (rollback_aborted_req). The shrink
|
||||||
# code path in cache_finished_req no longer exists.
|
# 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))
|
||||||
|
|||||||
Reference in New Issue
Block a user