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

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 9, 23:00 UTC

OffloadingConnector: Prevent redundant loads (#29087)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed Jan 21, 2026 at 03:15 UTC 7013e9ac8f65fc64bae3fe5bfab4bcd7b8ef8352

5 files changed +109 -8

tests/v1/kv_connector/unit/test_offloading_connector.py

+67 -1

@@ -213,7 +213,6 @@ class RequestRunner:

213 )

214

215 def new_request(self, token_ids: list[int]):

216 - assert not self.scheduler.requests

216 self.req_id += 1

217

218 req = Request(

@@ -338,11 +337,20 @@ class RequestRunner:

337 token_id=token_id or 0,

338 )

339

340 + prev_token_id = token_id

341 if self.scheduler.running:

342 token_id = next(tokens_iter, None)

343

344 self.scheduler.update_from_output(scheduler_output, model_runner_output)

345

346 + if (

347 + prev_token_id is EOS_TOKEN_ID

348 + and prev_token_id != token_id

349 + and self.scheduler.requests

350 + ):

351 + # continue for one more step to allow offloading to kick off

352 + continue

353 +

354 if token_id is None:

355 break

356

@@ -651,3 +659,61 @@ def test_request_preemption(request_runner):

659 decoded_tokens=[EOS_TOKEN_ID],

660 expected_stored_gpu_block_indexes=(9, 10, 11),

661 )

662 +

663 +

664 +def test_concurrent_lookups_of_the_same_prefix(request_runner):

665 + offloaded_block_size = 12

666 + gpu_block_size = 4

667 + num_gpu_blocks = 100

668 +

669 + runner = request_runner(

670 + offloaded_block_size=offloaded_block_size,

671 + gpu_block_size=gpu_block_size,

672 + num_gpu_blocks=num_gpu_blocks,

673 + )

674 +

675 + # store 1 blocks

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

677 + runner.manager.prepare_store.side_effect = (

678 + lambda block_hashes: generate_store_output(block_hashes)

679 + )

680 + runner.run(

681 + decoded_tokens=[EOS_TOKEN_ID],

682 + expected_stored_gpu_block_indexes=(0, 1, 2),

683 + )

684 +

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

686 + runner.scheduler.reset_prefix_cache()

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

688 + runner.manager.lookup.return_value = 1

689 + runner.run(

690 + decoded_tokens=[],

691 + complete_transfers=False,

692 + )

693 +

694 + # request triggered a load

695 + transfer_jobs = list(runner.offloading_spec.handler.transfer_specs)

696 + assert transfer_jobs

697 +

698 + # start a new request to load the same first block

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

700 + runner.manager.lookup.return_value = 1

701 + runner.run(

702 + decoded_tokens=[],

703 + complete_transfers=False,

704 + )

705 +

706 + # request did not trigger a load

707 + assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)

708 +

709 + # complete transfers

710 + runner.manager.prepare_store.side_effect = (

711 + lambda block_hashes: generate_store_output([])

712 + )

713 + runner.run(

714 + decoded_tokens=[EOS_TOKEN_ID],

715 + expected_loaded_gpu_block_indexes=(0, 1, 2),

716 + )

717 +

718 + # second request will use the GPU prefix cache

719 + assert transfer_jobs == list(runner.offloading_spec.handler.transfer_specs)

vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py

+36 -3

@@ -107,7 +107,7 @@ class OffloadingConnector(KVConnectorBase_V1):

107

108 def get_num_new_matched_tokens(

109 self, request: "Request", num_computed_tokens: int

110 - ) -> tuple[int, bool]:

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

111 assert self.connector_scheduler is not None

