@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 77cba0259f

[KV Offloading] Per-request tier filtering with TierFilter/TierMatcher (#48123)

Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>

Ronen Schaffer committed Jul 27, 2026 at 13:29 UTC 77cba0259f21dfc6f4f6298f54f3e43c9cc824c3

16 files changed +301 -52

tests/v1/kv_connector/unit/offloading_connector/test_events.py

+11 -13

@@ -9,8 +9,6 @@ from tests.v1.kv_connector.unit.utils import create_vllm_config

9 from vllm.config import KVEventsConfig, KVTransferConfig

10 from vllm.distributed.kv_events import (

11 MEDIUM_CPU,

12 - MEDIUM_FS,

13 - MEDIUM_OBJ,

12 BlockRemoved,

13 BlockStored,

14 )

@@ -33,6 +31,7 @@ from vllm.v1.kv_cache_interface import (

31 )

32 from vllm.v1.kv_offload.base import (

33 Locality,

34 + Medium,

35 OffloadingEvent,

36 OffloadingKVEventsConfig,

37 OffloadKey,

@@ -40,7 +39,7 @@ from vllm.v1.kv_offload.base import (

39 )

40 from vllm.v1.kv_offload.tiering.spec import TieringOffloadingSpec

43 -_CPU_MEDIUM = MEDIUM_CPU

42 +_CPU_MEDIUM = Medium.CPU

43 _FULL_ATTENTION_EVENT_SPEC = OffloadingEventGroupSpec(

44 kv_cache_spec_kind=KVCacheSpecKind.FULL_ATTENTION.value,

45 kv_cache_spec_sliding_window=None,

@@ -135,7 +134,7 @@ def _record_lookup_chunks(

134

135 def _stored_event(

136 keys: list[OffloadKey],

138 - medium: str = _CPU_MEDIUM,

137 + medium: Medium = _CPU_MEDIUM,

138 locality: Locality | None = None,

139 ) -> OffloadingEvent:

140 return OffloadingEvent(

@@ -148,7 +147,7 @@ def _stored_event(

147

148 def _removed_event(

149 keys: list[OffloadKey],

151 - medium: str = _CPU_MEDIUM,

150 + medium: Medium = _CPU_MEDIUM,

151 locality: Locality | None = None,

152 ) -> OffloadingEvent:

153 return OffloadingEvent(

@@ -181,7 +180,7 @@ def test_take_events_forwards_locality_to_rich_store():

180

181 events = list(

182 tracker.take_events(

184 - [_stored_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)]

183 + [_stored_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)]

184 )

185 )

186

@@ -199,7 +198,7 @@ def test_take_events_forwards_locality_to_placeholder_store():

198

199 events = list(

200 tracker.take_events(

202 - [_stored_event([key], locality=Locality.REMOTE, medium=MEDIUM_FS)]

201 + [_stored_event([key], locality=Locality.REMOTE, medium=Medium.STORAGE)]

202 )

203 )

204

@@ -216,7 +215,7 @@ def test_take_events_forwards_locality_to_remove():

215

216 events = list(

217 tracker.take_events(

219 - [_removed_event([key], locality=Locality.LOCAL, medium=MEDIUM_FS)]

218 + [_removed_event([key], locality=Locality.LOCAL, medium=Medium.STORAGE)]

219 )

220 )

221

@@ -240,7 +239,7 @@ def test_take_events_publishes_routable_block_stored():

239

240 for i, event in enumerate(batch1):

241 assert isinstance(event, BlockStored)

243 - assert event.medium == _CPU_MEDIUM

242 + assert event.medium == _CPU_MEDIUM.value

243 assert event.block_hashes == [_wire_hash(_hash(i))]

244 assert event.block_size == block_size

245 assert event.token_ids == list(

@@ -324,7 +323,7 @@ def test_lookup_promotion_factor_gt_1_store_and_remove():

323 assert len(removed) == 1

324 assert isinstance(removed[0], BlockRemoved)

325 assert removed[0].block_hashes == expected_hashes

327 - assert removed[0].medium == _CPU_MEDIUM

326 + assert removed[0].medium == _CPU_MEDIUM.value

327 assert removed[0].group_idx == 0

328 assert not tracker._pending_event_metadata

329

@@ -444,12 +443,11 @@ def test_pending_cpu_removal_consumes_hit_backfill_until_next_hit():

443 ]

444

445

447 -@pytest.mark.parametrize("medium", [MEDIUM_FS, MEDIUM_OBJ])

448 -def test_secondary_stored_event_does_not_mutate_cpu_metadata(medium: str):

446 +def test_secondary_stored_event_does_not_mutate_cpu_metadata():

447 tracker, _, _, key = _lookup_chunk()

448 expected_metadata = dict(tracker._pending_event_metadata)

449

452 - stored = list(tracker.take_events([_stored_event([key], medium)]))

450 + stored = list(tracker.take_events([_stored_event([key], Medium.STORAGE)]))

451 assert stored[0].token_ids == [1, 2, 3, 4]

452 assert tracker._pending_event_metadata == expected_metadata

453

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+4 -3

@@ -33,6 +33,7 @@ from vllm.v1.kv_cache_interface import (

33 )

34 from vllm.v1.kv_offload.base import (

35 LookupResult,

36 + Medium,

37 OffloadingEvent,

38 OffloadingManager,

39 OffloadPolicy,

@@ -299,7 +300,7 @@ def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(

300

301 assert not tracker._pending_event_metadata

302

302 - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))

303 + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False))

304 events = list(runner.connector_scheduler.take_events())

305 assert len(events) == 1

306 assert isinstance(events[0], BlockStored)

@@ -320,7 +321,7 @@ def test_abort_before_hit_uses_placeholder_then_later_hit_heals_removal(

321 )

322 assert key in tracker._pending_event_metadata

323

323 - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=True))

324 + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=True))

