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)
|
||||
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):
|
||||
|
||||
@@ -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))
|
||||
|
||||
Reference in New Issue
Block a user