112 return self.connector_scheduler.get_num_new_matched_tokens(

113 request, num_computed_tokens

@@ -161,6 +161,11 @@ class OffloadingConnectorScheduler:

161 # request blocks are stored in order

162 # index of next block (of size offloaded_block_size) to offload

163 self._next_stored_block_idx: dict[ReqId, int] = {}

164 + # if GPU prefix caching is enabled,

165 + # track loaded blocks to avoid redundant loads

166 + self._blocks_being_loaded: set[BlockHash] | None = (

167 + set() if spec.vllm_config.cache_config.enable_prefix_caching else None

168 + )

169

170 # request ID -> set(block hashes being stored/load)

171 self._reqs_being_stored = defaultdict[ReqId, set[BlockHash]](set)

@@ -181,7 +186,7 @@ class OffloadingConnectorScheduler:

186

187 def get_num_new_matched_tokens(

188 self, request: Request, num_computed_tokens: int

184 - ) -> tuple[int, bool]:

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

190 """

191 Get number of new tokens that can be loaded beyond the

192 num_computed_tokens.

@@ -195,6 +200,9 @@ class OffloadingConnectorScheduler:

200 A tuple with the following elements:

201 - The number of tokens that can be loaded beyond what is

202 already computed.

203 + If None, it means that the connector needs more time to

204 + determine the number of matched tokens, and the scheduler

205 + should query for this request again later.

206 - `True` if tokens will be loaded asynchronously

207 (between scheduler steps).

208 """

@@ -214,6 +222,9 @@ class OffloadingConnectorScheduler:

222 hits = self.manager.lookup(

223 self._get_block_hashes(request, start_idx=start_block_idx)

224 )

225 + if hits is None:

226 + # indicates a lookup that should be tried later

227 + return None, False

228 if hits == 0:

229 return 0, False

230

@@ -229,6 +240,22 @@ class OffloadingConnectorScheduler:

240 if num_hit_tokens < self.offloaded_block_size:

241 return 0, False

242

243 + if self._blocks_being_loaded:

244 + block_hashes = self._get_block_hashes(

245 + request, start_idx=start_block_idx, end_idx=start_block_idx + hits

246 + )

247 +

248 + if any(

249 + block_hash in self._blocks_being_loaded for block_hash in block_hashes

250 + ):

251 + # hit blocks are being loaded, delay request

252 + logger.debug(

253 + "Delaying request %s since some of its blocks are already"

254 + " being loaded",

255 + request.request_id,

256 + )

257 + return None, False

258 +

259 return num_hit_tokens, True

260

261 def update_state_after_alloc(

@@ -270,9 +297,13 @@ class OffloadingConnectorScheduler:

297 )

298

299 self._reqs_to_load[request.request_id] = (src_spec, dst_spec)

273 - self._reqs_being_loaded[request.request_id].update(block_hashes)

300 + req_blocks_being_loaded = self._reqs_being_loaded[request.request_id]

301 + req_blocks_being_loaded.update(block_hashes)

302 self._next_stored_block_idx[request.request_id] = num_blocks

303

304 + if self._blocks_being_loaded is not None:

305 + self._blocks_being_loaded.update(req_blocks_being_loaded)

306 +

307 def _get_reqs_to_store(self, scheduler_output: SchedulerOutput):

308 reqs_to_store: dict[ReqId, TransferSpec] = {}

309 # iterate over both new and cached requests

@@ -379,6 +410,8 @@ class OffloadingConnectorScheduler:

410 for req_id in connector_output.finished_recving or []:

411 block_hashes = self._reqs_being_loaded.pop(req_id, None)

412 if block_hashes:

413 + if self._blocks_being_loaded:

414 + self._blocks_being_loaded.difference_update(block_hashes)

415 self.manager.complete_load(block_hashes)

416

417 def request_finished(

vllm/v1/kv_offload/abstract.py

+4 -2

@@ -68,7 +68,7 @@ class OffloadingEvent:

69 class OffloadingManager(ABC):

70 @abstractmethod

71 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int:

71 + def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:

72 """

73 Finds the length of the maximal series of blocks, starting from the

74 first one, that are all offloaded.

@@ -78,7 +78,9 @@ class OffloadingManager(ABC):

79 Returns:

80 An integer representing the maximal number of blocks that

81 - are currently offloaded.

81 + are currently offloaded, or None if the lookup should be retried

82 + later. Returning None will delay the request handling by the vLLM

83 + scheduler.

84 """

85 pass

vllm/v1/kv_offload/arc_manager.py

+1 -1

@@ -63,7 +63,7 @@ class ARCOffloadingManager(OffloadingManager):

63 self.events: list[OffloadingEvent] | None = [] if enable_events else None

64 self.cache_capacity: int = self.backend.get_num_free_blocks()

66 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int:

66 + def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:

67 hit_count = 0

68 for block_hash in block_hashes:

69 block = self.t1.get(block_hash) or self.t2.get(block_hash)

vllm/v1/kv_offload/lru_manager.py

+1 -1

@@ -24,7 +24,7 @@ class LRUOffloadingManager(OffloadingManager):

24 self.blocks: OrderedDict[BlockHash, BlockStatus] = OrderedDict()

25 self.events: list[OffloadingEvent] | None = [] if enable_events else None

27 - def lookup(self, block_hashes: Iterable[BlockHash]) -> int:

27 + def lookup(self, block_hashes: Iterable[BlockHash]) -> int | None:

28 hit_count = 0

29 for block_hash in block_hashes:

30 block = self.blocks.get(block_hash)