325 [event] = runner.connector_scheduler.take_events()

326 assert isinstance(event, BlockRemoved)

327 assert event.medium == MEDIUM_CPU

@@ -355,7 +356,7 @@ def test_promotion_hit_precedes_stored_event_translation(

356 raw_events: list[OffloadingEvent] = []

357

358 def lookup(key, req_context):

358 - raw_events.append(OffloadingEvent(keys=[key], medium=MEDIUM_CPU, removed=False))

359 + raw_events.append(OffloadingEvent(keys=[key], medium=Medium.CPU, removed=False))

360 return LookupResult.HIT

361

362 def take_raw_events():

tests/v1/kv_offload/cpu/test_manager.py

+2 -2

@@ -6,10 +6,10 @@ from dataclasses import dataclass

6 import numpy as np

7 import pytest

9 -from vllm.distributed.kv_events import MEDIUM_CPU

9 from vllm.v1.kv_offload.base import (

10 LoadStoreSpec,

11 LookupResult,

12 + Medium,

13 OffloadingEvent,

14 OffloadKey,

15 PrepareStoreOutput,

@@ -115,7 +115,7 @@ def verify_events(

115 stores: list[set[OffloadKey]] = []

116 evictions: list[set[OffloadKey]] = []

117 for event in events:

118 - assert event.medium == MEDIUM_CPU

118 + assert event.medium == Medium.CPU

119 if event.removed:

120 evictions.append(set(event.keys))

121 else:

tests/v1/kv_offload/tiering/test_fs_tier.py

+3 -4

@@ -18,10 +18,10 @@ import numpy as np

18 import pytest

19 import torch

21 -from vllm.distributed.kv_events import MEDIUM_FS

21 from vllm.v1.kv_offload.base import (

22 Locality,

23 LookupResult,

24 + Medium,

25 OffloadingEvent,

26 OffloadingKVEventsConfig,

27 OffloadKey,

@@ -514,8 +514,7 @@ def test_successful_store_emits_stored_event(fs_tier_with_events):

514 events = list(tier.take_events())

515 assert len(events) == 1

516 assert events[0].keys == keys

517 - # Literal medium pins the wire contract, not just the constant choice.

518 - assert events[0].medium == "FS"

517 + assert events[0].medium == Medium.STORAGE

518 assert events[0].locality is Locality.LOCAL

519 assert not events[0].removed

520 # take_events drains the buffer.

@@ -685,7 +684,7 @@ def test_cascade_store_emits_fs_event_through_tiering_manager(tmp_path):

684 events.extend(manager.take_events())

685 time.sleep(0.01)

686

688 - fs_events = [e for e in events if e.medium == MEDIUM_FS]

687 + fs_events = [e for e in events if e.medium == Medium.STORAGE]

688 assert len(fs_events) == 1

689 assert set(fs_events[0].keys) == set(keys)

690 assert not fs_events[0].removed

tests/v1/kv_offload/tiering/test_obj_tier.py

+2 -2

@@ -21,6 +21,7 @@ import torch

21 from vllm.v1.kv_offload.base import (

22 Locality,

23 LookupResult,

24 + Medium,

25 OffloadingKVEventsConfig,

26 OffloadKey,

27 ReqContext,

@@ -473,8 +474,7 @@ class TestObjTierKVEvents:

474 events = list(self.tier.take_events())

475 assert len(events) == 1

476 assert events[0].keys == keys

476 - # Literal medium pins the wire contract, not just the constant choice.

477 - assert events[0].medium == "OBJ"

477 + assert events[0].medium == Medium.STORAGE

478 assert events[0].locality is Locality.REMOTE

479 assert not events[0].removed

480 # take_events drains the buffer.

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+136 -3

@@ -20,8 +20,13 @@ import torch

20 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

21 OffloadingConnectorStats,

22 )

23 +from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

24 + _parse_tier_filter,

25 +)

26 from vllm.v1.kv_offload.base import (

27 + Locality,

28 LookupResult,

29 + Medium,

30 OffloadingCounterMetadata,

31 OffloadingEvent,

32 OffloadKey,

@@ -29,6 +34,8 @@ from vllm.v1.kv_offload.base import (

34 ReqContext,

35 RequestOffloadingContext,

36 ScheduleEndContext,

37 + TierFilter,

38 + TierMatcher,

39 make_offload_key,

40 )

41 from vllm.v1.kv_offload.tiering.base import (

@@ -245,9 +252,9 @@ class TestTieringOffloadingManager:

252 self.manager.on_new_request(req_context)

253

254 def test_take_events_aggregates_tier_owned_events(self, manager_setup):

248 - primary_event = OffloadingEvent(to_keys([1]), "CPU", removed=False)

249 - secondary_event1 = OffloadingEvent(to_keys([2]), "tier-1", removed=False)

250 - secondary_event2 = OffloadingEvent(to_keys([3]), "tier-2", removed=True)

255 + primary_event = OffloadingEvent(to_keys([1]), Medium.CPU, removed=False)

256 + secondary_event1 = OffloadingEvent(to_keys([2]), Medium.STORAGE, removed=False)

257 + secondary_event2 = OffloadingEvent(to_keys([3]), Medium.STORAGE, removed=True)

258

259 self.primary_tier.take_events = MagicMock(return_value=[primary_event])

260 self.secondary_tier1.take_events = MagicMock(return_value=[secondary_event1])

@@ -973,6 +980,56 @@ class TestTieringOffloadingManager:

980 self.secondary_tier2.drain_jobs.assert_called_once()

981 assert self.manager._transfer_jobs == {}

982

983 + @pytest.mark.parametrize(

984 + "load_tier_filter",

985 + [

986 + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)),

987 + TierFilter(matchers=()),

988 + ],

989 + ids=["non_matching_medium", "empty_no_load"],

990 + )

991 + def test_tier_filter_skips_filtered_secondary(

992 + self, manager_setup, load_tier_filter

993 + ):

994 + """Filter excluding secondary medium returns MISS from secondaries

995 + even when they hold the block; primary is unaffected."""

996 + blocks = to_keys(range(2))

997 + # Put one block in primary, one only in secondary

998 + self._start_request()

999 + self.manager.prepare_store(blocks[:1], _CTX)

1000 + self.manager.complete_store(blocks[:1], _CTX, success=True)

1001 + self.secondary_tier1.blocks[blocks[1]] = True

1002 +

1003 + # Secondaries have medium=CPU, so load_tier_filter skips them.

1004 + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup)

1005 +

1006 + ctx = ReqContext(req_id="r1", load_tier_filter=load_tier_filter)

1007 + assert self.manager.lookup(blocks[0], ctx) is LookupResult.HIT

1008 + assert self.manager.lookup(blocks[1], ctx) is LookupResult.MISS

1009 + self.secondary_tier1.lookup.assert_not_called()

1010 +

1011 + @pytest.mark.parametrize(

1012 + "load_tier_filter",

1013 + [

1014 + TierFilter.ALL,

1015 + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)),

1016 + TierFilter(matchers=(TierMatcher(),)),

1017 + ],

1018 + ids=["all", "explicit_cpu", "unconstrained_matcher"],

1019 + )

