[KV Offload] Pass request context by omerpaz95 · Pull Request #39185 · vllm-project/vllm

omerpaz95 · 2026-04-07T12:10:20Z

Purpose

Pass per-request context (kv_transfer_params) through the CPU offloading manager interface.

The OffloadingManager ABC and CPUOffloadingManager now accept a ReqContext argument on lookup, prepare_load, and prepare_store. ReqContext wraps request.kv_transfer_params (a free-form dict set via sampling_params.extra_args) and is constructed at each call site in OffloadingConnectorScheduler.

This is wiring-only: ReqContext is accepted but not yet consumed by CPUOffloadingManager.

Structural changes:

ReqContext dataclass is defined in vllm/v1/kv_offload/abstract.py (shared layer) rather than cpu/manager.py, so other OffloadingManager implementations can use it without importing from the CPU-specific module.

Test Plan

# Lint
cd vllm
pre-commit run

# Scheduler unit tests
pytest tests/v1/kv_connector/test_offloading_scheduler.py -v -s

Test Result

Pre-commit (ruff, mypy, all hooks): all passed.

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

gemini-code-assist

Code Review

This pull request introduces a ReqContext dataclass to pass per-request parameters, such as kv_transfer_params, through the KV offloading pipeline. The OffloadingManager interface and its CPU implementation have been updated to include this context in lookup, prepare_load, and prepare_store methods. Feedback indicates that for full consistency and to support features like tenant-aware resource management, the ReqContext should also be added to the touch, complete_load, and complete_store methods. Furthermore, existing unit test mocks need to be updated to match the new method signatures to avoid TypeError exceptions.

gemini-code-assist · 2026-04-07T12:12:19Z

+    def lookup(
+        self,
+        block_hashes: Iterable[BlockHash],
+        req_context: ReqContext,


The ReqContext should also be added to the touch, complete_load, and complete_store methods. Since the stated goal of this PR is to support features like tenant routing and priority hints, these lifecycle management methods must also be context-aware. For example, a tenant-aware manager needs the context in complete_load to correctly release tenant-specific pins or quotas, and in touch to update the correct tenant's LRU state. Omitting the context here creates an inconsistent API that will likely require another breaking change soon.

gemini-code-assist · 2026-04-07T12:12:19Z

-
-        src_spec = self.manager.prepare_load(block_hashes)
+        req_context = ReqContext(kv_transfer_params=request.kv_transfer_params)
+        src_spec = self.manager.prepare_load(block_hashes, req_context)


Updating the prepare_load signature to require req_context will break the existing unit tests. Specifically, in tests/v1/kv_connector/unit/offloading_connector/utils.py, the MockOffloadingSpec initializes self.manager.prepare_load with a single-argument lambda (lambda block_hashes: ...). This will now raise a TypeError when called with two arguments in the scheduler. Please ensure the test mocks are updated to accommodate the new parameter.

mergify · 2026-04-08T18:16:14Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @omerpaz95.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

orozery · 2026-04-13T09:17:06Z

@@ -304,7 +309,8 @@ def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):
            assert len(req.block_hashes) >= num_gpu_blocks

            new_offload_keys = group_state.offload_keys[start_block_idx:num_blocks]
-            store_output = self.manager.prepare_store(new_offload_keys)
+            req_context = ReqContext(kv_transfer_params=req.kv_transfer_params)


Can we re-use ReqContext by introducing a field in RequestOffloadState?

