@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f774ba028a

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 22:30 UTC

[kv_offload+HMA][4/N]: Support sliding window lookup (#36645)

Signed-off-by: Or Ozeri <oro@il.ibm.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

Or Ozeri committed Apr 20, 2026 at 12:53 UTC f774ba028afe08fab1ca96ecf7c0b40c0fa45357

7 files changed +269 -78

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

+151 -9

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable

4 +from unittest.mock import MagicMock

6 import pytest

@@ -10,8 +11,16 @@ from tests.v1.kv_connector.unit.offloading_connector.utils import (

11 )

12 from tests.v1.kv_connector.unit.utils import EOS_TOKEN_ID

13 from vllm.distributed.kv_events import BlockRemoved, BlockStored

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

15 + OffloadingConnectorScheduler,

16 +)

17 from vllm.v1.core.kv_cache_utils import BlockHash

14 -from vllm.v1.kv_offload.abstract import OffloadingEvent

18 +from vllm.v1.kv_offload.abstract import (

19 + OffloadingEvent,

20 + OffloadingManager,

21 + ReqContext,

22 + get_offload_block_hash,

23 +)

24 from vllm.v1.request import RequestStatus

@@ -105,8 +114,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):

114 lambda keys, req_context: generate_store_output([])

115 )

116 runner.run(decoded_tokens=[EOS_TOKEN_ID])

108 - runner.manager.lookup.assert_called()

109 - assert len(list(runner.manager.lookup.call_args.args[0])) == 1

117 + runner.manager.lookup.assert_called_once()

118

119 # single block lookup with a hit

120 runner.scheduler.reset_prefix_cache()

@@ -114,7 +122,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):

122 runner.manager.prepare_store.side_effect = (

123 lambda keys, req_context: generate_store_output([])

124 )

117 - runner.manager.lookup.return_value = 1

125 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

126 runner.run(

127 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(0, 1, 2)

128 )

@@ -126,7 +134,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):

134 runner.manager.prepare_store.side_effect = (

135 lambda keys, req_context: generate_store_output([])

136 )

129 - runner.manager.lookup.return_value = 1

137 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

138 runner.run(

139 decoded_tokens=[EOS_TOKEN_ID], expected_loaded_gpu_block_indexes=(3, 4, 5)

140 )

@@ -210,7 +218,7 @@ def test_request_preemption(request_runner, async_scheduling: bool):

218

219 # request should now return from preemption

220 # re-load [0, ..., 8] from the CPU and store [9, 10, 11]

213 - runner.manager.lookup.return_value = 3

221 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 3

222 runner.manager.prepare_store.side_effect = (

223 lambda keys, req_context: generate_store_output(keys)

224 )

@@ -251,7 +259,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:

259 # start a request to load the first block, but don't complete

260 runner.scheduler.reset_prefix_cache()

261 runner.new_request(token_ids=[0] * offloaded_block_size)

254 - runner.manager.lookup.return_value = 1

262 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

