@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 000c7df9ff

[KV Offload] Add KVCR secondary-tier adapter (#53624)

Signed-off-by: Moein Khazraee <moein@nvidia.com> Signed-off-by: Moein <moein@nvidia.com> Co-authored-by: Olga Andreeva <124622579+oandreeva-nv@users.noreply.github.com> Co-authored-by: Olga Andreeva <oandreeva@nvidia.com>

Moein Khazraee committed Sep 14, 2026 at 23:55 UTC 000c7df9ffd3e470980fd4cd6b8ec1b0585500ff
7 files changed +1211 -1
.buildkite/scripts/install-kv-offload.sh new
+18
@@ -0,0 +1,18 @@
1 +#!/bin/bash
2 +# SPDX-License-Identifier: Apache-2.0
3 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4 +
5 +set -euo pipefail
6 +
7 +if [ -z "${CUDA_VERSION:-}" ]; then
8 + exit 0
9 +fi
10 +
11 +# TODO: Use the PyPI wheel once it is published.
12 +uv pip install --system \
13 + "kvcr @ git+https://github.com/ai-dynamo/kvcr.git@v0.1.0"
14 +
15 +# Keep only the NIXL wheel matching the CI image's CUDA runtime.
16 +NIXL_VERSION=$(uv pip show --system nixl | sed -n 's/^Version: //p')
17 +uv pip uninstall --system nixl-cu12 nixl-cu13
18 +uv pip install --system --no-deps "nixl-cu${CUDA_VERSION%%.*}==${NIXL_VERSION}"
.buildkite/test_areas/misc.yaml
+1
@@ -130,6 +130,7 @@ steps:
130 source_file_dependencies: *v1-core-kv-metrics-deps
131 commands:
132 - export VLLM_WORKER_MULTIPROC_METHOD=spawn
133 + - bash /vllm-workspace/.buildkite/scripts/install-kv-offload.sh
134 - pytest -v -s v1/kv_offload
135 - pytest -v -s v1/simple_kv_offload
136 mirror:
tests/v1/kv_offload/tiering/test_factory.py
+5 -1
@@ -10,6 +10,7 @@ These tests verify:
10 3. Error paths — missing tier_type, unknown tier_type, duplicate registration.
11 """
12
13 +from importlib.util import find_spec
14 from unittest.mock import MagicMock
15
16 import pytest
@@ -42,8 +43,11 @@ def _make_mock_args():
43
44
45 def test_pre_registered_tiers_can_be_imported():
45 - """CI sentinel: example/fs/obj paths must import and yield SecondaryTierManager."""
46 + """CI sentinel: registered tiers import and yield SecondaryTierManager."""
47 for tier_type in SecondaryTierFactory._registry:
48 + # KVCR is an optional external dependency and may not be installed.
49 + if tier_type == "kvcr" and find_spec("kvcr") is None:
50 + continue
51 cls = SecondaryTierFactory._registry[tier_type]()
52 assert issubclass(cls, SecondaryTierManager)
53
tests/v1/kv_offload/tiering/test_kvcr_tier.py new
+544
@@ -0,0 +1,544 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Collection, Iterable, Mapping
5 +from types import SimpleNamespace
6 +from typing import Any
7 +
8 +import numpy as np
9 +import pytest
10 +
11 +pytest.importorskip("kvcr")
12 +
13 +from kvcr import ROUTER_HINT_KEY, KVCRBindings
14 +from kvcr.config import G3Options, KVCRBackendConfigs, KVCRConfig, KVCRGuardConfig
15 +from kvcr.policy import FIFOPolicy, G3FIFOPolicy, G3LRUPolicy, LRUPolicy
16 +from kvcr.types import (
17 + BlockKey,
18 + CacheTier,
19 + InventoryEvent,
20 + MemDescriptor,
21 + OpEntryResult,
22 + OpEntryStatus,
23 + OpHandle,
24 + QueryStatus,
25 +)
26 +
27 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
28 + OffloadingConnectorStats,
29 +)
30 +from vllm.v1.kv_offload.base import (
31 + LookupResult,
32 + Medium,
33 + OffloadKey,
34 + ReqContext,
35 + make_offload_key,
36 +)
37 +from vllm.v1.kv_offload.tiering.base import JobResult, TransferJob
38 +from vllm.v1.kv_offload.tiering.kvcr import manager as kvcr_manager
39 +from vllm.v1.kv_offload.tiering.kvcr.manager import KVCRSecondaryTierManager
40 +
41 +
42 +def _op_entries(
43 + entries: Mapping[BlockKey, bool],
44 +) -> dict[BlockKey, OpEntryResult]:
45 + return {
46 + key: OpEntryResult(OpEntryStatus.SUCCESS if success else OpEntryStatus.FAILED)
47 + for key, success in entries.items()
48 + }
49 +
50 +
51 +class _StubControlChannel:
52 + def __init__(self, endpoint: str) -> None:
53 + self.endpoint = endpoint
54 +
55 + def send(self, endpoint: str, message: bytes) -> bool:
56 + return True
57 +
58 + def recv(self) -> list[bytes]:
59 + return []
60 +
61 + def close(self) -> None:
62 + pass
63 +
64 +
65 +class RecordingKVCR:
66 + def __init__(self) -> None:
67 + self.config: KVCRConfig | None = None
68 + self.guard_config: KVCRGuardConfig | None = None
69 + self.backend_configs = KVCRBackendConfigs()
70 + self.constructor_bindings: KVCRBindings | None = None
71 + self.nixl_agent_name = "recording"
72 + self.framework_control: _StubControlChannel | None = None
73 + self.inventory_sink = None
74 + self.query_status = QueryStatus.MISS
75 + self.stats: OffloadingConnectorStats | None = None
76 + self.submit_hint_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
77 + self.discard_hint_calls: list[str] = []
78 + self.deliver_calls: list[
79 + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]], str | None]
80 + ] = []
81 + self.deposit_calls: list[
82 + tuple[OpHandle, dict[BlockKey, list[MemDescriptor]]]
83 + ] = []
84 + self.completed: list[tuple[OpHandle, dict[BlockKey, OpEntryResult]]] = []
85 + self._next_op_handle = 1
86 +
87 + def submit_hint(self, *args: Any, **kwargs: Any) -> None:
88 + self.submit_hint_calls.append((args, kwargs))
89 +
90 + def discard_hint(self, request_id: str) -> None:
91 + self.discard_hint_calls.append(request_id)
92 +
93 + def query(
94 + self,
95 + keys: Collection[BlockKey],
96 + request_id: str | None = None,
97 + ) -> list[tuple[QueryStatus, CacheTier | None]]:
98 + tier = {
99 + QueryStatus.HIT: CacheTier.LOCAL_G2,
100 + QueryStatus.FETCHING: CacheTier.LOCAL_G2,
101 + QueryStatus.FETCHABLE: CacheTier.REMOTE_G2,
102 + QueryStatus.MISS: None,
103 + }[self.query_status]
104 + return [(self.query_status, tier) for _ in keys]
105 +
106 + def deliver(
107 + self,
108 + blocks: Mapping[BlockKey, list[MemDescriptor]],
109 + request_id: str | None = None,
110 + ) -> OpHandle:
111 + op_handle = self._next_op_handle
112 + self._next_op_handle += 1
113 + self.deliver_calls.append((op_handle, dict(blocks), request_id))
114 + self.complete(op_handle, {key: True for key in blocks})
115 + return op_handle
116 +
117 + def deposit(
118 + self,
119 + blocks: Mapping[BlockKey, list[MemDescriptor]],
120 + ) -> OpHandle:
121 + op_handle = self._next_op_handle
122 + self._next_op_handle += 1
123 + self.deposit_calls.append((op_handle, dict(blocks)))
124 + self.complete(op_handle, {key: True for key in blocks})
125 + return op_handle
126 +
127 + def complete(
128 + self,
129 + op_handle: OpHandle,
130 + entries: Mapping[BlockKey, bool],
131 + ) -> None:
132 + self.completed.append((op_handle, _op_entries(entries)))
133 +
134 + def poll_completed(
135 + self,
136 + ) -> Iterable[tuple[OpHandle, dict[BlockKey, OpEntryResult]]]:
137 + completed = self.completed
138 + self.completed = []
139 + return completed
140 +
141 + def get_stats(self) -> OffloadingConnectorStats | None:
142 + stats = self.stats
143 + self.stats = None
144 + return stats
145 +
146 + def close(self) -> None:
147 + pass
148 +
149 +
150 +class _ExternalPolicy(FIFOPolicy):
151 + pass
152 +
153 +
154 +def _job(
155 + job_id: int,
156 + req_context: ReqContext,
157 + key: OffloadKey | None = None,
158 + chunk_id: int = 0,
159 +) -> TransferJob:
160 + return TransferJob(
161 + job_id=job_id,
162 + keys=[key if key is not None else OffloadKey(b"k0")],
163 + chunk_ids=np.array([chunk_id], dtype=np.int64),
164 + is_promotion=True,
165 + req_context=req_context,
166 + )
167 +
168 +
169 +def _make_tier(
170 + monkeypatch,
171 + kvcr: RecordingKVCR,
172 + *,
173 + enable_telemetry: bool = False,
174 + secondary_g2_slots: int = 0,
175 + kvcr_service_socket_path: str | None = None,
176 + compatibility_digest: str | None = None,
177 + enable_kv_cache_events: bool = False,
178 + self_describing_kv_events: bool = False,
179 + policy: str | None = None,
180 + g3: dict[str, object] | None = None,
181 + control_ports: list[int] | None = None,
182 + data_parallel_rank_local: int | None = None,
183 +) -> KVCRSecondaryTierManager:
184 + def make_control(_bind_host, bind_port, advertise_host):
185 + return _StubControlChannel(f"tcp://{advertise_host}:{int(bind_port)}")
186 +
187 + def make_kvcr(config, bindings, backend_configs, guard_config):
188 + kvcr.config = config
189 + kvcr.guard_config = guard_config
190 + kvcr.nixl_agent_name = config.nixl_agent_name
191 + kvcr.backend_configs = backend_configs
192 + kvcr.constructor_bindings = bindings
193 + kvcr.framework_control = bindings.framework_control
194 + kvcr.inventory_sink = bindings.inventory_sink
195 + return kvcr
196 +
197 + monkeypatch.setattr(kvcr_manager, "KVCR", make_kvcr)
198 + monkeypatch.setattr(kvcr_manager, "ZmqPeerControlChannel", make_control)
199 + return KVCRSecondaryTierManager(
200 + offloading_spec=SimpleNamespace(
201 + config=SimpleNamespace(
202 + parallel=SimpleNamespace(
203 + data_parallel_rank_local=data_parallel_rank_local,
204 + )
205 + ),
206 + kv_events_config=SimpleNamespace(
207 + enable_kv_cache_events=enable_kv_cache_events,
208 + self_describing_kv_events=self_describing_kv_events,
209 + ),
210 + ),
211 + primary_kv_view=memoryview(np.zeros((4, 16), dtype=np.int8)),
212 + tier_type="kvcr",
213 + router_capabilities=["router_hint"],
214 + control_host="127.0.0.1",
215 + control_ports=control_ports if control_ports is not None else [7777],
216 + control_advertise_host="127.0.0.1",
217 + enable_telemetry=enable_telemetry,
218 + secondary_g2_slots=secondary_g2_slots,
219 + kvcr_service_socket_path=kvcr_service_socket_path,
220 + compatibility_digest=compatibility_digest,
221 + policy=policy,
222 + g3=g3,
223 + local_dram_backend="UCX",
224 + remote_fw_dram_backend="UCX",
225 + )
226 +
227 +
228 +def test_kvcr_tier_configures_service_for_local_dp_rank(monkeypatch):
229 + """Keep the control endpoint and guard pool aligned to the local DP rank."""
230 + kvcr = RecordingKVCR()
231 + tier = _make_tier(
232 + monkeypatch,
233 + kvcr,
234 + kvcr_service_socket_path="/tmp/kvcr.sock",
235 + compatibility_digest="Opaque-Digest",
236 + secondary_g2_slots=1,
237 + control_ports=[7001, 7002],
238 + data_parallel_rank_local=1,
239 + )
240 +
241 + assert kvcr.framework_control is not None
242 + assert kvcr.framework_control.endpoint == "tcp://127.0.0.1:7002"
243 + assert kvcr.guard_config == KVCRGuardConfig(
244 + kvcr_service_socket_path="/tmp/kvcr.sock",
245 + guard_index=1,
246 + compatibility_digest="Opaque-Digest",
247 + )
248 + assert kvcr.config is not None
249 + assert kvcr.config.pool_layouts == [("", tier._primary_row_stride)]
250 + assert kvcr.backend_configs.local_dram is None
251 +
252 +
253 +def test_kvcr_tier_converts_g3_paths(monkeypatch, tmp_path):
254 + """Convert user-provided G3 paths to KVCR's typed configuration."""
255 + kvcr = RecordingKVCR()
256 + path = tmp_path / "g3.data"
257 +
258 + _make_tier(
259 + monkeypatch,
260 + kvcr,
261 + g3={"paths": [str(path)], "capacity_bytes_per_file": 64},
262 + )
263 +
264 + assert kvcr.backend_configs.g3 == G3Options(
265 + paths=(path,),
266 + capacity_bytes_per_file=64,
267 + )
268 +
269 +
270 +def test_kvcr_tier_adapts_request_and_load(monkeypatch):
271 + """Check hint forwarding, key conversion, load descriptors, and cleanup."""
272 + kvcr = RecordingKVCR()
273 + tier = _make_tier(monkeypatch, kvcr)
274 + router_hint = {"opaque": True}
275 + ctx = ReqContext(
276 + req_id="req",
277 + kv_transfer_params={ROUTER_HINT_KEY: router_hint, "unrelated": object()},
278 + )
279 + key = make_offload_key((123).to_bytes(8, "big"), 0)
280 + same_hash_other_group = make_offload_key((123).to_bytes(8, "big"), 7)
281 + other_key = make_offload_key((124).to_bytes(8, "big"), 0)
282 +
283 + tier.on_new_request(ctx)
284 +
285 + assert kvcr.submit_hint_calls == [((), {"request_id": "req", "hints": router_hint})]
286 +
287 + bindings = kvcr.constructor_bindings
288 + assert bindings is not None
289 + assert bindings.key_adapter is not None
290 + monkeypatch.setenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "0")
291 + decode = bindings.key_adapter.decode
292 + hash_123 = (123).to_bytes(8, "big")
293 + assert decode(BlockKey(bytes(key))) == hash_123
294 + assert decode(BlockKey(bytes(same_hash_other_group))) == hash_123
295 + assert decode(BlockKey(bytes(other_key))) == (124).to_bytes(8, "big")
296 +
297 + tier.submit_load(_job(7, ctx, key=key, chunk_id=2))
298 +
299 + # Check load destinations and job completion independently of hint parsing.
300 + assert len(kvcr.submit_hint_calls) == 1
301 + _, blocks, request_id = kvcr.deliver_calls[0]
302 + assert request_id == "req"
303 + assert list(blocks) == [key]
304 + (descriptor,) = blocks[key]
305 + assert descriptor.end_point_name == kvcr.nixl_agent_name
306 + assert descriptor.addr == tier._primary_base_addr + 2 * 16
307 + assert descriptor.size == 16
308 + assert list(tier.get_finished_jobs()) == [JobResult(7, True)]
309 +
310 + # Here we verify that request cleanup discards the request-scoped hint in
311 + # KVCR.
312 + tier.on_request_finished(ctx)
313 + assert kvcr.discard_hint_calls == ["req"]
314 +
315 +
316 +def test_kvcr_tier_allows_request_without_router_hint(monkeypatch):
317 + """Keep router hints optional for requests from non-hint-aware routers."""
318 + kvcr = RecordingKVCR()
319 + tier = _make_tier(monkeypatch, kvcr)
320 +
321 + tier.on_new_request(ReqContext(req_id="req"))
322 +
323 + assert kvcr.submit_hint_calls == []
324 +
325 +
326 +@pytest.mark.parametrize(
327 + ("status", "expected"),
328 + [
329 + (QueryStatus.MISS, LookupResult.MISS),
330 + (QueryStatus.HIT, LookupResult.HIT),
331 + (QueryStatus.FETCHABLE, LookupResult.HIT),
332 + (QueryStatus.FETCHING, LookupResult.RETRY),
333 + ],
334 +)
335 +def test_kvcr_tier_maps_query_status(monkeypatch, status, expected):
336 + """Map KVCR cache states to the scheduler's lookup contract."""
337 + kvcr = RecordingKVCR()
338 + kvcr.query_status = status
339 + tier = _make_tier(monkeypatch, kvcr)
340 +
341 + assert tier.lookup(OffloadKey(b"k0"), ReqContext(req_id="req")) is expected
342 +
343 +
344 +def test_kvcr_tier_serves_primary_pin_request(monkeypatch):
345 + """Hold primary-tier hits until KVCR releases the corresponding pin."""
346 + kvcr = RecordingKVCR()
347 + tier = _make_tier(monkeypatch, kvcr)
348 + keys = (BlockKey(b"k0"), BlockKey(b"k1"), BlockKey(b"k2"))
349 + hit_keys = (keys[0], keys[2])
350 + chunk_ids = {keys[0]: 1, keys[2]: 5}
351 + lifecycle: list[str] = []
352 +
353 + class Parent:
354 + def on_new_request(self, req_context):
355 + lifecycle.append("new")
356 + return SimpleNamespace()
357 +
358 + def lookup(self, key, req_context):
359 + return LookupResult.HIT if key in hit_keys else LookupResult.MISS
360 +
361 + def create_store_job(self, requested_keys, req_context):
362 + return TransferJob(
363 + job_id=11,
364 + keys=requested_keys,
365 + chunk_ids=np.array([chunk_ids[key] for key in requested_keys]),
366 + is_promotion=True,
367 + req_context=req_context,
368 + )
369 +
370 + def on_request_finished(self, req_context):
371 + lifecycle.append("finished")
372 +
373 + bindings = kvcr.constructor_bindings
374 + assert bindings is not None
375 + request = bindings.request_pin(keys)
376 +
377 + tier.serve_external_requests(Parent())
378 +
379 + [(queued_request, result)] = bindings.poll_pin_results()
380 + assert queued_request == request
381 + assert result is not None
382 + pin_handle, descriptors = result
383 + assert descriptors[keys[1]] is None
384 + block_descriptors = descriptors[keys[2]]
385 + assert block_descriptors is not None
386 + (descriptor,) = block_descriptors
387 + assert descriptor.addr == (tier._primary_base_addr + 5 * tier._primary_row_stride)
388 + assert lifecycle == ["new", "finished"]
389 +
390 + polls = 0
391 +
392 + def release_pin():
393 + nonlocal polls
394 + polls += 1
395 + bindings.release_pin(pin_handle)
396 + return []
397 +
398 + monkeypatch.setattr(kvcr, "poll_completed", release_pin)
399 + tier.drain_jobs()
400 +
401 + assert polls == 1
402 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)]
403 +
404 +
405 +def test_kvcr_telemetry_is_opt_in_and_namespaced_at_vllm_boundary(monkeypatch):
406 + """Keep telemetry opt-in and namespace metrics only at the vLLM boundary."""
407 + assert KVCRSecondaryTierManager.build_metric_definitions({}) == {}
408 + definitions = KVCRSecondaryTierManager.build_metric_definitions(
409 + {"enable_telemetry": True}
410 + )
411 + assert "vllm:kvcr_duration_seconds" in definitions
412 + assert "vllm:kvcr_transfer_blocks" in definitions
413 +
414 + kvcr = RecordingKVCR()
415 + tier = _make_tier(monkeypatch, kvcr, enable_telemetry=True)
416 + bindings = kvcr.constructor_bindings
417 + assert bindings is not None
418 + assert bindings.stats_factory is not None
419 + stats = bindings.stats_factory()
420 + stats.increase_counter(
421 + "kvcr_transfer_blocks",
422 + 2,
423 + ("remote_deliver",),
424 + )
425 + kvcr.stats = stats
426 +
427 + returned = tier.get_stats()
428 +
429 + assert returned is stats
430 + assert returned.reduce() == {"vllm:kvcr_transfer_blocks:('remote_deliver',)": 2}
431 +
432 +
433 +@pytest.mark.parametrize(
434 + ("policy", "expected_type"),
435 + [
436 + ("fifo", FIFOPolicy),
437 + ("lru", LRUPolicy),
438 + ("g3_fifo", G3FIFOPolicy),
439 + ("g3_lru", G3LRUPolicy),
440 + (f"{__name__}._ExternalPolicy", _ExternalPolicy),
441 + ],
442 +)
443 +def test_kvcr_tier_passes_selected_policy(monkeypatch, policy, expected_type):
444 + """Resolve every built-in and fully qualified external policy."""
445 + kvcr = RecordingKVCR()
446 + _make_tier(monkeypatch, kvcr, policy=policy)
447 +
448 + bindings = kvcr.constructor_bindings
449 + assert bindings is not None
450 + assert type(bindings.policy) is expected_type
451 +
452 +
453 +@pytest.mark.parametrize(
454 + ("socket_path", "digest"),
455 + [("/tmp/kvcr.sock", None), (None, "Opaque-Digest")],
456 +)
457 +def test_kvcr_tier_requires_complete_service_config(monkeypatch, socket_path, digest):
458 + """Reject partial guard configuration before connecting to the service."""
459 + with pytest.raises(ValueError, match="configured together"):
460 + _make_tier(
461 + monkeypatch,
462 + RecordingKVCR(),
463 + kvcr_service_socket_path=socket_path,
464 + compatibility_digest=digest,
465 + )
466 +
467 +
468 +def test_kvcr_tier_stores_and_emits_inventory(monkeypatch):
469 + """Cover store descriptors and inventory translation at the KVCR boundary."""
470 + kvcr = RecordingKVCR()
471 + tier = _make_tier(
472 + monkeypatch,
473 + kvcr,
474 + secondary_g2_slots=2,
475 + enable_kv_cache_events=True,
476 + self_describing_kv_events=True,
477 + )
478 + local_dram = kvcr.backend_configs.local_dram
479 + assert local_dram is not None
480 + assert [(name, size) for name, _, size in local_dram.pools] == [
481 + ("", 2 * tier._primary_row_stride)
482 + ]
483 +
484 + key = OffloadKey(b"k0")
485 + tier.submit_store(_job(11, ReqContext(req_id="req"), key=key, chunk_id=2))
486 +
487 + _, blocks = kvcr.deposit_calls[0]
488 + assert blocks[key][0].addr == tier._primary_base_addr + 2 * 16
489 + assert list(tier.get_finished_jobs()) == [JobResult(11, True)]
490 +
491 + assert kvcr.inventory_sink is not None
492 + kvcr.inventory_sink(
493 + InventoryEvent((BlockKey(bytes(key)),), CacheTier.LOCAL_G2, False)
494 + )
495 + kvcr.inventory_sink(InventoryEvent((BlockKey(bytes(key)),), CacheTier.G3, False))
496 + events = list(tier.take_events())
497 + assert [
498 + (event.keys, event.medium, event.ownership, event.removed) for event in events
499 + ] == [
500 + ([key], Medium.CPU, "kvcr", False),
501 + ([key], Medium.STORAGE, "kvcr", False),
502 + ]
503 + assert all(event.removal_expected for event in events)
504 + tier.shutdown()
505 +
506 +
507 +def test_kvcr_tier_requires_self_describing_inventory_events(monkeypatch):
508 + """Require tier-aware events whenever KVCR owns local cache inventory."""
509 + with pytest.raises(ValueError, match="self_describing_kv_events"):
510 + _make_tier(
511 + monkeypatch,
512 + RecordingKVCR(),
513 + secondary_g2_slots=1,
514 + enable_kv_cache_events=True,
515 + )
516 +
517 +
518 +def test_kvcr_tier_accumulates_block_results(monkeypatch):
519 + """Wait for every block result and preserve partial success while draining."""
520 + kvcr = RecordingKVCR()
521 + tier = _make_tier(monkeypatch, kvcr)
522 + keys = [OffloadKey(b"k0"), OffloadKey(b"k1")]
523 + tier.submit_load(
524 + TransferJob(
525 + job_id=13,
526 + keys=keys,
527 + chunk_ids=np.array([0, 1], dtype=np.int64),
528 + is_promotion=True,
529 + req_context=ReqContext(req_id="req"),
530 + )
531 + )
532 +
533 + op_handle, _, _ = kvcr.deliver_calls[0]
534 + kvcr.completed.clear()
535 + kvcr.complete(op_handle, {keys[0]: True})
536 + assert list(tier.get_finished_jobs()) == []
537 +
538 + kvcr.complete(op_handle, {keys[1]: False})
539 + tier.drain_jobs()
540 +
541 + [result] = tier.get_finished_jobs()
542 + assert result.job_id == 13
543 + assert not result.success
544 + assert result.successful_keys == {keys[0]}
vllm/v1/kv_offload/tiering/factory.py
+6
@@ -123,3 +123,9 @@ SecondaryTierFactory.register_tier(
123 "vllm.v1.kv_offload.tiering.obj.manager",
124 "ObjectStoreSecondaryTierManager",
125 )
126 +
127 +SecondaryTierFactory.register_tier(
128 + "kvcr",
129 + "vllm.v1.kv_offload.tiering.kvcr.manager",
130 + "KVCRSecondaryTierManager",
131 +)
vllm/v1/kv_offload/tiering/kvcr/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/v1/kv_offload/tiering/kvcr/manager.py new
+635
@@ -0,0 +1,635 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""vLLM secondary-tier adapter for KVCR."""
4 +
5 +import ctypes
6 +import mmap
7 +import socket
8 +import time
9 +import uuid
10 +from collections.abc import Collection, Iterable, Mapping
11 +from pathlib import Path
12 +from typing import TYPE_CHECKING, Any, cast
13 +
14 +from kvcr import (
15 + DURATION_METRIC,
16 + KVCR,
17 + ROUTER_HINT_CAPABILITIES,
18 + ROUTER_HINT_KEY,
19 + STATE_METRIC,
20 + TRANSFER_BLOCKS_METRIC,
21 + TRANSFER_BYTES_METRIC,
22 + KVCRBindings,
23 +)
24 +from kvcr.config import (
25 + FrameworkDramInput,
26 + G3Options,
27 + KeyAdapter,
28 + KVCRBackendConfigs,
29 + KVCRConfig,
30 + KVCRGuardConfig,
31 + LocalDramOptions,
32 + RemoteFWDramOptions,
33 +)
34 +from kvcr.control_channels import ZmqPeerControlChannel
35 +from kvcr.policy import (
36 + FIFOPolicy,
37 + G3FIFOPolicy,
38 + G3LRUPolicy,
39 + KVCachePolicy,
40 + LRUPolicy,
41 +)
42 +from kvcr.types import (
43 + BlockKey,
44 + CacheTier,
45 + InventoryEvent,
46 + MemDescriptor,
47 + OpHandle,
48 + PinRequestId,
49 + PinResult,
50 + QueryStatus,
51 +)
52 +from typing_extensions import override
53 +
54 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
55 + OffloadingConnectorStats,
56 +)
57 +from vllm.logger import init_logger
58 +from vllm.utils.import_utils import resolve_obj_by_qualname
59 +from vllm.v1.core.kv_cache_utils import (
60 + BlockHash,
61 + maybe_convert_block_hash,
62 +)
63 +from vllm.v1.kv_offload.base import (
64 + LookupResult,
65 + Medium,
66 + OffloadingCounterMetadata,
67 + OffloadingEvent,
68 + OffloadingGaugeMetadata,
69 + OffloadingHistogramMetadata,
70 + OffloadingMetricMetadata,
71 + OffloadKey,
72 + ReqContext,
73 + RequestOffloadingContext,
74 + get_offload_block_hash,
75 +)
76 +from vllm.v1.kv_offload.tiering.base import (
77 + JobResult,
78 + ParentManager,
79 + SecondaryTierManager,
80 + TransferJob,
81 +)
82 +
83 +if TYPE_CHECKING:
84 + from vllm.v1.kv_offload.base import OffloadingSpec
85 +
86 +
87 +_REQUIRED_ROUTER_CAPABILITIES = ROUTER_HINT_CAPABILITIES
88 +
89 +logger = init_logger(__name__)
90 +
91 +_BUILTIN_POLICIES: dict[str, type[KVCachePolicy]] = {
92 + "fifo": FIFOPolicy,
93 + "lru": LRUPolicy,
94 + "g3_fifo": G3FIFOPolicy,
95 + "g3_lru": G3LRUPolicy,
96 +}
97 +
98 +
99 +def _resolve_policy(name: str | None) -> KVCachePolicy | None:
100 + if name is None:
101 + return None
102 + policy_type = _BUILTIN_POLICIES.get(name)
103 + if policy_type is None:
104 + if "." not in name:
105 + raise ValueError(
106 + f"Unknown KVCR policy: {name!r}. "
107 + f"Supported: {list(_BUILTIN_POLICIES)}; external policies "
108 + "must use a fully qualified module.Class path."
109 + )
110 + policy_type = resolve_obj_by_qualname(name)
111 + if not isinstance(policy_type, type) or not issubclass(
112 + policy_type, KVCachePolicy
113 + ):
114 + raise TypeError(f"{name} is not a KVCachePolicy")
115 + return policy_type()
116 +
117 +
118 +def _vllm_metric_name(name: str) -> str:
119 + return f"vllm:{name}"
120 +
121 +
122 +def _kvcr_metric_definitions() -> dict[str, OffloadingMetricMetadata]:
123 + return {
124 + _vllm_metric_name(DURATION_METRIC): OffloadingHistogramMetadata(
125 + documentation="KVCR operation and stage duration in seconds.",
126 + labelnames=("scope", "result"),
127 + ),
128 + _vllm_metric_name(TRANSFER_BYTES_METRIC): OffloadingCounterMetadata(
129 + documentation="Bytes transferred by KVCR operations.",
130 + labelnames=("operation",),
131 + ),
132 + _vllm_metric_name(TRANSFER_BLOCKS_METRIC): OffloadingCounterMetadata(
133 + documentation="KV blocks transferred by successful KVCR operations.",
134 + labelnames=("operation",),
135 + ),
136 + _vllm_metric_name(STATE_METRIC): OffloadingGaugeMetadata(
137 + documentation="Current KVCR metadata and operation counts.",
138 + labelnames=("resource",),
139 + ),
140 + }
141 +
142 +
143 +class _FrameworkPinAdapter:
144 + """Adapt KVCR framework pin callbacks to primary-tier store jobs."""
145 +
146 + def __init__(self, tier: "KVCRSecondaryTierManager") -> None:
147 + self._tier = tier
148 + self._next_request_id = 0
149 + self._next_pin_request_id = 0
150 + self._pending_pins: dict[PinRequestId, tuple[BlockKey, ...]] = {}
151 + self._completed_pins: list[tuple[PinRequestId, PinResult]] = []
152 + self._pin_jobs: dict[str, int] = {}
153 + self._pin_job_results: list[JobResult] = []
154 +
155 + def request_pin(self, keys: Collection[BlockKey]) -> PinRequestId:
156 + # The parent callback is valid only from serve_external_requests().
157 + request = PinRequestId(self._next_pin_request_id)
158 + self._next_pin_request_id += 1
159 + self._pending_pins[request] = tuple(keys)
160 + return request
161 +
162 + def process_pending(self, parent: ParentManager) -> None:
163 + pending = self._pending_pins
164 + self._pending_pins = {}
165 + self._completed_pins.extend(
166 + (request, self._resolve_pin(parent, keys))
167 + for request, keys in pending.items()
168 + )
169 +
170 + def poll_pin_results(self) -> list[tuple[PinRequestId, PinResult]]:
171 + completed = self._completed_pins
172 + self._completed_pins = []
173 + return completed
174 +
175 + def cancel_pin_request(self, request: PinRequestId) -> None:
176 + self._pending_pins.pop(request, None)
177 +
178 + def _resolve_pin(
179 + self, parent: ParentManager, keys: Collection[BlockKey]
180 + ) -> tuple[str, dict[BlockKey, list[MemDescriptor] | None]] | None:
181 + offload_keys = tuple(OffloadKey(bytes(key)) for key in keys)
182 + request_id = f"kvcr-source:{self._next_request_id}"
183 + self._next_request_id += 1
184 + req_context = ReqContext(req_id=request_id)
185 + started = False
186 + job: TransferJob | None = None
187 + pin_handle: str | None = None
188 + try:
189 + parent.on_new_request(req_context)
190 + started = True
191 + hit_keys = tuple(
192 + key
193 + for key in offload_keys
194 + if parent.lookup(key, req_context) is LookupResult.HIT
195 + )
196 + if not hit_keys:
197 + return None
198 +
199 + job = parent.create_store_job(hit_keys, req_context)
200 + job_keys = tuple(job.keys)
201 + chunk_ids = tuple(int(chunk_id) for chunk_id in job.chunk_ids)
202 + if len(job_keys) != len(chunk_ids) or set(job_keys) != set(hit_keys):
203 + return None
204 +
205 + descriptors: dict[BlockKey, list[MemDescriptor] | None] = {
206 + BlockKey(bytes(key)): None for key in offload_keys
207 + }
208 + descriptors.update(
209 + (
210 + BlockKey(bytes(key)),
211 + [self._tier._make_descriptor(chunk_id)],
212 + )
213 + for key, chunk_id in zip(job_keys, chunk_ids)
214 + )
215 + pin_handle = f"kvcr-job-{job.job_id}"
216 + self._pin_jobs[pin_handle] = job.job_id
217 + return pin_handle, descriptors
218 + except Exception:
219 + logger.warning("KVCR source acquisition failed", exc_info=True)
220 + return None
221 + finally:
222 + if started:
223 + try:
224 + parent.on_request_finished(req_context)
225 + except Exception:
226 + logger.warning("KVCR source request cleanup failed", exc_info=True)
227 + if job is not None and pin_handle is None:
228 + self._pin_job_results.append(
229 + JobResult(job_id=job.job_id, success=False)
230 + )
231 +
232 + def release_pin(self, pin_handle: str) -> bool:
233 + job_id = self._pin_jobs.pop(pin_handle, None)
234 + if job_id is not None:
235 + self._pin_job_results.append(JobResult(job_id=job_id, success=True))
236 + return True
237 +
238 + def has_active_pins(self) -> bool:
239 + return bool(self._pin_jobs)
240 +
241 + def take_pin_job_results(self) -> list[JobResult]:
242 + results = self._pin_job_results
243 + self._pin_job_results = []
244 + return results
245 +
246 +
247 +class _VllmKeyAdapter:
248 + """Adapt vLLM's local key format to KVCR."""
249 +
250 + def encode(self, framework_key: object) -> BlockKey:
251 + if not isinstance(framework_key, bytes):
252 + raise TypeError("vLLM offload keys must be bytes")
253 + return BlockKey(framework_key)
254 +
255 + def decode(self, key: BlockKey) -> int | bytes:
256 + block_hash = BlockHash(get_offload_block_hash(OffloadKey(key)))
257 + return maybe_convert_block_hash(block_hash)
258 +
259 +
260 +# Job ID, remaining blocks, aggregate success, and successful load keys.
261 +_JobState = tuple[int, int, bool, set[OffloadKey] | None]
262 +
263 +
264 +class KVCRSecondaryTierManager(SecondaryTierManager):
265 + """Secondary tier wrapper around the KVCR KV P2P API."""
266 +
267 + @classmethod
268 + @override
269 + def build_metric_definitions(
270 + cls, extra_config: dict[str, Any]
271 + ) -> dict[str, OffloadingMetricMetadata]:
272 + if not extra_config.get("enable_telemetry", False):
273 + return {}
274 + return _kvcr_metric_definitions()
275 +
276 + def __init__(
277 + self,
278 + offloading_spec: "OffloadingSpec",
279 + primary_kv_view: memoryview,
280 + tier_type: str,
281 + router_capabilities: Iterable[str] | None = None,
282 + control_host: str = "0.0.0.0",
283 + control_ports: list[int] | None = None,
284 + control_advertise_host: str | None = None,
285 + eager_ctrl_connect: bool = True,
286 + opportunistic_query: bool = True,
287 + enable_telemetry: bool = False,
288 + operation_timeout_ms: int = 1000,
289 + abandon_timeout_ms: int = 5000,
290 + metadata_retry_interval_ms: int = 100,
291 + secondary_g2_slots: int = 0,
292 + kvcr_service_socket_path: str | None = None,
293 + compatibility_digest: str | None = None,
294 + policy: str | None = None,
295 + g3: dict[str, Any] | None = None,
296 + local_dram_backend: str = "UCX",
297 + remote_fw_dram_backend: str = "UCX",
298 + ) -> None:
299 + super().__init__(offloading_spec, primary_kv_view, tier_type)
300 + selected_policy = _resolve_policy(policy)
301 + if (kvcr_service_socket_path is None) != (compatibility_digest is None):
302 + raise ValueError(
303 + "kvcr_service_socket_path and compatibility_digest must be "
304 + "configured together"
305 + )
306 + missing = _REQUIRED_ROUTER_CAPABILITIES - set(router_capabilities or ())
307 + if missing:
308 + raise ValueError(
309 + f"KVCR secondary tier requires router capabilities: {sorted(missing)}"
310 + )
311 + events_enabled = offloading_spec.kv_events_config.enable_kv_cache_events
312 + self_describing_events = (
313 + offloading_spec.kv_events_config.self_describing_kv_events
314 + )
315 + if (
316 + events_enabled
317 + and (
318 + secondary_g2_slots > 0
319 + or kvcr_service_socket_path is not None
320 + or g3 is not None
321 + )
322 + and not self_describing_events
323 + ):
324 + raise ValueError(
325 + "KVCR local G2/G3 inventory requires self_describing_kv_events"
326 + )
327 + self._key_adapter: KeyAdapter = _VllmKeyAdapter()
328 + advertise_host = control_advertise_host or socket.gethostname()
329 + dp_local_rank = offloading_spec.config.parallel.data_parallel_rank_local
330 + if dp_local_rank is None:
331 + if control_ports is None or len(control_ports) != 1:
332 + raise ValueError(
333 + "control_ports must contain exactly one port when "
334 + "data_parallel_rank_local is unset"
335 + )
336 + dp_local_rank = 0
337 + if control_ports is None or not 0 <= dp_local_rank < len(control_ports):
338 + raise ValueError(
339 + "control_ports must contain a port for "
340 + f"local data-parallel rank {dp_local_rank}"
341 + )
342 + if any(not 1 <= port <= 65535 for port in control_ports):
343 + raise ValueError("control_ports must be between 1 and 65535")
344 + control = ZmqPeerControlChannel(
345 + control_host,
346 + control_ports[dp_local_rank],
347 + advertise_host,
348 + )
349 + # Give colocated workers distinct transport listen ports.
350 + with socket.socket() as _s:
351 + _s.bind(("", 0))
352 + _nixl_listen_port = _s.getsockname()[1]
353 + self._primary_base_addr = ctypes.addressof(
354 + ctypes.c_char.from_buffer(primary_kv_view)
355 + )
356 + if primary_kv_view.strides is None:
357 + raise ValueError("primary KV memoryview must expose strides")
358 + self._primary_row_stride = int(primary_kv_view.strides[0])
359 + if secondary_g2_slots < 0:
360 + raise ValueError("secondary_g2_slots must be non-negative")
361 + if (
362 + secondary_g2_slots or kvcr_service_socket_path is not None
363 + ) and not events_enabled:
364 + logger.warning(
365 + "KVCR local DRAM is enabled, but KV cache events are disabled; "
366 + "local DRAM inventory will not be published"
367 + )
368 + local_mapping: mmap.mmap | None = None
369 + local_dram: LocalDramOptions | None = None
370 + if secondary_g2_slots:
371 + if kvcr_service_socket_path is not None:
372 + logger.warning(
373 + "secondary_g2_slots is ignored when "
374 + "kvcr_service_socket_path is configured"
375 + )
376 + else:
377 + local_mapping = mmap.mmap(
378 + -1, secondary_g2_slots * self._primary_row_stride
379 + )
380 + local_dram = LocalDramOptions(
381 + pools=[
382 + (
383 + "",
384 + ctypes.addressof(ctypes.c_char.from_buffer(local_mapping)),
385 + len(local_mapping),
386 + )
387 + ],
388 + backend=local_dram_backend,
389 + )
390 + guard_config = (
391 + KVCRGuardConfig(
392 + kvcr_service_socket_path=kvcr_service_socket_path,
393 + guard_index=dp_local_rank,
394 + compatibility_digest=compatibility_digest,
395 + )
396 + if kvcr_service_socket_path is not None and compatibility_digest is not None
397 + else None
398 + )
399 + g3_config = None
400 + if g3 is not None:
401 + g3_values = dict(g3)
402 + g3_values["paths"] = tuple(Path(path) for path in g3_values["paths"])
403 + g3_config = G3Options(**g3_values)
404 + nixl_agent_name = f"KVCR-{uuid.uuid4()}"
405 + self._framework_pin_adapter = _FrameworkPinAdapter(self)
406 + self._inventory_events: list[OffloadingEvent] = []
407 +
408 + try:
409 + self._kvcr = KVCR(
410 + KVCRConfig(
411 + pool_layouts=[("", self._primary_row_stride)],
412 + enable_telemetry=enable_telemetry,
413 + operation_timeout_ms=operation_timeout_ms,
414 + abandon_timeout_ms=abandon_timeout_ms,
415 + nixl_agent_name=nixl_agent_name,
416 + nixl_listen_port=_nixl_listen_port,
417 + ),
418 + KVCRBindings(
419 + request_pin=self._framework_pin_adapter.request_pin,
420 + poll_pin_results=self._framework_pin_adapter.poll_pin_results,
421 + release_pin=self._framework_pin_adapter.release_pin,
422 + cancel_pin_request=(self._framework_pin_adapter.cancel_pin_request),
423 + framework_control=control,
424 + key_adapter=self._key_adapter,
425 + inventory_sink=(self._record_inventory if events_enabled else None),
426 + stats_factory=(
427 + OffloadingConnectorStats if enable_telemetry else None
428 + ),
429 + policy=selected_policy,
430 + ),
431 + KVCRBackendConfigs(
432 + framework_dram=FrameworkDramInput(
433 + self._primary_base_addr, primary_kv_view.nbytes
434 + ),
435 + local_dram=local_dram,
436 + g3=g3_config,
437 + remote_fw_dram=RemoteFWDramOptions(
438 + eager_ctrl_connect=eager_ctrl_connect,
439 + opportunistic_query=opportunistic_query,
440 + metadata_retry_interval_ms=metadata_retry_interval_ms,
441 + backend=remote_fw_dram_backend,
442 + ),
443 + ),
444 + guard_config,
445 + )
446 + except BaseException:
447 + control.close()
448 + if local_mapping is not None:
449 + local_mapping.close()
450 + raise
451 + self._local_dram_mmap = local_mapping
452 + self._finished_jobs: list[JobResult] = []
453 + self._jobs_by_op: dict[OpHandle, _JobState] = {}
454 +
455 + @override
456 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:
457 + block_key = self._key_adapter.encode(key)
458 + status, _ = self._kvcr.query((block_key,), req_context.req_id)[0]
459 + if status is QueryStatus.FETCHING:
460 + return LookupResult.RETRY
461 + if status in (QueryStatus.HIT, QueryStatus.FETCHABLE):
462 + return LookupResult.HIT
463 + return LookupResult.MISS
464 +
465 + @override
466 + def submit_load(self, job_metadata: TransferJob) -> None:
467 + blocks = {
468 + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
469 + for key, chunk_id in zip(
470 + job_metadata.keys, job_metadata.chunk_ids, strict=True
471 + )
472 + }
473 + if not blocks:
474 + self._finished_jobs.append(
475 + JobResult(job_id=job_metadata.job_id, success=True)
476 + )
477 + return
478 +
479 + op_handle = self._kvcr.deliver(
480 + blocks, request_id=job_metadata.req_context.req_id
481 + )
482 + self._jobs_by_op[op_handle] = (
483 + job_metadata.job_id,
484 + len(blocks),
485 + True,
486 + set(),
487 + )
488 +
489 + @override
490 + def submit_store(self, job_metadata: TransferJob) -> None:
491 + blocks = {
492 + self._key_adapter.encode(key): [self._make_descriptor(int(chunk_id))]
493 + for key, chunk_id in zip(
494 + job_metadata.keys, job_metadata.chunk_ids, strict=True
495 + )
496 + }
497 + if not blocks:
498 + self._finished_jobs.append(
499 + JobResult(job_id=job_metadata.job_id, success=True)
500 + )
501 + return
502 + op_handle = self._kvcr.deposit(blocks)
503 + self._jobs_by_op[op_handle] = (
504 + job_metadata.job_id,
505 + len(blocks),
506 + True,
507 + None,
508 + )
509 +
510 + @override
511 + def get_finished_jobs(self) -> Iterable[JobResult]:
512 + results = self._finished_jobs
513 + self._finished_jobs = []
514 + results.extend(self._poll_finished_jobs())
515 + return results
516 +
517 + @override
518 + def get_stats(self) -> OffloadingConnectorStats | None:
519 + stats = cast(OffloadingConnectorStats | None, self._kvcr.get_stats())
520 + if stats is not None:
521 + # Add the framework namespace only at the vLLM boundary.
522 + for metrics in stats.data.values():
523 + for name in tuple(metrics):
524 + metrics[_vllm_metric_name(name)] = metrics.pop(name)
525 + return stats
526 +
527 + @override
528 + def serve_external_requests(self, parent: ParentManager) -> None:
529 + self._framework_pin_adapter.process_pending(parent)
530 +
531 + def _poll_finished_jobs(self) -> list[JobResult]:
532 + results: list[JobResult] = []
533 + for op_handle, entries in self._kvcr.poll_completed():
534 + job_state = self._jobs_by_op.get(op_handle)
535 + if job_state is None:
536 + continue
537 + job_id, remaining, success, successful_keys = job_state
538 + remaining -= len(entries)
539 + success = success and all(entry.success for entry in entries.values())
540 + if successful_keys is not None:
541 + successful_keys.update(
542 + OffloadKey(bytes(key))
543 + for key, entry in entries.items()
544 + if entry.success
545 + )
546 + if remaining > 0:
547 + self._jobs_by_op[op_handle] = (
548 + job_id,
549 + remaining,
550 + success,
551 + successful_keys,
552 + )
553 + continue
554 + self._jobs_by_op.pop(op_handle, None)
555 + results.append(
556 + JobResult(
557 + job_id=job_id,
558 + success=success,
559 + successful_keys=successful_keys if not success else None,
560 + )
561 + )
562 + results.extend(self._framework_pin_adapter.take_pin_job_results())
563 + return results
564 +
565 + @override
566 + def has_pending_work(self) -> bool:
567 + # Keep scheduler polling: control messages may arrive without a load job.
568 + return True
569 +
570 + @override
571 + def take_events(self) -> Iterable[OffloadingEvent]:
572 + events = self._inventory_events
573 + self._inventory_events = []
574 + return events
575 +
576 + @override
577 + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
578 + params = getattr(req_context, "kv_transfer_params", None)
579 + if isinstance(params, Mapping):
580 + hint = params.get(ROUTER_HINT_KEY)
581 + if hint is not None:
582 + self._kvcr.submit_hint(
583 + request_id=req_context.req_id,
584 + hints=hint,
585 + )
586 + return RequestOffloadingContext()
587 +
588 + @override
589 + def on_request_finished(self, req_context: ReqContext) -> None:
590 + self._kvcr.discard_hint(req_context.req_id)
591 +
592 + @override
593 + def drain_jobs(self) -> None:
594 + while self._jobs_by_op or self._framework_pin_adapter.has_active_pins():
595 + drained = self._poll_finished_jobs()
596 + if drained:
597 + self._finished_jobs.extend(drained)
598 + else:
599 + time.sleep(0.001)
600 +
601 + @override
602 + def shutdown(self) -> None:
603 + self.drain_jobs()
604 + # TODO: Keep registered buffers alive if KVCR.close() fails.
605 + self._kvcr.close()
606 + if self._local_dram_mmap is not None:
607 + self._local_dram_mmap.close()
608 + self._local_dram_mmap = None
609 +
610 + def _record_inventory(self, event: InventoryEvent) -> None:
611 + if event.tier is CacheTier.LOCAL_G2:
612 + medium = Medium.CPU
613 + elif event.tier is CacheTier.G3:
614 + medium = Medium.STORAGE
615 + else:
616 + raise ValueError(f"unsupported KVCR inventory tier: {event.tier}")
617 + self._inventory_events.append(
618 + OffloadingEvent(
619 + keys=[OffloadKey(bytes(key)) for key in event.keys],
620 + medium=medium,
621 + removed=event.removed,
622 + ownership="kvcr",
623 + removal_expected=True,
624 + )
625 + )
626 +
627 + def _make_descriptor(self, chunk_id: int) -> MemDescriptor:
628 + return MemDescriptor(
629 + end_point_name=self._kvcr.config.nixl_agent_name,
630 + mem_type="DRAM",
631 + addr=self._primary_base_addr + chunk_id * self._primary_row_stride,
632 + size=self._primary_row_stride,
633 + device_Id=0,
634 + info="",
635 + )