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

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 2, 07:15 UTC

[BugFix][kv_offload]: Prevent offloading stale sliding window blocks (#42959)

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

Or Ozeri committed Jun 2, 2026 at 05:59 UTC 480fadab1b6d6b2ee9bc039ee89b57fb64d66f60

2 files changed +156 -40

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

+76

@@ -1305,3 +1305,79 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):

1305 (1, 7),

1306 ),

1307 )

1308 +

1309 +

1310 +@pytest.mark.parametrize("async_scheduling", [True, False])

1311 +def test_stale_sliding_window_block_after_prepare_store_failure(

1312 + request_runner, async_scheduling: bool

1313 +):

1314 + """Regression test: when prepare_store fails (returns None), offloading is

1315 + delayed. Meanwhile, sliding window blocks get freed and reallocated to the

1316 + same request. On retry, the stale block_id must be detected and skipped.

1317 +

1318 + Without the fix, the stale block_id would either:

1319 + - Cause a KeyError in _remove_pending_job (duplicate in

1320 + _block_id_to_pending_jobs)

1321 + - Silently offload wrong data under a wrong key

1322 + """

1323 + block_size = 4

1324 + # sliding_window = 8 -> window of 2 blocks

1325 + sliding_window = 8

1326 + # Use a tight GPU block budget so freed sliding window blocks are

1327 + # immediately reused by the same request's new allocations.

1328 + num_gpu_blocks = 4

1329 +

1330 + kv_cache_groups = [

1331 + KVCacheGroupSpec(

1332 + ["layer0"],

1333 + SlidingWindowSpec(

1334 + block_size=block_size,

1335 + num_kv_heads=1,

1336 + head_size=1,

1337 + dtype=torch.float32,

1338 + sliding_window=sliding_window,

1339 + ),

1340 + ),

1341 + ]

1342 +

1343 + runner = request_runner(

1344 + block_size=block_size,

1345 + num_gpu_blocks=num_gpu_blocks,

1346 + async_scheduling=async_scheduling,

1347 + kv_cache_groups=kv_cache_groups,

1348 + )

1349 +

1350 + # Request with 3 blocks of prompt. Window = 2 blocks, so block 0 is

1351 + # outside the window but won't be freed until the next allocate_slots.

1352 + runner.new_request(token_ids=[0] * block_size * 3)

1353 +

1354 + # First step: prepare_store FAILS -> offloading delayed.

1355 + # next_stored_block_idx stays at 0, block_ids[0] still holds the

1356 + # original block_id for position 0.

1357 + runner.manager.prepare_store.side_effect = lambda keys, req_context: None

1358 + runner.run(decoded_tokens=[0])

1359 + runner.manager.prepare_store.assert_called()

1360 +

1361 + # Second step: decode more tokens -> block 3 allocated.

1362 + # allocate_slots calls remove_skipped_blocks which frees block 0

1363 + # (it's now outside the sliding window). With num_gpu_blocks=4,

1364 + # the freed block is immediately reused for the new allocation.

1365 + # prepare_store still fails so offloading is still delayed.

1366 + runner.manager.prepare_store.side_effect = lambda keys, req_context: None

1367 + runner.run(decoded_tokens=[0] * block_size)

1368 +

1369 + # Now prepare_store succeeds.

1370 + # Without the fix, the request would try to offload the stale block_id

1371 + # at position 0 (now reused at position 3), causing a duplicate in

1372 + # sliding_window_block_ids and eventually a KeyError.

1373 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

1374 + generate_store_output(keys)

1375 + )

1376 + # block_ids=[0, ?, 3, 1]: positions 0 and 1 are zeroed (stale blocks that

1377 + # were freed by the sliding window and reallocated). Only blocks at

1378 + # positions 2 and 3 (request offsets 2, 3) are stored.

1379 + runner.run(

1380 + decoded_tokens=[EOS_TOKEN_ID],

1381 + expected_stored=(2, 3),

1382 + expected_flushed=(2, 3) if not async_scheduling else (),

1383 + )

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

+80 -40

@@ -286,6 +286,8 @@ class OffloadingConnectorScheduler:

286 self._req_status: dict[ReqId, RequestOffloadState] = {}

287 self._current_batch_load_jobs: dict[int, TransferJob] = {}

288 self._current_batch_jobs_to_flush: set[int] = set()

289 + # GPU block IDs allocated in the current engine step

290 + self._current_batch_allocated_block_ids: set[int] = set()

291 # if GPU prefix caching is enabled,

