@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3a69636645

[Bugfix][HiSparse] Fix a chunked-prefill preemption livelock (#59494)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Matthew Bonanni committed Oct 1, 2026 at 13:02 UTC 3a6963664537ed21172e2ec12e96e3a2dcd3718c

3 files changed +106 -29

tests/v1/core/test_prefix_caching.py

+60 -1

@@ -1259,6 +1259,63 @@ def test_hisparse_cleanup_publishes_only_finalized_terminal_pages(

1259 assert not coordinator.has_pending_work()

1260

1261

1262 +@pytest.mark.parametrize("num_pages,reads_host", [(4, True), (3, False)])

1263 +def test_hisparse_prefill_reads_host_once_it_fills_admission_window(

1264 + num_pages, reads_host

1265 +):

1266 + """A prefill switches to host reads once it fills its admission window.

1267 +

1268 + Admission caps each request's resident pages at the in-flight window,

1269 + assuming older pages move to host, but pinned pages were released only

1270 + once free blocks fell below a fixed watermark. A chunked prefill larger

1271 + than the watermark outgrew the pool first and was preempted.

1272 + """

1273 + window_pages = 4

1274 + manager = make_hisparse_kv_cache_manager(

1275 + 32,

1276 + 16,

1277 + enable_caching=True,

1278 + max_in_flight_tokens=window_pages * HISPARSE_BLOCK_SIZE,

1279 + )

1280 + coordinator = get_hisparse_coordinator(manager)

1281 + tokens = list(range(2 * window_pages * HISPARSE_BLOCK_SIZE))

1282 + request = make_request("prefill", tokens, HISPARSE_BLOCK_SIZE, sha256)

1283 + assert manager.allocate_slots(request, num_pages * HISPARSE_BLOCK_SIZE) is not None

1284 + pool = coordinator.gpu_pool

1285 + assert pool is not None

1286 + assert pool.get_num_free_blocks() >= coordinator.transition_watermark

1287 +

1288 + hot_manager = coordinator.hot_managers[0]

1289 + assert (request.request_id in hot_manager.hot_required) == reads_host

1290 +

1291 +

1292 +def test_hisparse_preempted_prefill_resumes_from_durable_prefix():

1293 + """A prefill preempted mid-way must resume from the pages already on host.

1294 +

1295 + Host writes trail a chunked prefill by a chunk, and host pages were only

1296 + published once the whole computed prefix was durable. Each chunk moved

1297 + that target past the finished writes, so nothing was published before the

1298 + prefill ended: a prefill preempted each time it outgrew the GPU pool

1299 + recomputed from the start forever.

1300 + """

1301 + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)

1302 + coordinator = get_hisparse_coordinator(manager)

1303 + chunk = 2 * HISPARSE_BLOCK_SIZE

1304 + tokens = list(range(3 * chunk))

1305 + request = make_request("long", tokens, HISPARSE_BLOCK_SIZE, sha256)

1306 + assert manager.allocate_slots(request, chunk) is not None

1307 + request.num_computed_tokens = chunk

1308 + first_chunk_writes = coordinator.build_offload_command().page_transfers

1309 + assert manager.allocate_slots(request, chunk) is not None

1310 + request.num_computed_tokens = 2 * chunk

1311 + acks = {transfer.transfer_id: 1 for transfer in first_chunk_writes}

1312 + coordinator.update_spills(acks, acks)

1313 + manager.free(request)

1314 +

1315 + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1316 + assert manager.get_computed_blocks(resumed)[1] == chunk

1317 +

1318 +

1319 def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity():

1320 """Old completions retain their pages through pressure and request-ID reuse."""

1321 manager = make_hisparse_kv_cache_manager(32, 8, enable_caching=True)

@@ -1285,7 +1342,9 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity()

1342 assert not {block.block_id for block in old_blocks[:2]} & {

1343 block.block_id for block in pressure + new_blocks

1344 }

1288 - assert all(block.block_hash is None for block in old_blocks)

1345 + # Only the page every worker finished writing is published.

1346 + assert old_blocks[0].block_hash is not None

1347 + assert all(block.block_hash is None for block in old_blocks[1:])

1348 coordinator.update_spills({}, partial)

1349 repeated = make_request("probe", tokens, HISPARSE_BLOCK_SIZE, sha256)

1350 assert manager.get_computed_blocks(repeated)[1] == 2 * HISPARSE_BLOCK_SIZE

vllm/v1/core/single_type_kv_cache_manager.py

+7 -6

@@ -112,7 +112,7 @@ class SingleTypeKVCacheManager(ABC):

112 self.kv_cache_spec = kv_cache_spec

113 self.block_pool = block_pool

114 self.enable_caching = enable_caching

115 - self._max_admission_blocks_per_request = max_admission_blocks_per_request

115 + self.max_admission_blocks_per_request = max_admission_blocks_per_request

116 # Record newly allocated block ids only when worker-side zeroing will

117 # consume them and this manager holds a spec type that gets zeroed.

118 self._record_new_block_ids = (

@@ -202,7 +202,7 @@ class SingleTypeKVCacheManager(ABC):

202 model in spec decode). w/o spec decode, it is num_tokens;

203 with spec decode, it is num_tokens - num_lookahead_tokens.

204 apply_admission_cap: If True, clamp by `num_required_blocks` by

205 - `_max_admission_blocks_per_request`for recycling-aware specs

205 + `max_admission_blocks_per_request`for recycling-aware specs

206 (SWA, chunked-local).

207 prefill_end: The token index the request's prefill ends at, the

208 same value the scheduler splits chunks against. Mamba reserves

@@ -214,7 +214,7 @@ class SingleTypeKVCacheManager(ABC):

214

215 """

216 num_required_blocks = cdiv(num_tokens, self.block_size)

217 - if apply_admission_cap and self._max_admission_blocks_per_request is not None:

217 + if apply_admission_cap and self.max_admission_blocks_per_request is not None:

218 # Recycling-aware specs (SWA, chunked-local) cap the per-request

219 # reservation here so admission matches the startup pool sizer

220 # (`SlidingWindowSpec.max_admission_blocks_per_request` / its

@@ -225,7 +225,7 @@ class SingleTypeKVCacheManager(ABC):

225 # Drift between the two would re-introduce the deadlock from

226 # issue #39734 or, worse, mid-prefill OOM.

227 num_required_blocks = min(

228 - num_required_blocks, self._max_admission_blocks_per_request

228 + num_required_blocks, self.max_admission_blocks_per_request

229 )

230 num_req_blocks = len(self.req_to_blocks.get(request_id, ()))

231

@@ -2488,6 +2488,8 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager):

2488 class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2489 """Track GPU-resident pages for otherwise host-backed KV."""

2490

2491 + max_admission_blocks_per_request: int

2492 +

2493 def get_num_blocks_to_allocate(

2494 self,

2495 request_id: str,

@@ -2514,8 +2516,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2516 host_pages = cdiv(num_local_computed_tokens, self.block_size)

2517 required = cdiv(num_tokens, self.block_size)

2518 if apply_admission_cap:

2517 - assert self._max_admission_blocks_per_request is not None

2518 - required = min(required, self._max_admission_blocks_per_request)

2519 + required = min(required, self.max_admission_blocks_per_request)

2520 return max(required - max(existing, host_pages), 0)

2521

2522 def allocate_external_computed_blocks(

vllm/v1/hisparse/coordinator.py

+39 -22

@@ -318,6 +318,19 @@ class HiSparseCoordinator:

318 manager.has_hot(request_id) for manager in self.hot_managers

319 )

320

321 + def _fills_admission_window(self, request_id: str) -> bool:

322 + """Whether the request holds the resident pages admission reserved.

323 +

324 + Admission caps a request's resident pages at the in-flight window,

325 + assuming older pages move to host; keeping them pinned past it lets a

326 + chunked prefill outgrow the pool it was admitted into.

327 + """

328 + return any(

329 + len(manager.req_to_blocks.get(request_id, ()))

330 + >= manager.max_admission_blocks_per_request

331 + for manager in self.resident_managers

332 + )

333 +

334 def _resident_page_blocks(

335 self, request_id: str, page_idx: int

336 ) -> list[KVCacheBlock] | None:

@@ -390,8 +403,9 @@ class HiSparseCoordinator:

403

404 A request that can read from host releases every clean sealed page to

405 the pool. One that cannot keeps its pages pinned until the shared pool

393 - runs low, then asks for a hot region. Pages remain pinned until that region

394 - is allocated on a subsequent scheduling pass.

406 + runs low, or until it holds the resident pages admission reserved for

407 + it, then asks for a hot region. Pages remain pinned until that region is

408 + allocated on a subsequent scheduling pass.

409 """

410 if not self.resident_managers:

411 return

@@ -403,7 +417,10 @@ class HiSparseCoordinator:

417 state.pages_to_adopt = 0

418 if not self._can_read_from_host(request_id):

419 assert self.gpu_pool is not None

406 - if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark:

420 + if (

421 + self.gpu_pool.get_num_free_blocks() >= self.transition_watermark

422 + and not self._fills_admission_window(request_id)

423 + ):

424 return

425 for manager in self.hot_managers:

426 manager.require_hot(request_id)

@@ -453,47 +470,47 @@ class HiSparseCoordinator:

470 *,

471 replay_boundaries: Sequence[int],

472 ) -> None:

