Delete dead rematch path in SessionAwareCache.release_session (#22735)
This commit is contained in:
@@ -111,88 +111,6 @@ def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch):
|
||||
assert allocator.freed[0].tolist() == list(range(32, 40))
|
||||
|
||||
|
||||
def test_release_session_recomputes_current_tree_owned_prefix():
|
||||
page_size = 16
|
||||
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()
|
||||
|
||||
full_match = MatchResult(
|
||||
device_indices=torch.tensor(list(range(16)) + list(range(64, 96))),
|
||||
last_device_node="stale-expanded",
|
||||
last_host_node="stale-expanded",
|
||||
)
|
||||
protected_match = MatchResult(
|
||||
device_indices=torch.tensor(list(range(16))),
|
||||
last_device_node="current-protected",
|
||||
last_host_node="current-protected",
|
||||
)
|
||||
inner = _FakeInnerCache(
|
||||
req_to_token_pool,
|
||||
allocator,
|
||||
page_size,
|
||||
match_results=[full_match, protected_match],
|
||||
)
|
||||
tree_cache = SessionAwareCache(inner)
|
||||
|
||||
tree_cache.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=0,
|
||||
kv_committed_len=48,
|
||||
kv_allocated_len=48,
|
||||
last_node="outdated-node",
|
||||
cache_protected_len=32,
|
||||
)
|
||||
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
|
||||
|
||||
tree_cache.release_session("session-a", req)
|
||||
|
||||
assert inner.dec_lock_ref_calls == ["current-protected"]
|
||||
assert req_to_token_pool.free_slots == [0]
|
||||
assert len(allocator.freed) == 1
|
||||
assert allocator.freed[0].tolist() == list(range(16, 48))
|
||||
|
||||
|
||||
def test_release_session_never_grows_tree_owned_prefix():
|
||||
page_size = 16
|
||||
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()
|
||||
|
||||
overmatched = MatchResult(
|
||||
device_indices=torch.tensor(list(range(48))),
|
||||
last_device_node="overmatched-node",
|
||||
last_host_node="overmatched-node",
|
||||
)
|
||||
capped_match = MatchResult(
|
||||
device_indices=torch.tensor(list(range(16))),
|
||||
last_device_node="original-lock-node",
|
||||
last_host_node="original-lock-node",
|
||||
)
|
||||
inner = _FakeInnerCache(
|
||||
req_to_token_pool,
|
||||
allocator,
|
||||
page_size,
|
||||
match_results=[overmatched, capped_match],
|
||||
)
|
||||
tree_cache = SessionAwareCache(inner)
|
||||
|
||||
tree_cache.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=0,
|
||||
kv_committed_len=48,
|
||||
kv_allocated_len=48,
|
||||
last_node="outdated-node",
|
||||
cache_protected_len=16,
|
||||
)
|
||||
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
|
||||
|
||||
tree_cache.release_session("session-a", req)
|
||||
|
||||
assert inner.dec_lock_ref_calls == ["original-lock-node"]
|
||||
assert req_to_token_pool.free_slots == [0]
|
||||
assert len(allocator.freed) == 1
|
||||
assert allocator.freed[0].tolist() == list(range(16, 48))
|
||||
|
||||
|
||||
def test_match_prefix_abort_does_not_restore_live_session_slot():
|
||||
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
|
||||
Reference in New Issue
Block a user