1020 + def test_tier_filter_allows_matching_secondary(

1021 + self, manager_setup, load_tier_filter

1022 + ):

1023 + """Filter that matches the secondary's medium allows lookup."""

1024 + blocks = to_keys(range(1))

1025 + self.secondary_tier1.blocks[blocks[0]] = True

1026 +

1027 + self.secondary_tier1.lookup = MagicMock(wraps=self.secondary_tier1.lookup)

1028 +

1029 + ctx = ReqContext(req_id="r2", load_tier_filter=load_tier_filter)

1030 + assert self.manager.lookup(blocks[0], ctx) is LookupResult.RETRY

1031 + self.secondary_tier1.lookup.assert_called()

1032 +

1033

1034 class TestTieringOffloadingWithoutSecondaryTiers:

1035 """Test TieringOffloadingManager with no secondary tiers (backward compat)."""

@@ -999,5 +1056,81 @@ class TestTieringOffloadingWithoutSecondaryTiers:

1056 assert count_hits(manager, blocks) == 3

1057

1058

1059 +@pytest.mark.parametrize(

1060 + "raw,expected",

1061 + [

1062 + (

1063 + [{"medium": "storage"}],

1064 + TierFilter(matchers=(TierMatcher(medium=Medium.STORAGE),)),

1065 + ),

1066 + (

1067 + [{"medium": "CPU"}],

1068 + TierFilter(matchers=(TierMatcher(medium=Medium.CPU),)),

1069 + ),

1070 + (

1071 + [{}],

1072 + TierFilter(matchers=(TierMatcher(),)),

1073 + ),

1074 + (

1075 + [{"medium": "storage", "locality": "local"}],

1076 + TierFilter(

1077 + matchers=(TierMatcher(medium=Medium.STORAGE, locality=Locality.LOCAL),)

1078 + ),

1079 + ),

1080 + (

1081 + [{"medium": "cpu"}, {"medium": "storage"}],

1082 + TierFilter(

1083 + matchers=(

1084 + TierMatcher(medium=Medium.CPU),

1085 + TierMatcher(medium=Medium.STORAGE),

1086 + )

1087 + ),

1088 + ),

1089 + (

1090 + [],

1091 + TierFilter(matchers=()),

1092 + ),

1093 + ],

1094 + ids=[

1095 + "medium_storage",

1096 + "medium_cpu_uppercase",

1097 + "unconstrained",

1098 + "with_locality",

1099 + "multiple_matchers",

1100 + "empty_list_deny_all",

1101 + ],

1102 +)