263 runner.run(

264 decoded_tokens=[],

265 complete_transfers=False,

@@ -263,7 +271,7 @@ def test_concurrent_lookups_of_the_same_prefix(request_runner, async_scheduling:

271

272 # start a new request to load the same first block

273 runner.new_request(token_ids=[0] * offloaded_block_size)

266 - runner.manager.lookup.return_value = 1

274 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

275 runner.run(

276 decoded_tokens=[],

277 complete_transfers=False,

@@ -311,7 +319,7 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):

319 # start a request to load the first block, but don't complete

320 runner.scheduler.reset_prefix_cache()

321 runner.new_request(token_ids=[0] * offloaded_block_size)

314 - runner.manager.lookup.return_value = 1

322 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 1

323 runner.run(

324 decoded_tokens=[],

325 complete_transfers=False,

@@ -336,3 +344,137 @@ def test_abort_loading_requests(request_runner, async_scheduling: bool):

344

345 # assert request is deleted

346 assert req_id not in runner.scheduler.requests

347 +

348 +

349 +# ---------------------------------------------------------------------------

350 +# Unit tests for _maximal_prefix_lookup / _sliding_window_lookup

351 +# ---------------------------------------------------------------------------

352 +

353 +

354 +def _make_scheduler_with_lookup(

355 + lookup_results: dict[int, bool | None],

356 +) -> OffloadingConnectorScheduler:

357 + """Create an OffloadingConnectorScheduler with a mocked manager.lookup."""

358 + manager = MagicMock(spec=OffloadingManager)

359 + manager.lookup.side_effect = lambda key, req_context: lookup_results.get(

360 + int(get_offload_block_hash(key).decode()), False

361 + )

362 +

363 + scheduler = object.__new__(OffloadingConnectorScheduler)

364 + scheduler.manager = manager

365 + return scheduler

366 +

367 +

368 +_EMPTY_REQ_CTX = ReqContext()

369 +

370 +

371 +class TestMaximalPrefixLookup:

372 + def test_all_hit(self):

373 + sched = _make_scheduler_with_lookup({1: True, 2: True})

374 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2

375 +

376 + def test_all_miss(self):

377 + sched = _make_scheduler_with_lookup({})

378 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

379 +

380 + def test_partial_prefix(self):

381 + sched = _make_scheduler_with_lookup({1: True, 2: True})

382 + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2

383 +

384 + def test_miss_then_hit(self):

385 + sched = _make_scheduler_with_lookup({2: True})

386 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

387 +

388 + def test_single_hit(self):

389 + sched = _make_scheduler_with_lookup({1: True})

390 + assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1

391 +

392 + def test_empty(self):

393 + sched = _make_scheduler_with_lookup({})

394 + assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0

395 +

396 + def test_none_defers(self):

397 + sched = _make_scheduler_with_lookup({1: None, 2: True})

398 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None

399 +

400 + def test_none_after_hit_defers(self):

401 + sched = _make_scheduler_with_lookup({1: True, 2: None})

402 + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None

403 +

404 + def test_none_stops_at_miss(self):

405 + """None is treated as hit for iteration, but miss stops the scan."""

406 + sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True})

407 + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None

408 + # lookup should have been called for blocks 1 and 2 (stops at miss)

409 + assert sched.manager.lookup.call_count == 2

410 +

411 +

412 +class TestSlidingWindowLookup:

413 + def test_all_hit_exact_window(self):

414 + sched = _make_scheduler_with_lookup({1: True, 2: True})

415 + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2

416 +

417 + def test_all_miss(self):

418 + sched = _make_scheduler_with_lookup({})

419 + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0

420 +

421 + def test_window_at_end(self):

422 + sched = _make_scheduler_with_lookup({2: True, 3: True})

423 + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3

424 +

425 + def test_window_in_middle(self):

426 + sched = _make_scheduler_with_lookup({2: True, 3: True})

427 + assert (

428 + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3

429 + )

430 +

431 + def test_no_full_window_falls_back_to_prefix(self):

432 + sched = _make_scheduler_with_lookup({1: True, 2: True})

433 + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2

434 +

435 + def test_single_block_window(self):

436 + sched = _make_scheduler_with_lookup({2: True, 3: True})

437 + assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3

438 +

439 + def test_gap_resets_consecutive(self):

440 + sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True})

441 + # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3]

442 + assert (

443 + sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX)

444 + == 3

445 + )

446 +

447 + def test_window_prefers_rightmost(self):

448 + sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True})

449 + # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4

450 + # scans right-to-left, finds [4,5] first

451 + assert (

452 + sched._sliding_window_lookup(to_keys([1, 2, 3, 4, 5]), 2, _EMPTY_REQ_CTX)

453 + == 5

454 + )

455 +

456 + def test_prefix_fallback_with_gap(self):

457 + sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True})

458 + # window of 4 not found contiguously (gap at 1)

459 + assert (

460 + sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX)

461 + == 1

462 + )

463 +

464 + def test_empty(self):

465 + sched = _make_scheduler_with_lookup({})

466 + assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0

467 +

468 + def test_none_defers(self):

469 + sched = _make_scheduler_with_lookup({1: True, 2: None})

470 + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None

471 +

472 + def test_none_with_full_window_still_defers(self):

473 + """Even if a real window is found after a None, result is deferred."""

474 + # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window

475 + # but block 3 was None so defer_lookup is set

476 + sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True})

477 + assert (

478 + sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX)

479 + is None

480 + )

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

+12 -7

@@ -56,8 +56,12 @@ from vllm.v1.request import Request

56 from vllm.v1.structured_output import StructuredOutputManager

59 -def to_keys(int_ids: list[int]) -> list[OffloadKey]:

60 - return [make_offload_key(str(i).encode(), 0) for i in int_ids]

59 +def to_key(int_hash: int) -> OffloadKey:

60 + return make_offload_key(str(int_hash).encode(), 0)

61 +

62 +

63 +def to_keys(int_hashes: list[int]) -> list[OffloadKey]:

64 + return [to_key(i) for i in int_hashes]

