From 79fd1bc7eda1854ba6b243c737ecc4bad74ffb84 Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 13 May 2026 14:11:10 +0300 Subject: [PATCH] [kv_offload] Add req_id to ReqContext for per-request tracking (#42507) Signed-off-by: Ronen Schaffer --- .../unit/offloading_connector/test_scheduler.py | 2 +- tests/v1/kv_offload/cpu/test_manager.py | 6 ++++-- .../kv_transfer/kv_connector/v1/offloading/scheduler.py | 5 ++++- vllm/v1/kv_offload/base.py | 1 + 4 files changed, 10 insertions(+), 4 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 22eb1a86d4f..9a8783b084e 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -627,7 +627,7 @@ def _make_scheduler_with_lookup( return scheduler -_EMPTY_REQ_CTX = ReqContext() +_EMPTY_REQ_CTX = ReqContext(req_id="") class TestMaximalPrefixLookup: diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index 4641a3fe9f4..8bb77bf7631 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -19,9 +19,11 @@ from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy -def make_req_context(kv_transfer_params: dict | None = None) -> ReqContext: +def make_req_context( + req_id: str = "", kv_transfer_params: dict | None = None +) -> ReqContext: """Create a ReqContext as production code would, from a request's params.""" - return ReqContext(kv_transfer_params=kv_transfer_params) + return ReqContext(req_id=req_id, kv_transfer_params=kv_transfer_params) _EMPTY_REQ_CTX = make_req_context() 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 137eaef9788..ef016bb02eb 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -138,7 +138,10 @@ class RequestOffloadState: self.group_states = tuple( RequestGroupState() for _ in self.config.kv_group_configs ) - self.req_context = ReqContext(kv_transfer_params=self.req.kv_transfer_params) + self.req_context = ReqContext( + req_id=self.req.request_id, + kv_transfer_params=self.req.kv_transfer_params, + ) def update_offload_keys(self) -> None: for group_config, group_state in zip( diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index fed47cf5993..482865cb7cf 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -46,6 +46,7 @@ def get_offload_group_idx(key: OffloadKey) -> int: @dataclass class ReqContext: + req_id: str kv_transfer_params: dict[str, Any] | None = None