1103 +def test_parse_tier_filter_valid(raw, expected):

1104 + assert _parse_tier_filter(raw) == expected

1105 +

1106 +

1107 +@pytest.mark.parametrize(

1108 + "raw",

1109 + [

1110 + "not a list",

1111 + [{"medium": "unknown"}],

1112 + [{"locality": "nowhere"}],

1113 + ],

1114 + ids=["non_list", "invalid_medium", "invalid_locality"],

1115 +)

1116 +def test_parse_tier_filter_invalid_returns_all(raw):

1117 + assert _parse_tier_filter(raw) is TierFilter.ALL

1118 +

1119 +

1120 +def test_parse_tier_filter_skips_bad_entries():

1121 + result = _parse_tier_filter(

1122 + [

1123 + {"medium": "storage"},

1124 + "not a dict",

1125 + {"medium": "bogus"},

1126 + {"medium": "cpu"},

1127 + ]

1128 + )

1129 + assert result.matchers == (

1130 + TierMatcher(medium=Medium.STORAGE),

1131 + TierMatcher(medium=Medium.CPU),

1132 + )

1133 +

1134 +

1135 if __name__ == "__main__":

1136 pytest.main([__file__, "-v"])

vllm/distributed/kv_events.py

+1 -2

@@ -44,8 +44,7 @@ class KVCacheEvent(

45 MEDIUM_GPU = "GPU"

46 MEDIUM_CPU = "CPU"

47 -MEDIUM_FS = "FS"

48 -MEDIUM_OBJ = "OBJ"

47 +MEDIUM_STORAGE = "STORAGE"

50 class BlockStored(KVCacheEvent):

vllm/distributed/kv_transfer/kv_connector/v1/offloading/events.py

+17 -5

@@ -19,7 +19,13 @@ from collections.abc import Iterable

19 from dataclasses import dataclass

20 from typing import TYPE_CHECKING, Any, NamedTuple

22 -from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent

22 +from vllm.distributed.kv_events import (

23 + MEDIUM_CPU,

24 + MEDIUM_STORAGE,

25 + BlockRemoved,

26 + BlockStored,

27 + KVCacheEvent,

28 +)

29 from vllm.logger import init_logger

30 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash

31 from vllm.v1.kv_cache_interface import (

@@ -28,6 +34,7 @@ from vllm.v1.kv_cache_interface import (

34 get_kv_cache_spec_sliding_window,

35 )

36 from vllm.v1.kv_offload.base import (

37 + Medium,

38 OffloadingEvent,

39 OffloadingKVEventsConfig,

40 OffloadKey,

@@ -43,6 +50,11 @@ if TYPE_CHECKING:

51 logger = init_logger(__name__)

53 +_MEDIUM_TO_EVENT_STR: dict[Medium, str] = {

54 + Medium.CPU: MEDIUM_CPU,

55 + Medium.STORAGE: MEDIUM_STORAGE,

56 +}

57 +

59 class OffloadingEventGroupSpec(NamedTuple):

60 kv_cache_spec_kind: str | None

@@ -218,7 +230,7 @@ class OffloadingEventsTracker:

230 def _placeholder_stored(

231 self,

232 key: OffloadKey,

221 - medium: str,

233 + medium: Medium,

234 locality: str | None,

235 ) -> BlockStored:

236 return BlockStored(

@@ -229,7 +241,7 @@ class OffloadingEventsTracker:

241 token_ids=[],

242 lora_id=None,

243 block_size=0,

232 - medium=medium,

244 + medium=_MEDIUM_TO_EVENT_STR[medium],

245 lora_name=None,

246 group_idx=get_offload_group_idx(key),

247 locality=locality,

@@ -267,7 +279,7 @@ class OffloadingEventsTracker:

279 token_ids=list(meta.token_ids),

280 block_size=meta.block_size,

281 lora_id=meta.lora_id,

270 - medium=event.medium,

282 + medium=_MEDIUM_TO_EVENT_STR[event.medium],

283 lora_name=meta.lora_name,

284 extra_keys=(

285 list(meta.extra_keys) if meta.extra_keys is not None else None

@@ -308,7 +320,7 @@ class OffloadingEventsTracker:

320 for group_idx, hashes in by_group.items():

321 yield BlockRemoved(

322 block_hashes=hashes,

311 - medium=event.medium,

323 + medium=_MEDIUM_TO_EVENT_STR[event.medium],

324 group_idx=group_idx,

325 locality=locality,

326 )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+60 -1

@@ -39,7 +39,9 @@ from vllm.v1.kv_cache_interface import (

39 )

40 from vllm.v1.kv_offload.base import (

41 GPULoadStoreSpec,

42 + Locality,

43 LookupResult,

44 + Medium,

45 OffloadingManager,

46 OffloadingSpec,

47 OffloadKey,

@@ -47,6 +49,8 @@ from vllm.v1.kv_offload.base import (

49 ReqContext,

50 RequestOffloadingContext,

51 ScheduleEndContext,

52 + TierFilter,

53 + TierMatcher,

54 make_offload_key,

55 )

56 from vllm.v1.outputs import KVConnectorOutput

@@ -54,6 +58,10 @@ from vllm.v1.request import Request, RequestStatus

59 logger = init_logger(__name__)

61 +KV_LOAD_TIERS_KEY = "kv_load_tiers"

62 +MATCHER_MEDIUM_KEY = "medium"

63 +MATCHER_LOCALITY_KEY = "locality"

64 +

66 @dataclass(slots=True)

67 class TransferJobStatus:

@@ -376,10 +384,61 @@ class RequestOffloadState:

384 )

385

386

387 +def _parse_tier_filter(raw: Any) -> TierFilter:

388 + """Parse raw kv_transfer_params tier matchers into a TierFilter."""

389 + if not isinstance(raw, list):

390 + logger.warning(

391 + "_parse_tier_filter: expected list, got %s; ignoring",

392 + type(raw).__name__,

393 + )

394 + return TierFilter.ALL

395 + matchers: list[TierMatcher] = []

396 + for entry in raw:

397 + if not isinstance(entry, dict):

398 + logger.warning("_parse_tier_filter: entry is not a dict; skipping")

399 + continue

400 + medium: Medium | None = None

401 + locality: Locality | None = None

402 + raw_medium = entry.get(MATCHER_MEDIUM_KEY)

403 + if raw_medium is not None:

404 + try:

405 + medium = Medium(raw_medium.upper())

406 + except (ValueError, AttributeError):

407 + logger.warning(

408 + "_parse_tier_filter: unknown medium %r; skipping entry",

409 + raw_medium,

410 + )

411 + continue

412 + raw_locality = entry.get(MATCHER_LOCALITY_KEY)

413 + if raw_locality is not None:

414 + try:

415 + locality = Locality(raw_locality.upper())

416 + except (ValueError, AttributeError):

417 + logger.warning(

418 + "_parse_tier_filter: unknown locality %r; skipping entry",

419 + raw_locality,

420 + )

421 + continue

422 + matchers.append(TierMatcher(medium=medium, locality=locality))

423 + if not matchers:

424 + if not raw: # input was [] — user explicitly wants nothing

425 + return TierFilter(matchers=())

426 + # all entries were invalid — fall back to ALL

427 + return TierFilter.ALL

428 + return TierFilter(matchers=tuple(matchers))

429 +

430 +

431 def _create_req_context(req: Request) -> ReqContext:

432 + params = req.kv_transfer_params

433 + load_filter = TierFilter.ALL

434 + if params:

435 + raw = params.get(KV_LOAD_TIERS_KEY)

436 + if raw is not None:

437 + load_filter = _parse_tier_filter(raw)

438 return ReqContext(

439 req_id=req.request_id,

382 - kv_transfer_params=req.kv_transfer_params,

440 + kv_transfer_params=params,

441 + load_tier_filter=load_filter,

442 )

443

444

vllm/v1/kv_offload/base.py

+46 -9

@@ -8,7 +8,7 @@ from abc import ABC, abstractmethod

8 from collections.abc import Collection, Iterable, Sequence

9 from dataclasses import dataclass, field

10 from enum import Enum, auto

11 -from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeVar

11 +from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, NewType, TypeVar

13 import numpy as np

14 import torch

@@ -48,10 +48,54 @@ def get_offload_group_idx(key: OffloadKey) -> int:

48 _T = TypeVar("_T")

51 +class Medium(Enum):

52 + """Storage medium of an offloading tier."""

53 +

54 + CPU = "CPU"

55 + STORAGE = "STORAGE"

56 +

57 +

58 +class Locality(Enum):

59 + """Locality of a tier's storage relative to the publishing instance."""

60 +

61 + LOCAL = "LOCAL"

62 + REMOTE = "REMOTE"

63 +

64 +

65 +class TierMatcher(NamedTuple):

66 + medium: Medium | None = None

67 + locality: Locality | None = None

68 +

69 + def matches(self, medium: Medium | None, locality: Locality | None) -> bool:

70 + medium_matches = self.medium is None or medium is None or self.medium == medium

71 + locality_matches = (

72 + self.locality is None or locality is None or self.locality == locality

73 + )

74 + return medium_matches and locality_matches

75 +

76 +

77 +@dataclass(frozen=True)

78 +class TierFilter:

79 + """Per-request filter controlling which tiers participate."""

80 +

81 + matchers: tuple[TierMatcher, ...] = ()

82 +

83 + ALL: ClassVar["TierFilter"]

84 +

85 + def allows(self, medium: Medium | None, locality: Locality | None) -> bool:

86 + if self is TierFilter.ALL:

87 + return True

88 + return any(m.matches(medium, locality) for m in self.matchers)

89 +

90 +

91 +TierFilter.ALL = TierFilter(matchers=(TierMatcher(),))

92 +

93 +

94 @dataclass

95 class ReqContext:

96 req_id: str

97 kv_transfer_params: dict[str, Any] | None = None

98 + load_tier_filter: TierFilter = TierFilter.ALL

99 # Per-request scratch space keyed by value type, so a tier can parse

100 # kv_transfer_params once (in on_new_request) and read the result back

101 # on later calls for the same request.

@@ -110,17 +154,10 @@ class PrepareStoreOutput:

154 evicted_keys: list[OffloadKey]

155

156

113 -class Locality(Enum):

114 - """Locality of a tier's storage relative to the publishing instance."""

115 -

116 - LOCAL = "LOCAL"

117 - REMOTE = "REMOTE"

118 -

119 -

157 @dataclass

158 class OffloadingEvent:

159 keys: list[OffloadKey]

123 - medium: str

160 + medium: Medium

161 # True if blocks are removed, False if stored

162 removed: bool

163 locality: Locality | None = None

vllm/v1/kv_offload/cpu/manager.py

+2 -2

@@ -6,13 +6,13 @@ from typing import Literal

7 from typing_extensions import override

9 -from vllm.distributed.kv_events import MEDIUM_CPU

9 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (

10 OffloadingConnectorStats,

11 )

12 from vllm.v1.kv_offload.base import (

13 LoadStoreSpec,

14 LookupResult,

15 + Medium,

16 OffloadingEvent,

17 OffloadingManager,

18 OffloadKey,

@@ -52,7 +52,7 @@ class CPUOffloadingManager(OffloadingManager):

52 store_threshold: int = 1,

53 max_tracker_size: int = 64_000,

54 ):

55 - self.medium: str = MEDIUM_CPU

55 + self.medium: Medium = Medium.CPU

56 self._num_blocks: int = num_blocks

57 self._num_allocated_blocks: int = 0

58 self._free_list: list[int] = []

vllm/v1/kv_offload/tiering/base.py

+7 -1

@@ -7,12 +7,14 @@ Abstract interfaces and data types for the secondary tiering layer.

7 from abc import ABC, abstractmethod

8 from collections.abc import Collection, Iterable

9 from dataclasses import dataclass

10 -from typing import TYPE_CHECKING, Any

10 +from typing import TYPE_CHECKING, Any, ClassVar

12 import numpy as np

14 from vllm.v1.kv_offload.base import (

15 + Locality,

16 LookupResult,

17 + Medium,

18 OffloadingEvent,

19 OffloadingMetricMetadata,

20 OffloadKey,

@@ -27,6 +29,7 @@ if TYPE_CHECKING:

29 )

30 from vllm.v1.kv_offload.base import OffloadingSpec

32 +

33 # Type alias for job IDs used in async transfer tracking

34 JobId = int

@@ -108,6 +111,8 @@ class SecondaryTierManager(ABC):

111 async jobs; get_finished_jobs() polls for completion.

112 """

113

114 + medium: ClassVar[Medium | None] = None

115 +

116 def __init__(

117 self,

118 offloading_spec: "OffloadingSpec",

@@ -124,6 +129,7 @@ class SecondaryTierManager(ABC):

129 self._offloading_spec = offloading_spec

130 self._primary_kv_view: memoryview = primary_kv_view

131 self.tier_type = tier_type

132 + self.locality: Locality | None = None

133

134 @abstractmethod

135 def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult:

vllm/v1/kv_offload/tiering/example/manager.py

+4 -1

@@ -11,12 +11,13 @@ TieringOffloadingManager without requiring actual storage or network backends.

12 import logging

13 from collections.abc import Iterable

14 -from typing import TYPE_CHECKING

14 +from typing import TYPE_CHECKING, ClassVar

16 from typing_extensions import override

18 from vllm.v1.kv_offload.base import (

19 LookupResult,

20 + Medium,

21 OffloadKey,

22 ReqContext,

23 RequestOffloadingContext,

@@ -42,6 +43,8 @@ class ExampleSecondaryTierManager(SecondaryTierManager):

43 - Completes transfers immediately (synchronous)

44 """

46 + medium: ClassVar[Medium] = Medium.CPU

47 +

48 def __init__(

49 self,

50 offloading_spec: "OffloadingSpec",

vllm/v1/kv_offload/tiering/fs/manager.py

+2 -2

@@ -30,11 +30,11 @@ except ImportError:

31 from typing_extensions import override

33 -from vllm.distributed.kv_events import MEDIUM_FS

33 from vllm.logger import init_logger

34 from vllm.v1.kv_offload.base import (

35 Locality,

36 LookupResult,

37 + Medium,

38 OffloadingEvent,

39 OffloadKey,

40 ReqContext,

@@ -100,7 +100,7 @@ class FileSystemTierManager(SecondaryTierManager):

100 content.

101 """

102

103 - medium: ClassVar[str] = MEDIUM_FS

103 + medium: ClassVar[Medium] = Medium.STORAGE

104

105 def __init__(

106 self,

vllm/v1/kv_offload/tiering/manager.py

@@ -324,6 +324,8 @@ class TieringOffloadingManager(OffloadingManager):

324 for tier in self.secondary_tiers:

325 if tier is exclude_tier:

326 continue

327 + if not req_context.load_tier_filter.allows(tier.medium, tier.locality):

328 + continue

329 result = tier.lookup(key, req_context)

330 if result is LookupResult.HIT:

331 promoted = self._initiate_promotion(tier, key, req_context)

vllm/v1/kv_offload/tiering/obj/manager.py

+2 -2

@@ -7,13 +7,13 @@ import time

7 from collections.abc import Iterable

8 from typing import TYPE_CHECKING, ClassVar, NamedTuple

10 -from vllm.distributed.kv_events import MEDIUM_OBJ

10 from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent

11 from vllm.distributed.nixl_utils import nixl_agent_config

12 from vllm.logger import init_logger

13 from vllm.v1.kv_offload.base import (

14 Locality,

15 LookupResult,

16 + Medium,

17 OffloadingEvent,

18 OffloadKey,

19 ReqContext,

@@ -98,7 +98,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

98 primary tier. Object keys are formed as ``{prefix}/{hash_shard}/{hash}.bin``.

99 """

100

101 - medium: ClassVar[str] = MEDIUM_OBJ

101 + medium: ClassVar[Medium] = Medium.STORAGE

102

103 def __init__(

104 self,