orozery · 2026-04-14T07:11:42Z

        """
        Prepare the given blocks to be offloaded.
        The given blocks will be protected from eviction until
        complete_store is called.

        Args:
+            req_context: per-request context (e.g. kv_transfer_params).


Let's move this after keys to match the args order.

orozery · 2026-04-14T07:12:31Z

@@ -114,7 +123,11 @@ def complete_load(self, keys: Iterable[OffloadKey]) -> None:
            assert block.ref_cnt > 0, f"Block {key!r} ref_cnt is already 0"
            block.ref_cnt -= 1

-    def prepare_store(self, keys: Iterable[OffloadKey]) -> PrepareStoreOutput | None:
+    def prepare_store(
+        self, 


Suggested change

self,

self,

orozery · 2026-04-14T07:14:20Z

+    return ReqContext(kv_transfer_params=kv_transfer_params)
+
+
+_CTX = make_req_context()


Suggested change

_CTX = make_req_context()

_EMPTY_REQ_CTX = make_req_context()

orozery

Thanks! Please also make sure pre-commit passes following these changes.

orozery · 2026-04-14T11:00:14Z

-        hits = self.manager.lookup(offload_keys[start_block_idx:])
+        hits = self.manager.lookup(
+            offload_keys[start_block_idx:],
+            self._req_status[request.request_id].req_context,


Suggested change

self._req_status[request.request_id].req_context,

req_status.req_context,

orozery · 2026-04-14T11:00:38Z

@@ -249,7 +255,9 @@ def update_state_after_alloc(
        assert len(request.block_hashes) // self.config.block_size_factor >= num_blocks
        offload_keys = group_state.offload_keys[start_block_idx:num_blocks]

-        src_spec = self.manager.prepare_load(offload_keys)
+        src_spec = self.manager.prepare_load(
+            offload_keys, self._req_status[request.request_id].req_context


Suggested change

offload_keys, self._req_status[request.request_id].req_context

offload_keys, req_status.req_context

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

orozery

LGTM

…floading test lambdas for upstream req_context parameter (#1375) ## Root Cause Upstream vLLM PR vllm-project/vllm#39185 (commit `4353c9cb4a`) added a `req_context: ReqContext` parameter to `OffloadingManager.prepare_store()`, `prepare_load()`, and `lookup()` methods. The vllm-gaudi tests use mock `side_effect` lambdas and direct lambda assignments that only accepted the original single parameter (`block_hashes`/`keys`), but now receive two positional arguments. ## Fix Updated all 14 lambda signatures in `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` and 1 lambda in `utils.py` to accept the new `req_context` parameter (unused in test mocks). ## Files Changed - `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` — all `prepare_store.side_effect` lambdas: `lambda block_hashes:` → `lambda block_hashes, req_context:` - `tests/unit_tests/kv_offload/offloading_connector/utils.py` — `prepare_load` lambda: `lambda block_hashes:` → `lambda block_hashes, req_context:` ## Verification Verified on Gaudi3 HPU (g3 pod): - **Before fix**: 8/8 tests FAILED with `TypeError: lambda() takes 1 positional argument but 2 were given` - **After fix**: 8/8 tests PASSED (6.52s) Tests: `python -m pytest tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py -v` Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com>

…floading test lambdas for upstream req_context parameter (vllm-project#1375) ## Root Cause Upstream vLLM PR vllm-project/vllm#39185 (commit `4353c9cb4a`) added a `req_context: ReqContext` parameter to `OffloadingManager.prepare_store()`, `prepare_load()`, and `lookup()` methods. The vllm-gaudi tests use mock `side_effect` lambdas and direct lambda assignments that only accepted the original single parameter (`block_hashes`/`keys`), but now receive two positional arguments. ## Fix Updated all 14 lambda signatures in `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` and 1 lambda in `utils.py` to accept the new `req_context` parameter (unused in test mocks). ## Files Changed - `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` — all `prepare_store.side_effect` lambdas: `lambda block_hashes:` → `lambda block_hashes, req_context:` - `tests/unit_tests/kv_offload/offloading_connector/utils.py` — `prepare_load` lambda: `lambda block_hashes:` → `lambda block_hashes, req_context:` ## Verification Verified on Gaudi3 HPU (g3 pod): - **Before fix**: 8/8 tests FAILED with `TypeError: lambda() takes 1 positional argument but 2 were given` - **After fix**: 8/8 tests PASSED (6.52s) Tests: `python -m pytest tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py -v` Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com> Signed-off-by: Yeonsil Yoon <yeon.sil.yoon@intel.com>

…floading test lambdas for upstream req_context parameter (vllm-project#1375) ## Root Cause Upstream vLLM PR vllm-project/vllm#39185 (commit `4353c9cb4a`) added a `req_context: ReqContext` parameter to `OffloadingManager.prepare_store()`, `prepare_load()`, and `lookup()` methods. The vllm-gaudi tests use mock `side_effect` lambdas and direct lambda assignments that only accepted the original single parameter (`block_hashes`/`keys`), but now receive two positional arguments. ## Fix Updated all 14 lambda signatures in `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` and 1 lambda in `utils.py` to accept the new `req_context` parameter (unused in test mocks). ## Files Changed - `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` — all `prepare_store.side_effect` lambdas: `lambda block_hashes:` → `lambda block_hashes, req_context:` - `tests/unit_tests/kv_offload/offloading_connector/utils.py` — `prepare_load` lambda: `lambda block_hashes:` → `lambda block_hashes, req_context:` ## Verification Verified on Gaudi3 HPU (g3 pod): - **Before fix**: 8/8 tests FAILED with `TypeError: lambda() takes 1 positional argument but 2 were given` - **After fix**: 8/8 tests PASSED (6.52s) Tests: `python -m pytest tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py -v` Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com> Signed-off-by: bmyrcha <bartosz.myrcha@intel.com>

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: Avinash Singh <avinashsingh.rcoem@gmail.com>

…floading test lambdas for upstream req_context parameter (#1375) ## Root Cause Upstream vLLM PR vllm-project/vllm#39185 (commit `4353c9cb4a`) added a `req_context: ReqContext` parameter to `OffloadingManager.prepare_store()`, `prepare_load()`, and `lookup()` methods. The vllm-gaudi tests use mock `side_effect` lambdas and direct lambda assignments that only accepted the original single parameter (`block_hashes`/`keys`), but now receive two positional arguments. ## Fix Updated all 14 lambda signatures in `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` and 1 lambda in `utils.py` to accept the new `req_context` parameter (unused in test mocks). ## Files Changed - `tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py` — all `prepare_store.side_effect` lambdas: `lambda block_hashes:` → `lambda block_hashes, req_context:` - `tests/unit_tests/kv_offload/offloading_connector/utils.py` — `prepare_load` lambda: `lambda block_hashes:` → `lambda block_hashes, req_context:` ## Verification Verified on Gaudi3 HPU (g3 pod): - **Before fix**: 8/8 tests FAILED with `TypeError: lambda() takes 1 positional argument but 2 were given` - **After fix**: 8/8 tests PASSED (6.52s) Tests: `python -m pytest tests/unit_tests/kv_offload/offloading_connector/test_scheduler.py -v` Signed-off-by: Paweł Olejniczak <pawelx.olejniczak@intel.com> Signed-off-by: Artur Fierka <artur.fierka@intel.com>

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

omerpaz95 requested review from ApostaC, NickLucche and orozery as code owners April 7, 2026 12:10

mergify Bot added v1 kv-connector labels Apr 7, 2026

gemini-code-assist Bot reviewed Apr 7, 2026

View reviewed changes

mergify Bot added needs-rebase and removed needs-rebase labels Apr 8, 2026

orozery reviewed Apr 13, 2026

View reviewed changes

orozery reviewed Apr 14, 2026

View reviewed changes

orozery added the ready ONLY add when PR is ready to merge/full CI is needed label Apr 14, 2026

omerpaz95 requested review from DarkLight1337, LucasWilkinson, WoosukKwon, alexm-redhat, heheda12345, jikunshang, markmc, mgoin, njhill, noooop, pavanimajety, robertgshaw2-redhat, tjtanaa, tlrmchlsmth, vadiklyutiy, yewentao256 and ywang96 as code owners April 15, 2026 07:54

omerpaz95 added 2 commits April 15, 2026 16:13

Fixed PR comment.

34fd30a

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

Fixed tests to pass CI.

40f913b

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

omerpaz95 force-pushed the pass_request_context branch from 49aa8cf to 40f913b Compare April 15, 2026 13:15

omerpaz95 added 5 commits April 15, 2026 16:15

Merge branch 'main' into pass_request_context

5be8fc9

Merge branch 'main' into pass_request_context

d8b3eda

Merge branch 'main' into pass_request_context

bb21758

Merge branch 'main' into pass_request_context

bdc5652

Merge branch 'main' into pass_request_context

c091eec

omerpaz95 requested a review from xuechendi as a code owner April 16, 2026 18:08

Merge branch 'main' into pass_request_context

d01741c

orozery approved these changes Apr 19, 2026

View reviewed changes

github-project-automation Bot moved this to Ready in NVIDIA Apr 19, 2026

github-project-automation Bot moved this from To Triage to Ready in gpt-oss Issues & Enhancements Apr 19, 2026

orozery merged commit 4353c9c into vllm-project:main Apr 19, 2026
63 checks passed

github-project-automation Bot moved this from Ready to Done in gpt-oss Issues & Enhancements Apr 19, 2026

github-project-automation Bot moved this from Todo to Done in AMD Apr 19, 2026

github-project-automation Bot moved this from Ready to Done in NVIDIA Apr 19, 2026

pawel-olejniczak mentioned this pull request Apr 20, 2026

[FIX_FOR_VLLM_CUSTOM=d886c26d4d4fef7d079696beb4ece1cfb4b008a8] Fix offloading test lambdas for upstream req_context parameter vllm-project/vllm-gaudi#1375

Merged

avinashsingh77 pushed a commit to avinashsingh77/vllm that referenced this pull request Apr 27, 2026

[KV Offload] Pass request context (vllm-project#39185)

ab30629

Signed-off-by: omerpaz95 <omerpaz95@gmail.com> Signed-off-by: Avinash Singh <avinashsingh.rcoem@gmail.com>

ronensc mentioned this pull request Apr 30, 2026

[KV Offload] Pass ReqContext to touch(), complete_load(), and complete_store() #41366

Merged

3 of 4 tasks

mystous pushed a commit to mystous/vllm_hybrid that referenced this pull request May 10, 2026

[KV Offload] Pass request context (vllm-project#39185)

bf1ee0d

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026

[KV Offload] Pass request context (vllm-project#39185)

0c82b18

Signed-off-by: omerpaz95 <omerpaz95@gmail.com>

Navigation Menu

Repository navigation

[KV Offload] Pass request context - #39185

[KV Offload] Pass request context#39185
orozery merged 13 commits into
vllm-project:mainfrom
omerpaz95:pass_request_context

omerpaz95 commented Apr 7, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Apr 7, 2026

Uh oh!

gemini-code-assist Bot Apr 7, 2026

Uh oh!

mergify Bot commented Apr 8, 2026

Uh oh!

orozery Apr 13, 2026

Uh oh!

orozery Apr 14, 2026

Uh oh!

orozery Apr 14, 2026

Uh oh!

orozery Apr 14, 2026

Uh oh!

orozery left a comment

Uh oh!

orozery Apr 14, 2026

Uh oh!

orozery Apr 14, 2026

Uh oh!

orozery left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		return ReqContext(kv_transfer_params=kv_transfer_params)


		_CTX = make_req_context()

	_CTX = make_req_context()
	_EMPTY_REQ_CTX = make_req_context()

	self._req_status[request.request_id].req_context,
	req_status.req_context,

	offload_keys, self._req_status[request.request_id].req_context
	offload_keys, req_status.req_context

Conversation

omerpaz95 commented Apr 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Apr 7, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Apr 8, 2026

Uh oh!

orozery Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

orozery Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

orozery Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

orozery Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

orozery left a comment

Choose a reason for hiding this comment

Uh oh!

orozery Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

orozery Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

orozery left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

omerpaz95 commented Apr 7, 2026 •

edited

Loading