67 class MockLoadStoreSpec(LoadStoreSpec):

@@ -116,6 +120,7 @@ class MockOffloadingSpec(OffloadingSpec):

120 self.manager = MagicMock(spec=OffloadingManager)

121 self.manager.lookup.return_value = 0

122 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys)

123 + self.manager.lookup.return_value = False

124 self.handler = MockOffloadingHandler()

125

126 def get_manager(self) -> OffloadingManager:

@@ -228,14 +233,14 @@ class RequestRunner:

233 self.scheduler_connector: OffloadingConnector = scheduler_connector

234

235 # extract mocked OffloadingManager of scheduler connector

231 - connector_scheduler = scheduler_connector.connector_scheduler

232 - assert connector_scheduler is not None

233 - manager = connector_scheduler.manager

236 + self.connector_scheduler = scheduler_connector.connector_scheduler

237 + assert self.connector_scheduler is not None

238 + manager = self.connector_scheduler.manager

239 assert isinstance(manager, MagicMock)

240 self.manager: MagicMock = manager

241

237 - assert len(connector_scheduler.config.kv_group_configs) == 1

238 - kv_group_config = connector_scheduler.config.kv_group_configs[0]

242 + assert len(self.connector_scheduler.config.kv_group_configs) == 1

243 + kv_group_config = self.connector_scheduler.config.kv_group_configs[0]

244 assert kv_group_config.gpu_block_size == gpu_block_size

245 assert kv_group_config.offloaded_block_size == offloaded_block_size

246

tests/v1/kv_offload/test_cpu_manager.py

+36 -20

@@ -35,8 +35,12 @@ class ExpectedPrepareStoreOutput:

35 evicted_keys: list[int]

38 -def to_keys(int_ids: list[int]) -> list[OffloadKey]:

39 - return [make_offload_key(str(i).encode(), 0) for i in int_ids]

38 +def to_key(int_hash: int) -> OffloadKey:

39 + return make_offload_key(str(int_hash).encode(), 0)

40 +

41 +

42 +def to_keys(int_hashes: list[int]) -> list[OffloadKey]:

43 + return [to_key(i) for i in int_hashes]

46 def verify_store_output(

@@ -136,7 +140,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):

140 manager.complete_store(to_keys([2, 3, 4, 5]))

141

142 # block 2 must still be present in the cache

139 - assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1

143 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

144

145

146 def test_cpu_manager():

@@ -160,7 +164,8 @@ def test_cpu_manager():

164 )

165

166 # lookup [1, 2] -> not ready

163 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

167 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

168 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

169

170 # no events so far

171 assert list(cpu_manager.take_events()) == []

@@ -170,9 +175,9 @@ def test_cpu_manager():

175 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

176

177 # lookup [1, 2]

173 - assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1

174 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2

175 - assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2

178 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True

179 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

180 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

181

182 # prepare store [2, 3, 4, 5] -> evicts [1]

183 prepare_store_output = cpu_manager.prepare_store(

@@ -196,6 +201,14 @@ def test_cpu_manager():

201 # complete store [2, 3, 4, 5]

202 cpu_manager.complete_store(to_keys([2, 3, 4, 5]))

203

204 + # lookup (now that we have [2, 3, 4, 5])

205 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

206 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

207 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True

208 + assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True

209 + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True

210 + assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False

211 +

212 # prepare load [2, 3]

213 prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX)

214 verify_load_output(prepare_load_output, [1, 2])

@@ -238,8 +251,8 @@ def test_cpu_manager():

251 cpu_manager.complete_store(to_keys([7, 9]), success=False)

252

253 # assert [7] is still stored, but [9] is not

241 - assert cpu_manager.lookup(to_keys([7]), _EMPTY_REQ_CTX) == 1

242 - assert cpu_manager.lookup(to_keys([9]), _EMPTY_REQ_CTX) == 0

254 + assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True

255 + assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False

256

257 verify_events(

258 cpu_manager.take_events(),

@@ -284,7 +297,8 @@ class TestARCPolicy:

297 )

298

299 # lookup [1, 2] -> not ready

287 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

300 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

301 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

302

303 # no events so far

304 assert list(cpu_manager.take_events()) == []

@@ -294,9 +308,9 @@ class TestARCPolicy:

308 verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))

309

310 # lookup [1, 2]

297 - assert cpu_manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1

298 - assert cpu_manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2

299 - assert cpu_manager.lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2

311 + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True

312 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

313 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

314

315 # blocks should be in T1 (recent)

