From 36f5ff3172eb75433e54e496d64c473e951f8ecc Mon Sep 17 00:00:00 2001 From: Moein Khazraee <33970824+mkhazraee@users.noreply.github.com> Date: Tue, 27 Oct 2026 16:41:00 -0700 Subject: [PATCH] [Perf][KV Offload] Reduce KVCR adapter overhead (#63889) Signed-off-by: Moein Khazraee Co-authored-by: vllm-agent --- tests/v1/kv_offload/tiering/test_kvcr_tier.py | 50 ++++++++++++++----- vllm/v1/kv_offload/tiering/kvcr/manager.py | 32 ++++++------ 2 files changed, 52 insertions(+), 30 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_kvcr_tier.py b/tests/v1/kv_offload/tiering/test_kvcr_tier.py index a6f67f80f3f..beca574d888 100644 --- a/tests/v1/kv_offload/tiering/test_kvcr_tier.py +++ b/tests/v1/kv_offload/tiering/test_kvcr_tier.py @@ -390,6 +390,7 @@ def test_kvcr_tier_adapts_request_and_load(monkeypatch, backpressure): assert list(blocks) == [key] (descriptor,) = blocks[key] assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) + assert type(descriptor.element_index) is int assert list(tier.get_finished_jobs()) == [JobResult(7, True)] # Here we verify that request cleanup discards the request-scoped hint in @@ -479,13 +480,24 @@ def test_kvcr_tier_maps_query_status(monkeypatch, status, expected): assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected -def test_kvcr_tier_serves_primary_pin_request(monkeypatch): - """Hold primary-tier hits until KVCR releases the corresponding pin.""" +@pytest.mark.parametrize( + ("job_keys", "chunk_ids", "success"), + [ + (None, [1, 3], True), + ((b"k2", b"k0"), [3, 1], True), + (None, [1], False), + ((b"k0", b"other"), [1, 3], False), + ], + ids=["original-keys", "reordered-keys", "short-chunks", "unexpected-key"], +) +def test_kvcr_tier_serves_primary_pin_request( + monkeypatch, job_keys, chunk_ids, success +): + """Validate pin jobs and hold accepted hits until KVCR releases the pin.""" kvcr = RecordingKVCR() tier = _make_tier(monkeypatch, kvcr) keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2")) hit_keys = (keys[0], keys[2]) - chunk_ids = {keys[0]: 1, keys[2]: 3} lifecycle: list[str] = [] class Parent: @@ -499,8 +511,8 @@ def test_kvcr_tier_serves_primary_pin_request(monkeypatch): def create_store_job(self, requested_keys, req_context): return TransferJob( job_id=11, - keys=requested_keys, - chunk_ids=np.array([chunk_ids[key] for key in requested_keys]), + keys=requested_keys if job_keys is None else job_keys, + chunk_ids=np.array(chunk_ids, dtype=np.int32), is_promotion=True, req_context=req_context, ) @@ -516,6 +528,11 @@ def test_kvcr_tier_serves_primary_pin_request(monkeypatch): [(queued_request, result)] = bindings.poll_pin_results() assert queued_request == request + assert lifecycle == ["new", "finished"] + if not success: + assert result is None + assert list(tier.get_finished_jobs()) == [JobResult(11, False)] + return assert result is not None pin_handle, descriptors = result assert descriptors[keys[1]] is None @@ -523,7 +540,7 @@ def test_kvcr_tier_serves_primary_pin_request(monkeypatch): assert block_descriptors is not None (descriptor,) = block_descriptors assert descriptor == MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=3) - assert lifecycle == ["new", "finished"] + assert type(descriptor.element_index) is int polls = 0 @@ -626,6 +643,7 @@ def test_kvcr_tier_stores_and_emits_inventory(monkeypatch): assert blocks[key] == [ MemoryRef(end_point_name=kvcr.nixl_agent_name, element_index=2) ] + assert type(blocks[key][0].element_index) is int assert list(tier.get_finished_jobs()) == [JobResult(11, True)] assert kvcr.inventory_sink is not None @@ -655,7 +673,16 @@ def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch): ) -def test_kvcr_tier_accumulates_block_results(monkeypatch): +@pytest.mark.parametrize( + ("successes", "expected"), + [ + ((True, False), JobResult(13, False, {OffloadKey(b"k0")})), + ((False, True), JobResult(13, False, {OffloadKey(b"k1")})), + ((True, True), JobResult(13, True)), + ], + ids=["success-then-failure", "failure-then-success", "all-success"], +) +def test_kvcr_tier_accumulates_block_results(monkeypatch, successes, expected): """Wait for every block result and preserve partial success while draining.""" kvcr = RecordingKVCR() tier = _make_tier(monkeypatch, kvcr) @@ -672,13 +699,10 @@ def test_kvcr_tier_accumulates_block_results(monkeypatch): op_handle, _, _ = kvcr.deliver_calls[0] kvcr.completed.clear() - kvcr.complete(op_handle, {keys[0]: True}) + kvcr.complete(op_handle, {keys[0]: successes[0]}) assert list(tier.get_finished_jobs()) == [] - kvcr.complete(op_handle, {keys[1]: False}) + kvcr.complete(op_handle, {keys[1]: successes[1]}) tier.drain_jobs() - [result] = tier.get_finished_jobs() - assert result.job_id == 13 - assert not result.success - assert result.successful_keys == {keys[0]} + assert list(tier.get_finished_jobs()) == [expected] diff --git a/vllm/v1/kv_offload/tiering/kvcr/manager.py b/vllm/v1/kv_offload/tiering/kvcr/manager.py index 1f3d9638b49..e24531120b5 100644 --- a/vllm/v1/kv_offload/tiering/kvcr/manager.py +++ b/vllm/v1/kv_offload/tiering/kvcr/manager.py @@ -201,20 +201,17 @@ class _FrameworkPinAdapter: job = parent.create_store_job(hit_keys, req_context) job_keys = tuple(job.keys) - chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids) - if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys): + if len(job_keys) != len(job.chunk_ids) or ( + job_keys is not hit_keys and set(job_keys) != set(hit_keys) + ): return None - descriptors: dict[BlockKey, list[MemoryRef] | None] = { - BlockKey(bytes(key)): None for key in offload_keys - } - descriptors.update( - ( - BlockKey(bytes(key)), - [self._tier._make_descriptor(chunk_id)], - ) - for key, chunk_id in zip(job_keys, chunk_ids) + descriptors: dict[BlockKey, list[MemoryRef] | None] = dict.fromkeys( + cast(tuple[BlockKey, ...], offload_keys) ) + make_descriptor = self._tier._make_descriptor + for key, chunk_id in zip(job_keys, job.chunk_ids.tolist()): + descriptors[BlockKey(key)] = [make_descriptor(chunk_id)] pin_handle = f"kvcr-job-{job.job_id}" self._pin_jobs[pin_handle] = job.job_id return pin_handle, descriptors @@ -476,6 +473,7 @@ class KVCRSecondaryTierManager(SecondaryTierManager): if local_mapping is not None: local_mapping.close() raise + self._nixl_agent_name = self._kvcr.config.nixl_agent_name self._local_dram_mmap = local_mapping self._finished_jobs: list[JobResult] = [] self._jobs_by_op: dict[OpHandle, _JobState] = {} @@ -493,9 +491,9 @@ class KVCRSecondaryTierManager(SecondaryTierManager): @override def submit_load(self, job_metadata: TransferJob) -> None: blocks = { - self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + self._key_adapter.encode(key): [self._make_descriptor(chunk_id)] for key, chunk_id in zip( - job_metadata.keys, job_metadata.chunk_ids, strict=True + job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True ) } if not blocks: @@ -517,9 +515,9 @@ class KVCRSecondaryTierManager(SecondaryTierManager): @override def submit_store(self, job_metadata: TransferJob) -> None: blocks = { - self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))] + self._key_adapter.encode(key): [self._make_descriptor(chunk_id)] for key, chunk_id in zip( - job_metadata.keys, job_metadata.chunk_ids, strict=True + job_metadata.keys, job_metadata.chunk_ids.tolist(), strict=True ) } if not blocks: @@ -566,7 +564,7 @@ class KVCRSecondaryTierManager(SecondaryTierManager): job, remaining, success, successful_keys = job_state remaining -= len(entries) success = success and all(entry.success for entry in entries.values()) - if successful_keys is not None: + if successful_keys is not None and (remaining > 0 or not success): successful_keys.update( OffloadKey(bytes(key)) for key, entry in entries.items() @@ -706,6 +704,6 @@ class KVCRSecondaryTierManager(SecondaryTierManager): def _make_descriptor(self, chunk_id: int) -> MemoryRef: return MemoryRef( - end_point_name=self._kvcr.config.nixl_agent_name, + end_point_name=self._nixl_agent_name, element_index=chunk_id, )