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

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 10, 09:00 UTC

[KVConnector] OffloadingConnector: Fix bug in handling of preemptions (#29870)

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

Or Ozeri committed Jan 11, 2026 at 10:05 UTC 4c16ba617f76b342dd0e62deba1f96ed6cee74fa

7 files changed +255 -64

tests/v1/kv_connector/unit/test_offloading_connector.py

+175 -64

@@ -64,8 +64,11 @@ class MockLoadStoreSpec(LoadStoreSpec):

65 class MockOffloadingHandler(OffloadingHandler):

66 def __init__(self):

67 + self.transfer_specs: dict[int, TransferSpec] = {}

68 self.completed_transfers: list[TransferResult] = []

68 - self.completed_specs: list[TransferSpec] = []

69 + self.waiting_jobs: set[int] = set()

70 + self.completed_jobs: list[int] = []

71 + self.flushed_jobs: set[int] = set()

73 def get_finished(self) -> list[TransferResult]:

74 finished = self.completed_transfers

@@ -73,10 +76,21 @@ class MockOffloadingHandler(OffloadingHandler):

76 return finished

78 def transfer_async(self, job_id: int, spec: TransferSpec) -> bool:

76 - self.completed_specs.append(spec)

77 - self.completed_transfers.append((job_id, True))

79 + self.transfer_specs[job_id] = spec

80 + self.waiting_jobs.add(job_id)

81 return True

83 + def complete_jobs(self, job_ids: set[int]) -> None:

84 + for job_id in job_ids:

85 + if job_id in self.waiting_jobs:

86 + self.waiting_jobs.remove(job_id)

87 + self.completed_jobs.append(job_id)

88 + self.completed_transfers.append((job_id, True))

89 +

90 + def wait(self, job_ids: set[int]) -> None:

91 + self.flushed_jobs |= job_ids

92 + self.complete_jobs(job_ids)

93 +

95 class MockOffloadingSpec(OffloadingSpec):

96 def __init__(self, vllm_config: VllmConfig):

@@ -98,9 +112,22 @@ class MockOffloadingSpec(OffloadingSpec):

112 yield GPULoadStoreSpec, MockLoadStoreSpec, self.handler

113 yield MockLoadStoreSpec, GPULoadStoreSpec, self.handler

114

115 + def complete_transfers(self):

116 + self.handler.complete_jobs(self.handler.waiting_jobs.copy())

117 +

118 def get_completed_transfers(self) -> list[TransferSpec]:

102 - specs = self.handler.completed_specs

103 - self.handler.completed_specs = []

119 + specs = [

120 + self.handler.transfer_specs[job_id]

121 + for job_id in self.handler.completed_jobs

122 + ]

123 + self.handler.completed_jobs.clear()

124 + return specs

125 +

126 + def get_flushed_transfers(self):

127 + specs = [

128 + self.handler.transfer_specs[job_id] for job_id in self.handler.flushed_jobs

129 + ]

130 + self.handler.flushed_jobs.clear()

131 return specs

132

133

@@ -170,12 +197,9 @@ class RequestRunner:

197 # mapping (offloading address) -> gpu_block_index

198 self.offloaded: dict[Any, int] = {}

199

173 - self.pending_loads_count: int = 0

174 - self.pending_stores_count: int = 0

175 - self.unsubmitted_stores_count = 0

176 -

200 self.completed_loads: list[TransferSummary] = []

201 self.completed_stores: list[TransferSummary] = []

202 + self.flushed_gpu_block_indexes: set[int] = set()

203

204 # maps {block_id: block_offset}

205 self.gpu_block_index: dict[int, int] = {}

@@ -202,54 +226,60 @@ class RequestRunner:

226

227 self.scheduler.add_request(req)

228

205 - def _wait_for_transfers(self):

229 + def _parse_transfers(self):

230 + for transfer_spec in self.offloading_spec.get_flushed_transfers():

231 + src_spec, dst_spec = transfer_spec

232 + assert isinstance(src_spec, GPULoadStoreSpec)

233 +

234 + for block_id in src_spec.block_ids:

235 + self.flushed_gpu_block_indexes.add(

236 + self.gpu_block_index[block_id.item()]

237 + )

238 +

239 block_size_factor = self.offloaded_block_size // self.gpu_block_size

240

208 - while self.pending_loads_count or self.pending_stores_count:

209 - for transfer_spec in self.offloading_spec.get_completed_transfers():

210 - src_spec, dst_spec = transfer_spec

211 -

212 - if isinstance(src_spec, GPULoadStoreSpec):

213 - store = True

214 - gpu_spec = src_spec

215 - offload_spec = dst_spec

216 - else:

217 - store = False

218 - gpu_spec = dst_spec

219 - offload_spec = src_spec

220 -

221 - assert isinstance(offload_spec, MockLoadStoreSpec)

222 - assert isinstance(gpu_spec, GPULoadStoreSpec)

223 -

224 - gpu_block_indices: list[int] = []

225 - for block_id in gpu_spec.block_ids:

226 - gpu_block_indices.append(self.gpu_block_index[block_id.item()])

227 -

228 - # list of (block_hash, sub_block_offset)

229 - offload_addresses: list[Any] = []

230 - for block_hash in offload_spec.block_hashes:

231 - for sub_block_idx in range(block_size_factor):

232 - offload_addresses.append((block_hash, sub_block_idx))

233 -

234 - if store:

235 - assert len(gpu_block_indices) == len(offload_addresses)

236 -

237 - self.completed_stores.append(

238 - TransferSummary(gpu_block_indices, offload_addresses)

239 - )

240 - self.pending_stores_count -= 1

241 - else:

242 - remainder_sub_block_count = len(offload_addresses) - len(

243 - gpu_block_indices

244 - )

245 - assert remainder_sub_block_count >= 0

246 - assert remainder_sub_block_count < block_size_factor

247 - offload_addresses = offload_addresses[remainder_sub_block_count:]

248 -

249 - self.completed_loads.append(

250 - TransferSummary(gpu_block_indices, offload_addresses)

251 - )

252 - self.pending_loads_count -= 1

241 + for transfer_spec in self.offloading_spec.get_completed_transfers():

242 + src_spec, dst_spec = transfer_spec

243 +

244 + if isinstance(src_spec, GPULoadStoreSpec):

245 + store = True

246 + gpu_spec = src_spec

247 + offload_spec = dst_spec

248 + else:

249 + store = False

250 + gpu_spec = dst_spec

251 + offload_spec = src_spec

252 +

253 + assert isinstance(offload_spec, MockLoadStoreSpec)

254 + assert isinstance(gpu_spec, GPULoadStoreSpec)

255 +

256 + gpu_block_indices: list[int] = []

257 + for block_id in gpu_spec.block_ids:

258 + gpu_block_indices.append(self.gpu_block_index[block_id.item()])

259 +

260 + # list of (block_hash, sub_block_offset)

261 + offload_addresses: list[Any] = []

262 + for block_hash in offload_spec.block_hashes:

263 + for sub_block_idx in range(block_size_factor):

264 + offload_addresses.append((block_hash, sub_block_idx))

265 +

266 + if store:

267 + assert len(gpu_block_indices) == len(offload_addresses)

268 +

269 + self.completed_stores.append(

270 + TransferSummary(gpu_block_indices, offload_addresses)

271 + )

272 + else:

273 + remainder_sub_block_count = len(offload_addresses) - len(

274 + gpu_block_indices

275 + )

276 + assert remainder_sub_block_count >= 0

277 + assert remainder_sub_block_count < block_size_factor

278 + offload_addresses = offload_addresses[remainder_sub_block_count:]

279 +

280 + self.completed_loads.append(

281 + TransferSummary(gpu_block_indices, offload_addresses)

282 + )

283

284 def _update_gpu_block_idx(self):

285 for blocks in self.scheduler.kv_cache_manager.coordinator.single_type_managers[

@@ -258,18 +288,19 @@ class RequestRunner:

288 for block_idx, block in enumerate(blocks):

289 self.gpu_block_index[block.block_id] = block_idx

290

261 - def _run(self, decoded_tokens: list[int]):

291 + def _run(self, decoded_tokens: list[int], complete_transfers: bool):

292 """

293 Runs multiple engine (scheduler + worker) steps.

294 Assumes a single request is running.

295

296 Args:

297 decoded_tokens: the tokens to yield at each step.

298 + complete_transfers: complete transfers immediately

299 """

300

301 tokens_iter = iter(decoded_tokens)

302 token_id = next(tokens_iter, None)

272 - while token_id is not None:

303 + while True:

304 assert self.scheduler.requests

305

306 scheduler_output = self.scheduler.schedule()

@@ -279,10 +310,10 @@ class RequestRunner:

310 assert kv_connector_metadata is not None

311 assert isinstance(kv_connector_metadata, OffloadingConnectorMetadata)

312

282 - self.pending_loads_count += len(kv_connector_metadata.reqs_to_load)

283 -

284 - self.pending_stores_count += self.unsubmitted_stores_count

285 - self.unsubmitted_stores_count = len(kv_connector_metadata.reqs_to_store)

313 + if scheduler_output.preempted_req_ids:

314 + self.worker_connector.handle_preemptions(

315 + scheduler_output.preempted_req_ids

316 + )

317

318 self.worker_connector.bind_connector_metadata(kv_connector_metadata)

319 self.worker_connector.start_load_kv(self._dummy_ctx)

@@ -290,6 +321,9 @@ class RequestRunner:

321 if scheduler_output.total_num_scheduled_tokens > 0:

322 self.worker_connector.wait_for_save()

323

324 + if complete_transfers:

325 + self.offloading_spec.complete_transfers()

326 +

327 finished_sending, finished_recving = self.worker_connector.get_finished(

328 scheduler_output.finished_req_ids

329 )

@@ -300,7 +334,7 @@ class RequestRunner:

334 reqs=self.scheduler.running,

335 finished_sending=finished_sending,

336 finished_recving=finished_recving,

303 - token_id=token_id,

337 + token_id=token_id or 0,

338 )

339

340 if self.scheduler.running:

@@ -308,7 +342,10 @@ class RequestRunner:

342

343 self.scheduler.update_from_output(scheduler_output, model_runner_output)

344

311 - self._wait_for_transfers()

345 + if token_id is None:

346 + break

347 +

348 + self._parse_transfers()

349

350 # run one more step to update finished stored

351 if EOS_TOKEN_ID in decoded_tokens:

@@ -333,8 +370,10 @@ class RequestRunner:

370 def run(

371 self,

372 decoded_tokens: list[int],

373 + complete_transfers: bool = True,

374 expected_stored_gpu_block_indexes: tuple[int, ...] = (),

375 expected_loaded_gpu_block_indexes: tuple[int, ...] = (),

376 + expected_flushed_gpu_block_indexes: tuple[int, ...] = (),

377 ):

378 """

379 Runs multiple engine (scheduler + worker) steps.

@@ -342,14 +381,17 @@ class RequestRunner:

381

382 Args:

383 decoded_tokens: the tokens to yield at each step.

384 + complete_transfers: complete transfers immediately

385 expected_stored_gpu_block_indexes: GPU block indexes

386 that are expected to be written during the run.

387 expected_loaded_gpu_block_indexes: GPU block indexes

388 that are expected to be loaded during the run.

389 + expected_flushed_gpu_block_indexes: GPU block indexes

390 + that are expected to be flushed during the run.

391 """

392

393 self.manager.reset_mock()

352 - self._run(decoded_tokens)

394 + self._run(decoded_tokens, complete_transfers)

395

396 loaded_gpu_block_indexes: set[int] = set()

397 for transfer in self.completed_loads:

@@ -373,6 +415,9 @@ class RequestRunner:

415 assert set(expected_stored_gpu_block_indexes) == stored_gpu_block_indexes

416 self.completed_stores.clear()

417

418 + assert set(expected_flushed_gpu_block_indexes) == self.flushed_gpu_block_indexes

419 + self.flushed_gpu_block_indexes.clear()

420 +

421

422 @pytest.fixture

423 def request_runner():

@@ -539,3 +584,69 @@ def test_offloading_connector(request_runner):

584 assert isinstance(event, BlockRemoved)

585 assert event.block_hashes == to_hashes([4, 5, 6])

586 assert event.medium == "B"

587 +

588 +

589 +def test_request_preemption(request_runner):

590 + offloaded_block_size = 12

591 + gpu_block_size = 4

592 + num_gpu_blocks = 100

593 +

594 + runner = request_runner(

595 + offloaded_block_size=offloaded_block_size,

596 + gpu_block_size=gpu_block_size,

597 + num_gpu_blocks=num_gpu_blocks,

598 + )

599 +

600 + free_block_queue = runner.scheduler.kv_cache_manager.block_pool.free_block_queue

601 + num_free_blocks_empty = free_block_queue.num_free_blocks

602 +

603 + # 2 blocks, store all, without flushing

604 + # blocks = [0, 1, 2], [3, 4, 5]

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

606 + runner.manager.prepare_store.side_effect = (

607 + lambda block_hashes: generate_store_output(block_hashes)

608 + )

609 + runner.run(

610 + decoded_tokens=[0],

611 + complete_transfers=False,

612 + )

613 +

614 + # decode 2 more blocks - 1 gpu block, storing [6, 7, 8] (no flush)

615 + runner.manager.prepare_store.side_effect = (

616 + lambda block_hashes: generate_store_output(block_hashes)

617 + )

618 + runner.run(

619 + decoded_tokens=[0] * (2 * offloaded_block_size - gpu_block_size),

620 + complete_transfers=False,

621 + )

622 +

623 + # simulate KV cache running out of space

624 + free_block_queue.num_free_blocks = 0

625 +

626 + # request should be preempted now

627 + runner.run(

628 + decoded_tokens=[],

629 + complete_transfers=False,

630 + expected_flushed_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),

631 + expected_stored_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),

