@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 569adb5a97

[Bugfix][KV Offload] Fix SWA store reachability during chunked prefill (#54362)

Signed-off-by: Will Hampson <whamp@ggl.slmail.me> Co-authored-by: Will Hampson <whamp@ggl.slmail.me> Co-authored-by: Or Ozeri <oro@il.ibm.com>

Will Hampson committed Sep 6, 2026 at 03:50 UTC 569adb5a9780f9c02d22a6b29826acf711512356

4 files changed +385 -15

tests/v1/core/test_prefix_caching.py

+35

@@ -4389,6 +4389,41 @@ def test_mamba_shared_prefix_reuse_under_zero_retention():

4389 assert last_req_hit(retention=0, pin=True) == 2 * block_size

4390

4391

4392 +@pytest.mark.skip_global_cleanup

4393 +def test_swa_reachable_block_mask_final_partial_segment():

4394 + """Only a final SWA horizon makes its incomplete segment tail reachable."""

4395 + from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager

4396 +

4397 + spec = SlidingWindowSpec(

4398 + block_size=4,

4399 + num_kv_heads=1,

4400 + head_size=1,

4401 + dtype=torch.float32,

4402 + sliding_window=16,

4403 + )

4404 + start_block = 243

4405 + end_block = 250

4406 +

4407 + def reachable_blocks(

4408 + final_segment_end_block: int, *, use_eagle: bool = False

4409 + ) -> set[int]:

4410 + mask = SlidingWindowManager.reachable_block_mask(

4411 + start_block=start_block,

4412 + end_block=end_block,

4413 + alignment_tokens=36,

4414 + kv_cache_spec=spec,

4415 + use_eagle=use_eagle,

4416 + final_segment_end_block=final_segment_end_block,

4417 + )

4418 + assert mask is not None

4419 + return {start_block + i for i, reachable in enumerate(mask) if reachable}

4420 +

4421 + assert reachable_blocks(300) == {248, 249}

4422 + assert reachable_blocks(250) == {246, 247, 248, 249}

4423 + assert reachable_blocks(300, use_eagle=True) == {243, 248, 249}

4424 + assert reachable_blocks(250, use_eagle=True) == {243, 248, 249}

4425 +

4426 +

4427 def test_swa_reachable_block_mask_pins_shared_prefix():

4428 """SWA analog of the Mamba pin: the shared-prefix junction must keep the

4429 ``need``-block sliding-window tail ending on that boundary (not a single

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

+215

@@ -48,6 +48,7 @@ from vllm.v1.core.sched.output import (

48 from vllm.v1.kv_cache_interface import (

49 ChunkedLocalAttentionSpec,

50 FullAttentionSpec,

51 + KVCacheConfig,

52 KVCacheGroupSpec,

53 SlidingWindowSpec,

54 )

@@ -3390,6 +3391,220 @@ class TestEagle:

3391 f"interior hole in stored blocks: {offsets}"

3392 )

3393

3394 + @pytest.mark.skip_global_cleanup

3395 + @pytest.mark.parametrize(

3396 + "abort_after_first_step,async_scheduling,store_horizon_tokens",

3397 + [

3398 + pytest.param(False, False, 1200, id="completed-prompt"),

3399 + pytest.param(True, False, 1000, id="aborted-sync"),

3400 + pytest.param(True, True, 1000, id="aborted-async"),

3401 + ],

3402 + )

3403 + def test_chunked_prefill_uses_final_swa_store_horizon(

3404 + self,

3405 + request_runner,

3406 + abort_after_first_step: bool,

3407 + async_scheduling: bool,

3408 + store_horizon_tokens: int,

3409 + ):

3410 + """SWA stores must use the completed or aborted request horizon."""

3411 + full_attn_block_size = 36

3412 + swa_block_size = 4

3413 + sliding_window = 16

3414 + num_tokens = 1200

3415 +

3416 + kv_cache_groups = [

3417 + KVCacheGroupSpec(

3418 + ["layer0"],

3419 + FullAttentionSpec(

3420 + block_size=full_attn_block_size,

3421 + num_kv_heads=1,

3422 + head_size=1,

3423 + dtype=torch.float32,

3424 + ),

3425 + ),

3426 + KVCacheGroupSpec(

3427 + ["layer1"],

3428 + SlidingWindowSpec(

3429 + block_size=swa_block_size,

3430 + num_kv_heads=1,

3431 + head_size=1,

3432 + dtype=torch.float32,

3433 + sliding_window=sliding_window,

3434 + ),

3435 + ),

3436 + ]

3437 + runner = request_runner(

3438 + block_size=swa_block_size,

3439 + num_gpu_blocks=1000,

3440 + async_scheduling=async_scheduling,

3441 + kv_cache_groups=kv_cache_groups,

3442 + )

3443 +

3444 + runner.new_request(token_ids=[0] * num_tokens)

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

3446 + generate_store_output(keys)

3447 + )

3448 +

3449 + def abort_after_prefill_step() -> None:

3450 + if runner.scheduler.running:

3451 + runner.scheduler.finish_requests(

3452 + (str(runner.req_id),), RequestStatus.FINISHED_ABORTED

3453 + )

3454 +

3455 + runner._run(

3456 + [1] if abort_after_first_step else [1, 1, EOS_TOKEN_ID],

3457 + complete_transfers=True,

3458 + post_step_fn=(abort_after_prefill_step if abort_after_first_step else None),

3459 + )

3460 +

3461 + stored_swa_chunks = sorted(

3462 + block.request_block_offset

3463 + for transfer in runner.completed_stores

3464 + for block in transfer.gpu_blocks

3465 + if block.group_idx == 1

3466 + )

3467 + store_horizon_chunks = store_horizon_tokens // swa_block_size

3468 + alignment_chunks = full_attn_block_size // swa_block_size

3469 + sliding_window_chunks = sliding_window // swa_block_size

3470 + expected_swa_chunks = []

3471 + for chunk in range(store_horizon_chunks):

3472 + segment_start = chunk - chunk % alignment_chunks

3473 + segment_length = min(alignment_chunks, store_horizon_chunks - segment_start)

3474 + if chunk % alignment_chunks >= segment_length - sliding_window_chunks:

3475 + expected_swa_chunks.append(chunk)

3476 + assert stored_swa_chunks == expected_swa_chunks

3477 +

3478 + @pytest.mark.skip_global_cleanup

3479 + def test_active_decode_does_not_advance_swa_final_horizon(self, request_runner):

3480 + """Active decode chunks are not final partial SWA segments."""

3481 + full_attn_block_size = 36

3482 + swa_block_size = 4

3483 + prompt_tokens = 1200

3484 + kv_cache_groups = [

3485 + KVCacheGroupSpec(

3486 + ["layer0"],

3487 + FullAttentionSpec(

3488 + block_size=full_attn_block_size,

3489 + num_kv_heads=1,

3490 + head_size=1,

3491 + dtype=torch.float32,

3492 + ),

3493 + ),

3494 + KVCacheGroupSpec(

3495 + ["layer1"],

3496 + SlidingWindowSpec(

3497 + block_size=swa_block_size,

3498 + num_kv_heads=1,

3499 + head_size=1,

3500 + dtype=torch.float32,

3501 + sliding_window=16,

3502 + ),

3503 + ),

3504 + ]

3505 + runner = request_runner(

3506 + block_size=swa_block_size,

3507 + num_gpu_blocks=1000,

3508 + async_scheduling=False,

3509 + kv_cache_groups=kv_cache_groups,

3510 + )

3511 + runner.new_request(token_ids=[0] * prompt_tokens)

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

3513 + generate_store_output(keys)

3514 + )

3515 +

3516 + runner._run([1] * 12, complete_transfers=True)

3517 +

3518 + assert runner.scheduler.running

3519 + stored_swa_chunks = {

3520 + block.request_block_offset

3521 + for transfer in runner.completed_stores

3522 + for block in transfer.gpu_blocks

3523 + if block.group_idx == 1

3524 + }

3525 + first_decode_chunk = prompt_tokens // swa_block_size

3526 + assert first_decode_chunk not in stored_swa_chunks

3527 +

3528 + def test_dcp_abort_reconsiders_final_swa_tail(self):

3529 + """DCP-expanded SWA blocks are not scaled twice during abort."""

3530 + vllm_config = _make_vllm_config(

3531 + extra_config={"offload_prompt_only": False},

3532 + tensor_parallel_size=2,

3533 + decode_context_parallel_size=2,

3534 + )

3535 + vllm_config.speculative_config = None

3536 + vllm_config.kv_events_config = KVEventsConfig(

3537 + enable_kv_cache_events=True, publisher="null"

3538 + )

3539 + kv_cache_config = KVCacheConfig(

3540 + num_blocks=0,

3541 + kv_cache_tensors=[],

3542 + kv_cache_groups=[

3543 + KVCacheGroupSpec(

3544 + ["full"],

3545 + FullAttentionSpec(

3546 + block_size=36,

3547 + num_kv_heads=1,

3548 + head_size=1,

3549 + dtype=torch.float32,

3550 + ),

3551 + ),

3552 + KVCacheGroupSpec(

3553 + ["swa"],

3554 + SlidingWindowSpec(

3555 + block_size=4,

3556 + num_kv_heads=1,

3557 + head_size=1,

3558 + dtype=torch.float32,

3559 + sliding_window=16,

3560 + ),

3561 + ),

3562 + ],

3563 + )

3564 + spec = MockOffloadingSpec(build_offloading_config(vllm_config, kv_cache_config))

3565 + scheduler = OffloadingConnectorScheduler(spec, vllm_config, kv_cache_config)

3566 + scheduler.manager.prepare_store.side_effect = (

3567 + lambda keys, req_context: generate_store_output(keys)

3568 + )

3569 +

3570 + request = MagicMock()

3571 + request.request_id = "req"

3572 + request.kv_transfer_params = None

3573 + request.num_prompt_tokens = 1200

3574 + request.num_tokens = 1200

3575 + request.num_computed_tokens = 0

3576 + request.block_hashes = [BlockHash(f"h{i}".encode()) for i in range(150)]

3577 + request.all_token_ids = list(range(1200))

3578 + request.lora_request = None

3579 + request.shared_prefix_boundary = 0

3580 + request.status = RequestStatus.RUNNING

3581 + request.is_finished.return_value = False

3582 + scheduler.on_new_request(request)

3583 +

3584 + full_block_ids = list(range(1, 15))

3585 + swa_block_ids = list(range(1001, 1126))

3586 + first_output = SchedulerOutput.make_empty()

3587 + first_output.scheduled_new_reqs = [

3588 + SimpleNamespace(req_id="req", block_ids=(full_block_ids, swa_block_ids))

3589 + ]

3590 + first_output.num_scheduled_tokens = {"req": 1000}

3591 + first_output.total_num_scheduled_tokens = 1000

3592 + scheduler.build_connector_meta(first_output)

3593 +

3594 + request.num_computed_tokens = 1000

3595 + request.status = RequestStatus.FINISHED_ABORTED

3596 + request.is_finished.return_value = True

3597 + abort_output = SchedulerOutput.make_empty()

3598 + abort_output.finished_req_ids = {"req"}

3599 + abort_meta = scheduler.build_connector_meta(abort_output)

3600 +

3601 + assert len(abort_meta.store_jobs) == 1

3602 + [store_job] = abort_meta.store_jobs.values()

3603 + assert isinstance(store_job.src_spec, GPULoadStoreSpec)

3604 + assert store_job.src_spec.block_ids.tolist() == [swa_block_ids[123]]

3605 + assert store_job.src_spec.group_sizes == [0, 1]

3606 + assert store_job.src_spec.block_indices == [0, 123]

3607 +

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

3609 def test_full_attn_store_then_load(self, request_runner, async_scheduling: bool):

3610 """Eagle group constrains load: convergence tightens both groups.

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

+117 -15

@@ -1326,6 +1326,46 @@ class OffloadingConnectorScheduler:

1326

1327 return store_jobs

1328

1329 + def _reachable_store_block_mask(

1330 + self,

1331 + group_config: GroupOffloadConfig,

1332 + start_chunk_idx: int,

1333 + end_chunk_idx: int,

1334 + final_segment_end_chunk_idx: int | None,

1335 + reachable_boundaries: tuple[int, ...],

1336 + ) -> list[bool] | None:

1337 + """Build the block mask for a range of candidate offload chunks."""

1338 + blocks_per_chunk = self.config.blocks_per_chunk

1339 + return group_config.manager_cls.reachable_block_mask(

1340 + start_block=start_chunk_idx * blocks_per_chunk,

1341 + end_block=end_chunk_idx * blocks_per_chunk,

1342 + alignment_tokens=self.config.alignment_tokens,

1343 + kv_cache_spec=group_config.kv_cache_spec,

1344 + use_eagle=group_config.is_eagle_group,

1345 + retention_interval=self.config.retention_interval,

1346 + reachable_boundaries=reachable_boundaries,

1347 + dcp_world_size=self.config.dcp_world_size,

1348 + final_segment_end_block=(

1349 + final_segment_end_chunk_idx * blocks_per_chunk

1350 + if final_segment_end_chunk_idx is not None

1351 + else None

1352 + ),

1353 + )

1354 +

1355 + def _final_swa_alignment_blocks(

1356 + self, group_config: GroupOffloadConfig

1357 + ) -> int | None:

1358 + """Return representable final-tail alignment for a non-EAGLE SWA group."""

1359 + alignment_tokens = self.config.alignment_tokens

1360 + if (

1361 + alignment_tokens is None

1362 + or group_config.is_eagle_group

1363 + or not isinstance(group_config.kv_cache_spec, SlidingWindowSpec)

1364 + or alignment_tokens % group_config.tokens_per_block != 0

1365 + ):

1366 + return None

1367 + return alignment_tokens // group_config.tokens_per_block

1368 +

1369 def _build_store_jobs(

1370 self,

1371 scheduler_output: SchedulerOutput,

@@ -1352,10 +1392,14 @@ class OffloadingConnectorScheduler:

1392 num_offloadable_tokens = self._calc_num_offloadable_tokens(

1393 req_status, num_tokens_after_batch

1394 )

1395 + prompt_offloadable_tokens = self._calc_num_offloadable_tokens(

1396 + req_status, req.num_prompt_tokens

1397 + )

1398

1399 # Filter out chunks skipped due to sliding window attention / SSM

1400 # or unreachable by the load path's alignment constraints.

1401 new_offload_keys: list[OffloadKey] = []

1402 + group_store_ranges: list[tuple[int, int]] = []

1403

1404 reachable_boundaries: tuple[int, ...] = ()

1405 if self.config.retention_interval is not None:

@@ -1370,10 +1414,47 @@ class OffloadingConnectorScheduler:

1414 group_config, group_state, num_offloadable_tokens

1415 )

1416

1417 + start_chunk_idx = group_state.next_stored_chunk_idx

1418 + prompt_horizon_chunks = (

1419 + prompt_offloadable_tokens // group_config.tokens_per_chunk

1420 + )

1421 + final_swa_alignment_blocks = self._final_swa_alignment_blocks(

1422 + group_config

1423 + )

1424 + reconsider_final_swa_tail = (

1425 + req.is_finished()

1426 + and num_chunks != prompt_horizon_chunks

1427 + and self.config.retention_interval is None

1428 + and final_swa_alignment_blocks is not None

1429 + )

1430 + if (

1431 + self.config.retention_interval is not None

1432 + or group_config.is_eagle_group

1433 + ):

1434 + store_horizon_chunks = None

1435 + elif req.is_finished():

1436 + store_horizon_chunks = num_chunks

1437 + elif num_chunks <= prompt_horizon_chunks:

1438 + store_horizon_chunks = prompt_horizon_chunks

1439 + else:

1440 + # An active decode frontier is not a final request

1441 + # boundary. Only fixed alignment tails are reachable.

1442 + store_horizon_chunks = None

1443 + if reconsider_final_swa_tail:

1444 + assert final_swa_alignment_blocks is not None

1445 + horizon_blocks = num_chunks * blocks_per_chunk

1446 + partial_segment_start_block = (

1447 + horizon_blocks - horizon_blocks % final_swa_alignment_blocks

1448 + )

1449 + partial_segment_start_chunk = (

1450 + partial_segment_start_block // blocks_per_chunk

1451 + )

1452 + start_chunk_idx = min(start_chunk_idx, partial_segment_start_chunk)

1453 + group_store_ranges.append((start_chunk_idx, num_chunks))

1454 +

1455 if group_config.requires_cow_source:

1456 continue

1457

1376 - start_chunk_idx = group_state.next_stored_chunk_idx

1458 if num_chunks <= start_chunk_idx:

1459 continue

1460 offload_keys = group_state.offload_keys[start_chunk_idx:num_chunks]

@@ -1393,22 +1474,44 @@ class OffloadingConnectorScheduler:

1474 # (SWA/Mamba sparsity + retention interval).

1475 # reachable_block_mask operates in KV-block coordinates,

1476 # so convert chunk indices to block indices.

1396 - block_mask = group_config.manager_cls.reachable_block_mask(

1397 - start_block=start_chunk_idx * blocks_per_chunk,

1398 - end_block=num_chunks * blocks_per_chunk,

1399 - alignment_tokens=self.config.alignment_tokens,

1400 - kv_cache_spec=group_config.kv_cache_spec,

1401 - use_eagle=group_config.is_eagle_group,

1402 - retention_interval=self.config.retention_interval,

1477 + block_mask = self._reachable_store_block_mask(

1478 + group_config=group_config,

1479 + start_chunk_idx=start_chunk_idx,

1480 + end_chunk_idx=num_chunks,

1481 + final_segment_end_chunk_idx=store_horizon_chunks,

1482 reachable_boundaries=reachable_boundaries,

1404 - dcp_world_size=self.config.dcp_world_size,

1483 )

1484

1485 + prompt_horizon_block_mask: list[bool] | None = None

1486 + if reconsider_final_swa_tail:

1487 + prompt_horizon_block_mask = self._reachable_store_block_mask(

1488 + group_config=group_config,

1489 + start_chunk_idx=start_chunk_idx,

1490 + end_chunk_idx=num_chunks,

1491 + final_segment_end_chunk_idx=prompt_horizon_chunks,

1492 + reachable_boundaries=reachable_boundaries,

1493 + )

1494 +

1495 for key_idx, (offload_key, block_id) in enumerate(

1496 zip(offload_keys, offload_block_ids)

1497 ):

1498 if block_id == 0:

1499 continue

1500 + abs_chunk_idx = start_chunk_idx + key_idx

1501 + if (

1502 + reconsider_final_swa_tail

1503 + and abs_chunk_idx < group_state.next_stored_chunk_idx

1504 + and (

1505 + prompt_horizon_block_mask is None

1506 + or any(

1507 + prompt_horizon_block_mask[

1508 + key_idx * blocks_per_chunk + block_idx

1509 + ]

1510 + for block_idx in range(blocks_per_chunk)

1511 + )

1512 + )

1513 + ):

1514 + continue

1515 # A chunk is reachable if any of its constituent

1516 # blocks is reachable.

1517 if block_mask is not None and not any(

@@ -1445,16 +1548,15 @@ class OffloadingConnectorScheduler:

1548 src_block_ids: list[int] = []

1549 fenced_block_ids: list[int] = []

1550 deferred_fence_block_ids: list[int] = []

1448 - for group_config, group_state in zip(

1449 - self.config.kv_group_configs, req_status.group_states

1551 + for group_config, group_state, store_range in zip(

1552 + self.config.kv_group_configs,

1553 + req_status.group_states,

1554 + group_store_ranges,

1555 ):

1556 is_sliding_window = (

1557 group_config.sliding_window_size_in_chunks is not None

1558 )

1454 - num_chunks = req_status.storable_chunks(

1455 - group_config, group_state, num_offloadable_tokens

1456 - )

1457 - start_chunk_idx = group_state.next_stored_chunk_idx

1559 + start_chunk_idx, num_chunks = store_range

1560 block_ids = group_state.block_ids

1561 num_group_blocks = 0

1562 start_gpu_block_idx: int | None = None

vllm/v1/core/single_type_kv_cache_manager.py

+18

@@ -508,6 +508,7 @@ class SingleTypeKVCacheManager(ABC):

508 retention_interval: int | None = None,

509 reachable_boundaries: Sequence[int] = (),

510 dcp_world_size: int = 1,

511 + final_segment_end_block: int | None = None,

512 ) -> list[bool] | None:

513 """Per-block mask for ``cache_full_blocks``. ``None`` means cache

514 every (non-null) block — the default for full attention.

@@ -516,6 +517,10 @@ class SingleTypeKVCacheManager(ABC):

517 blocks that can never serve a hit at any alignment-aligned prefix length.

518 ``reachable_boundaries`` are token positions whose reachable tail must be

519 retained; the base (dense) policy ignores them.

520 + ``final_segment_end_block`` is the exclusive end of the request's final

521 + segment. It may be later than ``end_block`` while prefill is still in

522 + progress. Non-EAGLE sparse managers may keep its reachable tail. The

523 + base policy already keeps every block, so it ignores this value.

524 """

525 return None

526

@@ -1033,6 +1038,7 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

1038 retention_interval: int | None = None,

1039 reachable_boundaries: Sequence[int] = (),

1040 dcp_world_size: int = 1,

1041 + final_segment_end_block: int | None = None,

1042 ) -> list[bool] | None:

1043 assert isinstance(kv_cache_spec, SlidingWindowSpec)

1044 if alignment_tokens is None:

@@ -1079,6 +1085,17 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

1085 if i >= shift and (i - shift) % per_segment >= per_segment - need:

1086 mask[i - start_block] = True

1087

1088 + if final_segment_end_block is not None and not use_eagle:

1089 + final_segment_size = final_segment_end_block % per_segment

1090 + if final_segment_size:

1091 + final_segment_start = final_segment_end_block - final_segment_size

1092 + final_tail_start = max(

1093 + final_segment_start, final_segment_end_block - need

1094 + )

1095 + final_tail_end = min(end_block, final_segment_end_block)

1096 + for i in range(max(start_block, final_tail_start), final_tail_end):

1097 + mask[i - start_block] = True

1098 +

1099 # (2) Reachable-boundary tails: the replay boundary (``num_prompt - 1``,

1100 # capped by ``get_computed_blocks``) and any shared-prefix junction. Both

1101 # land before segments would cover them under sparse retention, so keep

@@ -1513,6 +1530,7 @@ class MambaManager(SingleTypeKVCacheManager):

1530 retention_interval: int | None = None,

1531 reachable_boundaries: Sequence[int] = (),

1532 dcp_world_size: int = 1,

1533 + final_segment_end_block: int | None = None,

1534 ) -> list[bool] | None:

1535 """Sparse Mamba state-snapshot retention.

1536