292 # track loaded blocks to avoid redundant loads

293 self._blocks_being_loaded: set[OffloadKey] | None = (

@@ -589,6 +591,10 @@ class OffloadingConnectorScheduler:

591 req_status.group_states,

592 blocks.blocks,

593 ):

594 + self._current_batch_allocated_block_ids.update(

595 + block.block_id for block in group_blocks if block.block_id != 0

596 + )

597 +

598 gpu_block_size = group_config.gpu_block_size

599 offloaded_block_size = group_config.offloaded_block_size

600 offload_keys = group_state.offload_keys

@@ -638,17 +644,6 @@ class OffloadingConnectorScheduler:

644 if req_status.offloading_context.policy == OffloadPolicy.BLOCK_LEVEL:

645 group_state.next_stored_block_idx = num_blocks

646

641 - # Fence dst blocks against finished-request pending stores.

642 - if (

643 - self._block_id_to_pending_jobs

644 - and not self._block_id_to_pending_jobs.keys().isdisjoint(dst_block_ids)

645 - ):

646 - self._current_batch_jobs_to_flush.update(

647 - jid

648 - for bid in dst_block_ids

649 - for jid in self._block_id_to_pending_jobs.get(bid, ())

650 - )

651 -

647 src_spec = self.manager.prepare_load(keys_to_load, req_status.req_context)

648 dst_spec = GPULoadStoreSpec(

649 dst_block_ids, group_sizes=group_sizes, block_indices=block_indices

@@ -672,37 +667,68 @@ class OffloadingConnectorScheduler:

667 if self._blocks_being_loaded is not None:

668 self._blocks_being_loaded.update(keys_to_load)

669

675 - def _build_store_jobs(

676 - self,

677 - scheduler_output: SchedulerOutput,

678 - ) -> dict[int, TransferJob]:

679 - block_size_factor = self.config.block_size_factor

680 - store_jobs: dict[int, TransferJob] = {}

681 - # iterate over both new and cached requests

670 + def _update_req_states(self, scheduler_output: SchedulerOutput) -> None:

671 + """

672 + Update request states from the Scheduler's output.

673 + """

674 +

675 + # new_block_ids_end[req_id][i] = end of pre-existing block_ids for

676 + # the i-th sliding window group (before this step's extend).

677 + # Used to detect sliding window blocks that got re-allocated.

678 + new_block_ids_end: dict[str, tuple[int, ...]] = {}

679 +

680 for req_id, new_block_id_groups, preempted in yield_req_data(scheduler_output):

681 req_status = self._req_status[req_id]

682 req_status.update_offload_keys()

685 - req = req_status.req

683

684 if preempted:

685 for group_state in req_status.group_states:

686 group_state.block_ids.clear()

687

688 if new_block_id_groups:

689 + if self._sliding_window_groups:

690 + new_block_ids_end[req_id] = tuple(

691 + len(req_status.group_states[grp_idx].block_ids)

692 + for grp_idx in self._sliding_window_groups

693 + )

694 req_status.update_block_id_groups(new_block_id_groups)

693 - # Fence new blocks against in-flight stores.

694 - if self._block_id_to_pending_jobs:

695 - new_blocks_flat = [

696 - bid for new_blocks in new_block_id_groups for bid in new_blocks

697 - ]

698 - if not self._block_id_to_pending_jobs.keys().isdisjoint(

699 - new_blocks_flat

700 - ):

701 - self._current_batch_jobs_to_flush.update(

702 - jid

703 - for bid in new_blocks_flat

704 - for jid in self._block_id_to_pending_jobs.get(bid, ())

705 - )

695 + for new_blocks in new_block_id_groups:

696 + for bid in new_blocks:

697 + if bid != 0:

698 + self._current_batch_allocated_block_ids.add(bid)

699 +

700 + # Zero out stale block_ids in sliding window groups' pending-store

701 + # positions. Only sliding window groups can have stale entries (blocks

702 + # freed by remove_skipped_blocks then reallocated). Only positions in

703 + # [next_stored_block_idx * bsf, end) need checking where end is the

704 + # pre-extend length: earlier positions were already offloaded, later

705 + # ones are fresh allocations from this step.

706 + if self._sliding_window_groups and self._current_batch_allocated_block_ids:

707 + block_size_factor = self.config.block_size_factor

708 + for req_id, req_status in self._req_status.items():

709 + ends = new_block_ids_end.get(req_id)

710 + for i, grp_idx in enumerate(self._sliding_window_groups):

711 + group_state = req_status.group_states[grp_idx]

712 + start = group_state.next_stored_block_idx * block_size_factor

713 + end = ends[i] if ends is not None else len(group_state.block_ids)

714 + for j in range(start, end):

715 + if (

716 + group_state.block_ids[j]

717 + in self._current_batch_allocated_block_ids

718 + ):

719 + group_state.block_ids[j] = 0

720 +

721 + def _build_store_jobs(

722 + self,

723 + scheduler_output: SchedulerOutput,

724 + ) -> dict[int, TransferJob]:

725 + block_size_factor = self.config.block_size_factor

726 + store_jobs: dict[int, TransferJob] = {}

727 + for req_id in scheduler_output.num_scheduled_tokens:

728 + req_status = self._req_status.get(req_id)

729 + if req_status is None:

730 + continue

731 + req = req_status.req

732

733 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]