632 + )

633 +

634 + # restore KV cache space and reset GPU prefix cache

635 + free_block_queue.num_free_blocks = num_free_blocks_empty

636 + runner.scheduler.reset_prefix_cache()

637 +

638 + # request should now return from preemption

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

640 + runner.manager.lookup.return_value = 3

641 + runner.manager.prepare_store.side_effect = (

642 + lambda block_hashes: generate_store_output(block_hashes)

643 + )

644 + runner.run(

645 + decoded_tokens=[0] * gpu_block_size,

646 + expected_loaded_gpu_block_indexes=(0, 1, 2, 3, 4, 5, 6, 7, 8),

647 + )

648 +

649 + runner.run(

650 + decoded_tokens=[EOS_TOKEN_ID],

651 + expected_stored_gpu_block_indexes=(9, 10, 11),

652 + )

tests/v1/kv_offload/test_worker.py

+12

@@ -63,6 +63,12 @@ class OffloadingHandler1To2(OffloadingHandler):

63 del self.transfers[job_id]

64 return finished

66 + def wait(self, job_ids: set[int]) -> None:

67 + for job_id in job_ids:

68 + spec = self.transfers.get(job_id)

69 + if spec:

70 + assert spec.finished

71 +

73 class OffloadingHandler2To1(OffloadingHandler):