316 assert len(arc_policy.t1) == 2

@@ -500,7 +514,7 @@ class TestARCPolicy:

514 cpu_manager.complete_store(to_keys([5]), success=False)

515

516 # block 5 should not be in cache

503 - assert cpu_manager.lookup(to_keys([5]), _EMPTY_REQ_CTX) == 0

517 + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False

518 # block 5 should not be in T1 or T2

519 assert to_keys([5])[0] not in arc_policy.t1

520 assert to_keys([5])[0] not in arc_policy.t2

@@ -541,8 +555,8 @@ class TestARCPolicy:

555 cpu_manager.complete_store(to_keys([6]))

556

557 # verify blocks 2, 3 (in T2) are still present

544 - assert cpu_manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 1

545 - assert cpu_manager.lookup(to_keys([3]), _EMPTY_REQ_CTX) == 1

558 + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True

559 + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True

560

561 # verify events

562 events = list(cpu_manager.take_events())

@@ -562,7 +576,8 @@ def test_filter_reused_manager():

576 )

577

578 # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet

565 - assert manager.lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0

579 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

580 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

581

582 # prepare store [1, 2] -> should be filtered

583 prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

@@ -570,7 +585,7 @@ def test_filter_reused_manager():

585 assert prepare_store_output.keys_to_store == []

586

587 # Lookup [1] -> 2nd time, eligible now

573 - assert manager.lookup(to_keys([1]), _EMPTY_REQ_CTX) == 0

588 + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False

589

590 # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered

591 prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX)

@@ -579,12 +594,13 @@ def test_filter_reused_manager():

594

595 # Lookup [3, 4] -> 1st time

596 # (evicts [2] from tracker since max_size is 3 and tracker has [1])

582 - assert manager.lookup(to_keys([3, 4]), _EMPTY_REQ_CTX) == 0

597 + assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False

598 + assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False

599 # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4])

600 assert to_keys([2])[0] not in manager.counts

601

602 # Lookup [2] again -> (this adds [2] back to the tracker as 1st time)

587 - assert manager.lookup(to_keys([2]), _EMPTY_REQ_CTX) == 0

603 + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False

604 # Verify [2] was re-added with count=1 (not eligible yet)

605 assert manager.counts.get(to_keys([2])[0]) == 1

606

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

+48 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections import defaultdict

4 -from collections.abc import Iterable

4 +from collections.abc import Iterable, Sequence

5 from dataclasses import dataclass, field

6 from itertools import islice

7 from typing import Any, NamedTuple

@@ -132,6 +132,49 @@ class OffloadingConnectorScheduler:

132 self._reqs_being_stored = defaultdict[ReqId, set[OffloadKey]](set)

133 self._reqs_being_loaded = defaultdict[ReqId, set[OffloadKey]](set)

134

135 + def _maximal_prefix_lookup(

136 + self, keys: Iterable[OffloadKey], req_context: ReqContext

137 + ) -> int | None:

138 + """Find the length of the maximal prefix of offloaded blocks."""

139 + hit_count = 0

140 + defer_lookup = False

141 + for key in keys:

142 + result = self.manager.lookup(key, req_context)

143 + if result is None:

144 + defer_lookup = True

145 + # continue lookup to allow manager to kick-off async lookups

146 + # for all blocks (until a miss is detected)

147 + result = True

148 + if not result:

149 + break

150 + hit_count += 1

151 + return hit_count if not defer_lookup else None

152 +

153 + def _sliding_window_lookup(

154 + self,

155 + keys: Sequence[OffloadKey],

156 + sliding_window_size: int,

157 + req_context: ReqContext,

158 + ) -> int | None:

159 + """Find the maximal ending position of consecutive offloaded blocks

160 + within a sliding window."""

161 + defer_lookup = False

162 + consecutive_hits = 0

163 + for idx in range(len(keys) - 1, -1, -1):

164 + result = self.manager.lookup(keys[idx], req_context)

165 + if result is None:

166 + defer_lookup = True

167 + # continue lookup to allow manager to kick-off async lookups

168 + # for all blocks (until a hit is detected)

169 + result = False

170 + if not result:

171 + consecutive_hits = 0

172 + else:

173 + consecutive_hits += 1

174 + if consecutive_hits == sliding_window_size:

175 + return idx + sliding_window_size if not defer_lookup else None

176 + return consecutive_hits if not defer_lookup else None

177 +

178 def get_num_new_matched_tokens(

179 self, request: Request, num_computed_tokens: int

180 ) -> tuple[int | None, bool]:

