From fada48bb2a199d278a36e95532441e31375d1e68 Mon Sep 17 00:00:00 2001 From: Giannis Ntovas Date: Sun, 11 Oct 2026 16:21:19 +0300 Subject: [PATCH] [Bugfix][KV Offload] Store the chunk of a loaded partial tail After a load that ends in a partial tail, update_state_after_alloc sets next_stored_chunk_idx to cdiv(num_cached_tokens, tokens_per_chunk). This skips the chunk that holds the partial tail. The request then computes the rest of that chunk, but the connector never stores it. The next lookup of the same prefix stops at this chunk, finds the same partial tail again, and the request computes all tokens after it again. For a full-attention group, start the store at the chunk of the partial tail. Sliding-window and Mamba groups keep the skip: their block for that chunk can hold the tail state, which is not the state at the chunk end. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Giannis Ntovas --- .../offloading_connector/test_scheduler.py | 35 +++++++++++++++++++ .../kv_connector/v1/offloading/scheduler.py | 10 ++++-- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index e71c3dbf18d4..523d8fc28732 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -462,6 +462,41 @@ def test_partial_lookup_returns_exact_boundary_and_group_load_keys(): assert req_status.partial_tail_boundary is None +def test_partial_tail_load_keeps_tail_chunk_storable(): + """The chunk of a loaded partial tail is not a hit chunk. + + Geometry: 16 tokens per chunk, a load of 28 tokens. Chunk 0 is a full hit, + and chunk 1 holds the partial tail (tokens 16-27). The request computes the + rest of chunk 1, so the full-attention group must store it when it is full. + If the store starts at chunk 2, chunk 1 is never stored, and every later + lookup of this prefix stops at the same partial tail. The Mamba group keeps + the skip: its block for chunk 1 holds the state at token 28. + """ + scheduler = _make_partial_tail_scheduler() + request = _make_partial_tail_request(scheduler) + req_status = scheduler._req_status["req"] + req_status.num_locally_computed_tokens = 0 + req_status.update_offload_keys() + + scheduler.manager.lookup.return_value = LookupResult.HIT + assert scheduler._lookup(req_status) == 28 + + scheduler.update_state_after_alloc( + request, + KVCacheBlocks( + ( + [KVCacheBlock(31), KVCacheBlock(32)], + [KVCacheBlock(0, is_null=True), KVCacheBlock(41)], + ) + ), + num_external_tokens=28, + ) + + full_state, mamba_state = req_status.group_states + assert full_state.next_stored_chunk_idx == 1 + assert mamba_state.next_stored_chunk_idx == 2 + + def test_lookup_cap_stops_at_authoritative_prefix_boundary(): scheduler = _make_partial_tail_scheduler() request = _make_partial_tail_request(scheduler) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 806267dcf3dd..edb06e2e752a 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -1249,9 +1249,15 @@ def update_state_after_alloc( # Skip prefix-hit chunks for block-level policy; for # request-level, next_stored_chunk_idx stays at 0 so all - # chunks (including hits) are offloaded. + # chunks (including hits) are offloaded. In a full-attention group + # the chunk of a partial tail is not a hit: the request computes the + # rest of it, so store it when it is full. Sliding-window and Mamba + # blocks of that chunk can hold the tail state, so they keep the skip. if req_status.offloading_context.policy == OffloadPolicy.CHUNK_LEVEL: - group_state.next_stored_chunk_idx = num_chunks + group_state.next_stored_chunk_idx = num_chunks - ( + partial_tail_boundary is not None + and group_config.sliding_window_size_in_chunks is None + ) src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context) dst_spec = GPULoadStoreSpec(