From f39bcf1e30693f2e620e3a2ec732c5f7192408fb Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 6 May 2026 17:31:21 +0300 Subject: [PATCH] [KV Offload] Return None from lookup() for in-flight blocks (#41795) Signed-off-by: Ronen Schaffer --- tests/v1/kv_offload/cpu/test_manager.py | 12 ++++++------ vllm/v1/kv_offload/cpu/manager.py | 6 +++++- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index ef5d61e7b3..e043590a41 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -163,9 +163,9 @@ def test_cpu_manager(): ), ) - # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + # lookup [1, 2] -> write in-flight, not yet ready + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None # no events so far assert list(cpu_manager.take_events()) == [] @@ -296,9 +296,9 @@ class TestARCPolicy: ), ) - # lookup [1, 2] -> not ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + # lookup [1, 2] -> write in-flight, not yet ready + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None # no events so far assert list(cpu_manager.take_events()) == [] diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index 90d56915cc..80bcb568f9 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -86,7 +86,11 @@ class CPUOffloadingManager(OffloadingManager): def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: block = self._policy.get(key) - return block is not None and block.is_ready + if block is None: + return False + if not block.is_ready: + return None # write in-flight; caller should retry + return True def prepare_load( self,