@@ -184,9 +227,10 @@ class OffloadingConnectorScheduler:

227 return 0, False

228

229 start_block_idx = num_computed_tokens // group_config.offloaded_block_size

187 - hits = self.manager.lookup(

188 - offload_keys[start_block_idx:],

189 - req_status.req_context,

230 + # Full attention relays on all previous KV cache blocks.

231 + # Thus, we search for a maximal prefix of KV cache which are all cached.

232 + hits = self._maximal_prefix_lookup(

233 + offload_keys[start_block_idx:], req_status.req_context

234 )

235 if hits is None:

236 # indicates a lookup that should be tried later

vllm/v1/kv_offload/abstract.py

+7 -13

@@ -7,8 +7,7 @@ This class runs in the scheduler, tracks which blocks are offloaded

7 and their address.

9 The class provides the following primitives:

10 - lookup() - find the length of the maximal series of blocks,

11 - starting from the first one, that are all offloaded.

10 + lookup() - check whether a single block is offloaded and ready.

11 prepare_load() - prepare given blocks to be read.

12 The given blocks will be protected from eviction.

13 This function returns a LoadSpec which encapsulates

@@ -91,23 +90,18 @@ class OffloadingEvent:

91 class OffloadingManager(ABC):

92 @abstractmethod

94 - def lookup(

95 - self,

96 - keys: Iterable[OffloadKey],

97 - req_context: ReqContext,

98 - ) -> int | None:

93 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

94 """

100 - Finds the length of the maximal series of blocks, starting from the

101 - first one, that are all offloaded.

95 + Checks whether a single block is offloaded and ready to be read.

97 Args:

104 - keys: the keys identifying the blocks to lookup.

98 + key: the key identifying the block to lookup.

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

100

101 Returns:

108 - An integer representing the maximal number of blocks that

109 - are currently offloaded, or None if the lookup should be retried

110 - later. Returning None will delay the request handling by the vLLM

102 + True if the block is offloaded and ready, False if not,

103 + or None if the lookup should be retried later.

104 + Returning None will delay the request handling by the vLLM

105 scheduler.

106 """

107 pass

vllm/v1/kv_offload/cpu/manager.py

+3 -12

@@ -84,18 +84,9 @@ class CPUOffloadingManager(OffloadingManager):

85 # --- OffloadingManager interface ---

87 - def lookup(

88 - self,

89 - keys: Iterable[OffloadKey],

90 - req_context: ReqContext,

91 - ) -> int | None:

92 - hit_count = 0

93 - for key in keys:

94 - block = self._policy.get(key)

95 - if block is None or not block.is_ready:

96 - break

97 - hit_count += 1

98 - return hit_count

87 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

88 + block = self._policy.get(key)

89 + return block is not None and block.is_ready

91 def prepare_load(

92 self,

vllm/v1/kv_offload/reuse_manager.py

+12 -13

@@ -27,8 +27,9 @@ class FilterReusedOffloadingManager(OffloadingManager):

27 All methods are delegated to the *backing* manager. Two methods are

28 intercepted:

30 - * ``lookup`` — records each visited key in an internal LRU counter.

30 * ``prepare_store`` — filters out keys that have not yet

31 + * ``lookup`` — records the visited key in an internal LRU

32 + counter, then delegates to the backing manager.

33 crossed the threshold *before* calling the backing

34 ``prepare_store``.

@@ -66,18 +67,16 @@ class FilterReusedOffloadingManager(OffloadingManager):

67 # Intercepted methods

68 # ------------------------------------------------------------------

69 - def lookup(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> int | None:

70 - """Record each key, then delegate lookup to backing manager."""

71 - keys = list(keys)

72 - for key in keys:

73 - if key in self.counts:

74 - self.counts.move_to_end(key)

75 - self.counts[key] += 1

76 - else:

77 - if len(self.counts) >= self.max_tracker_size:

78 - self.counts.popitem(last=False) # evict LRU

79 - self.counts[key] = 1

80 - return self._backing.lookup(keys, req_context)

70 + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None:

71 + """Record the key, then delegate lookup to backing manager."""

72 + if key in self.counts:

73 + self.counts.move_to_end(key)

74 + self.counts[key] += 1

75 + else:

76 + if len(self.counts) >= self.max_tracker_size:

77 + self.counts.popitem(last=False) # evict LRU

78 + self.counts[key] = 1

79 + return self._backing.lookup(key, req_context)

81 def prepare_store(

82 self, keys: Iterable[OffloadKey], req_context: ReqContext