734 num_tokens_after_batch = req.num_computed_tokens + num_scheduled_tokens

@@ -735,11 +761,8 @@ class OffloadingConnectorScheduler:

761 # For each block to offload, take the last corresponding GPU block.

762 # e.g. if block size factor is 3 and GPU block IDs are

763 # 1 5 6 7 2 4 9 3 8 then we'll take blocks 6 4 8.

738 - # We will use these GPU blocks to determine if the block needs

739 - # offloading, or (if the GPU block ID is 0) this block should

740 - # be skipped due to sliding window attention / SSM.

741 - # We know that if a block is skipped, then all the previous blocks

742 - # are skipped as well. This is why we take the last of each block.

764 + # A block_id of 0 means either a sliding window / SSM skip

765 + # or a stale entry that was zeroed out — skip it either way.

766 offload_block_ids = group_state.block_ids[

767 start_block_idx * block_size_factor

768 + block_size_factor

@@ -814,10 +837,8 @@ class OffloadingConnectorScheduler:

837 for i in range(block_size_factor):

838 block_id = block_ids[gpu_block_idx + i]

839 if block_id == 0:

817 - # skipped blocks cannot appear after non-skipped blocks

818 - assert start_gpu_block_idx is None

840 continue

820 - elif start_gpu_block_idx is None:

841 + if start_gpu_block_idx is None:

842 start_gpu_block_idx = gpu_block_idx + i

843 src_block_ids.append(block_id)

844 num_group_blocks += 1

@@ -875,6 +896,9 @@ class OffloadingConnectorScheduler:

896 def build_connector_meta(

897 self, scheduler_output: SchedulerOutput

898 ) -> KVConnectorMetadata:

899 + self._update_req_states(scheduler_output)

900 +

901 + # Flush jobs for preempted requests.

902 for req_id in scheduler_output.preempted_req_ids or ():

903 req_status = self._req_status.get(req_id)

904 if req_status is None or not req_status.transfer_jobs:

@@ -883,6 +907,20 @@ class OffloadingConnectorScheduler:

907 assert self._jobs[any_jid].is_store

908 self._current_batch_jobs_to_flush.update(req_status.transfer_jobs)

909

910 + # Flush jobs that contain re-allocated blocks.

911 + if (

912 + self._block_id_to_pending_jobs

913 + and not self._block_id_to_pending_jobs.keys().isdisjoint(

914 + self._current_batch_allocated_block_ids

915 + )

916 + ):

917 + self._current_batch_jobs_to_flush.update(

918 + jid

919 + for bid in self._current_batch_allocated_block_ids

920 + if bid in self._block_id_to_pending_jobs

921 + for jid in self._block_id_to_pending_jobs[bid]

922 + )

923 +

924 # If all tracked requests are finished, flush all pending jobs

925 # (both store and load) - there might not be a future scheduler

926 # step to trigger their completion.

@@ -898,6 +936,7 @@ class OffloadingConnectorScheduler:

936 )

937 self._current_batch_load_jobs = {}

938 self._current_batch_jobs_to_flush = set()

939 + self._current_batch_allocated_block_ids = set()

940 return meta

941

942 def update_connector_output(self, connector_output: KVConnectorOutput):

@@ -1013,6 +1052,7 @@ class OffloadingConnectorScheduler:

1052 # reset_cache cannot be called in the middle of a schedule step

1053 assert not self._current_batch_load_jobs

1054 assert not self._current_batch_jobs_to_flush

1055 + assert not self._current_batch_allocated_block_ids

1056

1057 # Flush all in-flight jobs

1058 self._current_batch_jobs_to_flush.update(self._jobs.keys())