[Perf][KV Offload] Reduce KVCR adapter overhead (#63889)

Signed-off-by: Moein Khazraee <moein@nvidia.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

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]

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,
         )