74 def __init__(self):

@@ -84,6 +90,12 @@ class OffloadingHandler2To1(OffloadingHandler):

90 del self.transfers[job_id]

91 return finished

93 + def wait(self, job_ids: set[int]) -> None:

94 + for job_id in job_ids:

95 + spec = self.transfers.get(job_id)

96 + if spec:

97 + assert spec.finished

98 +

100 def test_offloading_worker():

101 """

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

+10

@@ -25,6 +25,9 @@ The class provides the following primitives:

26 Worker-side: runs in each worker, loads/saves KV cache to/from

27 the Connector based on the metadata.

28 + handle_preemptions() - called if there are preempted requests,

29 + before their blocks are overwritten

30 +

31 start_load_kv() - starts loading all KVs (maybe async)

32 wait_for_layer_load() - blocks until layer i load is done

@@ -262,6 +265,13 @@ class KVConnectorBase_V1(ABC):

265 """

266 return

267

268 + def handle_preemptions(self, preempted_req_ids: set[str]):

269 + """

270 + Handle preempted requests BEFORE their blocks are overwritten.

271 + Needed for connectors which use async saves (e.g., OffloadingConnector)

272 + """

273 + return

274 +

275 @abstractmethod

276 def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None:

277 """

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

+24

@@ -75,6 +75,10 @@ class OffloadingConnector(KVConnectorBase_V1):

75 assert self.connector_worker is not None

76 self.connector_worker.register_cross_layers_kv_cache(kv_cache, attn_backend)

78 + def handle_preemptions(self, preempted_req_ids: set[str]):

79 + assert self.connector_worker is not None

80 + self.connector_worker.handle_preemptions(preempted_req_ids)

81 +

82 def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:

83 assert self.connector_worker is not None

84 assert isinstance(self._connector_metadata, OffloadingConnectorMetadata)

@@ -348,6 +352,15 @@ class OffloadingConnectorScheduler:

352 reqs_to_store=self._get_reqs_to_store(scheduler_output),

353 )

354 self._reqs_to_load = {}

355 +

356 + # NOTE (orozery): we should move this logic to update_connector_output

357 + # once KVConnectorOutput allows us to report completed transfers

358 + for req_id in scheduler_output.preempted_req_ids or ():

359 + block_hashes = self._reqs_being_stored.get(req_id)

360 + if block_hashes:

361 + self.manager.complete_store(block_hashes)

362 + block_hashes.clear()

363 +

364 return meta

365

366 def update_connector_output(self, connector_output: KVConnectorOutput):

@@ -466,6 +479,17 @@ class OffloadingConnectorWorker:

479 attn_backends = {cross_layer_name: attn_backend}

480 self._register_handlers(kv_caches, attn_backends)

481

482 + def handle_preemptions(self, preempted_req_ids: set[str]):

483 + for job_id, transfer_spec in self._unsubmitted_store_jobs:

484 + success = self.worker.transfer_async(job_id, transfer_spec)

485 + assert success

486 + self._unsubmitted_store_jobs.clear()

487 +

488 + for req_id in preempted_req_ids:

489 + job_ids = self._store_jobs.get(req_id)

490 + if job_ids:

491 + self.worker.wait(job_ids)

492 +

493 def start_kv_transfers(self, metadata: OffloadingConnectorMetadata):

494 for job_id, transfer_spec in self._unsubmitted_store_jobs:

495 success = self.worker.transfer_async(job_id, transfer_spec)

vllm/v1/kv_offload/worker/cpu_gpu.py

+10

@@ -96,6 +96,8 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

96 assert len(src_tensors) > 0

97 self.gpu_to_cpu: bool = self.src_tensors[0].is_cuda

99 + # job_id -> event

100 + self._transfer_events: dict[int, torch.Event] = {}

101 # queue of transfers (job_id, stream, event)

102 self._transfers: deque[tuple[int, torch.cuda.Stream, torch.Event]] = deque()

103 # list of CUDA streams available for re-use

@@ -152,6 +154,7 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

154 ops.swap_blocks(src_tensor, dst_tensor, src_to_dst_tensor)

155 event.record(stream)

156

157 + self._transfer_events[job_id] = event

158 self._transfers.append((job_id, stream, event))

159

160 # success

@@ -164,8 +167,15 @@ class SingleDirectionOffloadingHandler(OffloadingHandler):

167 results.append((job_id, True))

168 self._stream_pool.append(stream)

169 self._event_pool.append(event)

170 + del self._transfer_events[job_id]

171 return results

172

173 + def wait(self, job_ids: set[int]):

174 + for job_id in job_ids:

175 + event = self._transfer_events.get(job_id)

176 + if event is not None:

177 + event.synchronize()

178 +

179

180 class CpuGpuOffloadingHandlers:

181 def __init__(

vllm/v1/kv_offload/worker/worker.py

+19

@@ -53,6 +53,15 @@ class OffloadingHandler(ABC):

53 """

54 pass

56 + @abstractmethod

57 + def wait(self, job_ids: set[int]) -> None:

58 + """

59 + Wait for jobs to finish (blocking).

60 +

61 + Args:

62 + job_ids: The set of job IDs to wait for.

63 + """

64 +

66 class OffloadingWorker:

67 """

@@ -142,3 +151,13 @@ class OffloadingWorker:

151 for handler in self.handlers:

152 finished.extend(handler.get_finished())

153 return finished

154 +

155 + def wait(self, job_ids: set[int]) -> None:

156 + """

157 + Wait for jobs to finish (blocking).

158 +

159 + Args:

160 + job_ids: The set of job IDs to wait for.

161 + """

162 + for handler in self.handlers:

163 + handler.wait(job_ids)

vllm/v1/worker/gpu_model_runner.py

@@ -3112,6 +3112,11 @@ class GPUModelRunner(

3112 "after execute_model() returns None."

3113 )

3114

3115 + if scheduler_output.preempted_req_ids and has_kv_transfer_group():

3116 + get_kv_transfer_group().handle_preemptions(

3117 + scheduler_output.preempted_req_ids

3118 + )

3119 +

3120 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

3121 with (

3122 record_function_or_nullcontext("gpu_model_runner: preprocess"),