456 - """Publish host-source hashes only after their pages are durable."""

473 + """Publish host-source hashes of the pages that are already durable."""

474 manager = self.host_manager

475 if manager is None:

476 return

460 - num_pages = num_computed_tokens // manager.block_size

477 request_id = request.request_id

462 - state = self._get_request_state(request_id)

463 - if state.ready_prefix_pages >= num_pages:

464 - manager.publish_blocks(

465 - request,

466 - num_computed_tokens,

467 - retention_interval=retention_interval,

468 - replay_boundaries=replay_boundaries,

469 - )

470 - self._record_copies(request_id, num_computed_tokens)

471 - state.publication = None

472 - return

473 - state.publication = _PendingPublication(

478 + self._get_request_state(request_id).publication = _PendingPublication(

479 request=request,

480 num_computed_tokens=num_computed_tokens,

476 - num_pages=num_pages,

481 + num_pages=num_computed_tokens // manager.block_size,

482 retention_interval=retention_interval,

483 replay_boundaries=replay_boundaries,

484 )

485 + self._publish_host_blocks_if_ready(request_id)

486

487 def _publish_host_blocks_if_ready(self, request_id: str) -> None:

488 + """Publish the durable prefix of the pending publication.

489 +

490 + Host writes trail a long prefill by about a chunk, so waiting for the

491 + whole computed prefix would publish nothing until the prefill ends, and

492 + a prefill preempted before then would recompute from the start.

493 + """

494 state = self.request_states.get(request_id)

495 if state is None:

496 return

497 publication = state.publication

486 - if publication is None or state.ready_prefix_pages < publication.num_pages:

498 + if publication is None:

499 return

500 assert self.host_manager is not None

501 + num_tokens = min(

502 + publication.num_computed_tokens,

503 + state.ready_prefix_pages * self.host_manager.block_size,

504 + )

505 self.host_manager.publish_blocks(

506 publication.request,

491 - publication.num_computed_tokens,

507 + num_tokens,

508 retention_interval=publication.retention_interval,

509 replay_boundaries=publication.replay_boundaries,

510 )

495 - self._record_copies(request_id, publication.num_computed_tokens)

496 - state.publication = None

511 + self._record_copies(request_id, num_tokens)

512 + if state.ready_prefix_pages >= publication.num_pages:

513 + state.publication = None

514

515 def record_pending_host_import(self, request_id: str, num_tokens: int) -> None:

516 """Note a prefix an external load is populating in host pages."""