@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 25268f0c9c

[KVConnector] Add retention interval to OffloadingConnector (#51886)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

bnellnm committed Sep 3, 2026 at 22:42 UTC 25268f0c9c173aa6557a3eebf7d7a87fa91fec3a

9 files changed +543 -139

tests/v1/core/test_prefix_caching.py

+155

@@ -4433,6 +4433,161 @@ def test_swa_reachable_block_mask_pins_shared_prefix():

4433 assert retained(0, 0, block_size) == {14}

4434

4435

4436 +def test_swa_reachable_block_mask_with_dcp_scaling():

4437 + """DCP shards each block's KV across ranks, scaling the effective block size.

4438 + Verify that dcp_world_size > 1 scales block size in reachability calculations,

4439 + producing different masks than dcp_world_size=1."""

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

4441 +

4442 + block_size = 16

4443 + sliding_window = 32 # need = cdiv(31, 16) = 2

4444 +

4445 + def get_mask(dcp_world_size):

4446 + spec = SlidingWindowSpec(

4447 + block_size=block_size,

4448 + num_kv_heads=1,

4449 + head_size=1,

4450 + dtype=torch.float32,

4451 + sliding_window=sliding_window,

4452 + )

4453 + m = SlidingWindowManager.reachable_block_mask(

4454 + start_block=0,

4455 + end_block=16,

4456 + alignment_tokens=64, # alignment constraint

4457 + kv_cache_spec=spec,

4458 + use_eagle=False,

4459 + retention_interval=64, # sparse: 1 state per 64-token segment

4460 + reachable_boundaries=(255,),

4461 + dcp_world_size=dcp_world_size,

4462 + )

4463 + return None if m is None else {i for i, v in enumerate(m) if v}

4464 +

4465 + # With dcp_world_size=1: effective block_size = 16

4466 + # per_segment = 64 // 16 = 4 blocks per segment

4467 + # need = 2, so 2 < 4 -> sparse retention with coarser granularity

4468 + mask_no_dcp = get_mask(dcp_world_size=1)

4469 +

4470 + # With dcp_world_size=2: effective block_size = 32

4471 + # per_segment = 64 // 32 = 2 blocks per segment

4472 + # need = 2, so 2 >= 2 -> different retention granularity

4473 + mask_with_dcp = get_mask(dcp_world_size=2)

4474 +

4475 + # DCP scaling changes block_size, altering segment granularity

4476 + assert mask_no_dcp is not None

4477 + assert mask_with_dcp is not None

4478 + # Different DCP world sizes should produce different masks

4479 + assert mask_no_dcp != mask_with_dcp, (

4480 + "DCP scaling should change retention granularity"

4481 + )

4482 +

4483 +

4484 +@pytest.mark.parametrize(

4485 + "block_size,dcp_world_size,alignment_tokens",

4486 + [

4487 + (16, 1, 24), # 24 % 16 != 0

4488 + (16, 3, 64), # effective block_size 48; 64 % 48 != 0

4489 + ],

4490 +)

4491 +def test_swa_reachable_block_mask_sub_block_alignment_is_dense(

4492 + block_size, dcp_world_size, alignment_tokens

4493 +):

4494 + """The mask is block-granular, so a sub-block alignment_tokens (not a

4495 + multiple of the DCP-scaled block size) cannot be represented exactly. This

4496 + happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the

4497 + full-attention chunk size and need not divide the SWA block size. The mask

4498 + must fall back to dense (None) rather than raise."""

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

4500 +

4501 + spec = SlidingWindowSpec(

4502 + block_size=block_size,

4503 + num_kv_heads=1,

4504 + head_size=1,

4505 + dtype=torch.float32,

4506 + sliding_window=32,

4507 + )

4508 +

4509 + mask = SlidingWindowManager.reachable_block_mask(

4510 + start_block=0,

4511 + end_block=16,

4512 + alignment_tokens=alignment_tokens,

4513 + kv_cache_spec=spec,

4514 + use_eagle=False,

4515 + retention_interval=alignment_tokens,

4516 + reachable_boundaries=(),

4517 + dcp_world_size=dcp_world_size,

4518 + )

4519 + assert mask is None

4520 +

4521 +

4522 +def test_mamba_reachable_block_mask_ignores_dcp():

4523 + """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a

4524 + state block spans kv_cache_spec.block_size tokens regardless of DCP. The

4525 + mask must not scale by dcp_world_size, so retention granularity is

4526 + identical for any DCP world size."""

4527 + from vllm.v1.core.single_type_kv_cache_manager import MambaManager

4528 +

4529 + spec = MambaSpec(

4530 + block_size=16,

4531 + shapes=(1, 1),

4532 + dtypes=(torch.float32,),

4533 + mamba_cache_mode="align",

4534 + )

4535 +

4536 + def get_mask(dcp_world_size):

4537 + m = MambaManager.reachable_block_mask(

4538 + start_block=0,

4539 + end_block=16,

4540 + alignment_tokens=64,

4541 + kv_cache_spec=spec,

4542 + use_eagle=False,

4543 + retention_interval=64, # sparse: one state per 64-token segment

4544 + reachable_boundaries=(255,),

4545 + dcp_world_size=dcp_world_size,

4546 + )

4547 + return None if m is None else {i for i, v in enumerate(m) if v}

4548 +

4549 + # block_size stays 16 regardless of DCP: per_segment = 64 // 16 = 4.

4550 + mask_no_dcp = get_mask(dcp_world_size=1)

4551 + mask_with_dcp = get_mask(dcp_world_size=2)

4552 +

4553 + assert mask_no_dcp is not None

4554 + assert mask_with_dcp is not None

4555 + assert mask_no_dcp == mask_with_dcp, "DCP must not change Mamba retention"

4556 +

4557 +

4558 +def test_mamba_reachable_block_mask_large_dcp_stays_sparse():

4559 + """A large DCP world size must not scale the Mamba block size, so it can

4560 + never collapse sparse retention into dense caching."""

4561 + from vllm.v1.core.single_type_kv_cache_manager import MambaManager

4562 +

4563 + spec = MambaSpec(

4564 + block_size=8,

4565 + shapes=(1, 1),

4566 + dtypes=(torch.float32,),

4567 + mamba_cache_mode="align",

4568 + )

4569 +

4570 + def get_mask(dcp_world_size):

4571 + return MambaManager.reachable_block_mask(

4572 + start_block=0,

4573 + end_block=16,

4574 + alignment_tokens=64,

4575 + kv_cache_spec=spec,

4576 + use_eagle=False,

4577 + retention_interval=64,

4578 + reachable_boundaries=(255,),

4579 + dcp_world_size=dcp_world_size,

4580 + )

4581 +

4582 + # block_size stays 8 regardless of DCP: per_segment = 64 // 8 = 8, sparse.

4583 + mask_no_dcp = get_mask(dcp_world_size=1)

4584 + mask_large_dcp = get_mask(dcp_world_size=16)

4585 +

4586 + assert mask_no_dcp is not None, "Sparse retention expected"

4587 + assert mask_large_dcp is not None, "Large DCP must not force dense caching"

4588 + assert mask_no_dcp == mask_large_dcp, "DCP must not change Mamba retention"

4589 +

4590 +

4591 def test_swa_shared_prefix_reuse_under_zero_retention():

4592 """SWA cross-request analog: a partial shared prefix's sliding-window tail

4593 must stay reusable under ``prefix_cache_retention_interval=0``. Without

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

+14 -13

@@ -15,7 +15,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (

15 )

16 from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

17 SchedulerOffloadConfig,

18 - is_store_reachable_swa_chunk,

18 )

19 from vllm.platforms import current_platform

20 from vllm.v1.kv_cache_interface import (

@@ -31,6 +30,7 @@ from vllm.v1.kv_cache_interface import (

30 SlidingWindowSpec,

31 UniformTypeKVCacheSpecs,

32 )

33 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

36 def _make_vllm_config(

@@ -46,6 +46,7 @@ def _make_vllm_config(

46 config.cache_config.enable_prefix_caching = True

47 config.cache_config.prefix_match_unit = None

48 config.cache_config.cache_dtype = torch.float16

49 + config.cache_config.prefix_cache_retention_interval = None

50 config.model_config.model = "test-model"

51 config.model_config.use_mla = False

52 # _full_attention_spec's heads at tp=1: the parallelism-agnostic gate

@@ -375,18 +376,18 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks():

376 _make_mamba_hybrid_kv_cache_config(),

377 )

378 mamba_group = scheduler_config.kv_group_configs[1]

378 - assert mamba_group.alignment_chunk_count == 2

379 - assert [

380 - chunk_idx

381 - for chunk_idx in range(4)

382 - if is_store_reachable_swa_chunk(

383 - chunk_idx,

384 - 4,

385 - mamba_group.alignment_chunk_count,

386 - mamba_group.sliding_window_size_in_chunks,

387 - mamba_group.is_eagle_group,

388 - )

389 - ] == [1, 3]

379 + assert mamba_group.sliding_window_size_in_chunks == 1

380 + assert scheduler_config.alignment_tokens is not None

381 + manager_cls = KVCacheSpecRegistry.get_manager_class(mamba_group.kv_cache_spec)

382 + assert manager_cls is not None

383 + block_mask = manager_cls.reachable_block_mask(

384 + start_block=0,

385 + end_block=4,

386 + alignment_tokens=scheduler_config.alignment_tokens,

387 + kv_cache_spec=mamba_group.kv_cache_spec,

388 + use_eagle=mamba_group.is_eagle_group,

389 + )

390 + assert block_mask is None or [i for i in range(4) if block_mask[i]] == [1, 3]

391

392

393 @pytest.mark.parametrize("dcp_size,expected", [(1, 16), (2, 32)])

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

+11

@@ -29,6 +29,7 @@ from vllm.v1.kv_cache_interface import (

29 KVCacheGroupSpec,

30 KVCacheSpecKind,

31 )

32 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

33 from vllm.v1.kv_offload.base import (

34 Locality,

35 Medium,

@@ -88,12 +89,22 @@ def _group_config(

89 tokens_per_hash = block_size

90 tokens_per_chunk = block_size * blocks_per_chunk

91 assert tokens_per_chunk % tokens_per_hash == 0

92 + kv_spec = FullAttentionSpec(

93 + block_size=block_size,

94 + num_kv_heads=1,

95 + head_size=1,

96 + dtype=torch.float32,

97 + )

98 + manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec)

99 + assert manager_cls is not None

100 return GroupOffloadConfig(

101 group_idx=group_idx,

102 tokens_per_block=block_size,

103 tokens_per_chunk=tokens_per_chunk,

104 hashes_per_chunk=tokens_per_chunk // tokens_per_hash,

105 sliding_window_size_in_chunks=sliding_window_size_in_chunks,

106 + kv_cache_spec=kv_spec,

107 + manager_cls=manager_cls,

108 kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC,

109 )

110

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

+238 -70

@@ -33,7 +33,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (

33 OffloadingConnectorScheduler,

34 RequestOffloadState,

35 get_sliding_window_size_in_chunks,

36 - is_store_reachable_swa_chunk,

36 )

37 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

38 from vllm.v1.core.kv_cache_utils import (

@@ -52,6 +51,7 @@ from vllm.v1.kv_cache_interface import (

51 KVCacheGroupSpec,

52 SlidingWindowSpec,

53 )

54 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

55 from vllm.v1.kv_offload.base import (

56 GPULoadStoreSpec,

57 LookupResult,

@@ -498,55 +498,6 @@ def test_scheduler_reports_lookup_sync_delay(request_runner):

498 assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0

499

500

501 -@pytest.mark.parametrize(

502 - (

503 - "absolute_chunk_index",

504 - "storable_chunk_count",

505 - "alignment_chunk_count",

506 - "sliding_window_chunks",

507 - "is_eagle_group",

508 - "expected",

509 - ),

510 - [

511 - # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61.

512 - (61, 64, 64, 2, False, False),

513 - (62, 64, 64, 2, False, True),

514 - (60, 64, 64, 2, True, False),

515 - (61, 64, 64, 2, True, True),

516 - # Partial 48-of-64 segment: the reachable tail ends at chunk 47.

517 - (45, 48, 64, 2, False, False),

518 - (46, 48, 64, 2, False, True),

519 - (44, 48, 64, 2, True, False),

520 - (45, 48, 64, 2, True, True),

521 - # A later partial segment uses its own actual end (chunks 64-79).

522 - (76, 80, 64, 3, False, False),

523 - (77, 80, 64, 3, False, True),

524 - # No alignment means no store-pruning optimization.

525 - (0, 1, None, None, False, True),

526 - # A tail at least as large as the segment keeps every chunk.

527 - (0, 2, 64, 2, False, True),

528 - ],

529 -)

530 -def test_is_store_reachable_swa_chunk(

531 - absolute_chunk_index: int,

532 - storable_chunk_count: int,

533 - alignment_chunk_count: int | None,

534 - sliding_window_chunks: int | None,

535 - is_eagle_group: bool,

536 - expected: bool,

537 -):

538 - assert (

539 - is_store_reachable_swa_chunk(

540 - absolute_chunk_index,

541 - storable_chunk_count,

542 - alignment_chunk_count,

543 - sliding_window_chunks,

544 - is_eagle_group,

545 - )

546 - is expected

547 - )

548 -

549 -

501 def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner):

502 """A deferred lookup reports its async delay once it resolves."""

503 runner = request_runner(

@@ -1732,20 +1683,15 @@ def test_sliding_window_demand_is_store_reachable(

1683 for lookup_call in sched.manager.lookup.call_args_list

1684 }

1685 assert demanded, "the scan must query something"

1735 - unsuppliable = sorted(

1736 - chunk_idx

1737 - for chunk_idx in demanded

1738 - if not is_store_reachable_swa_chunk(

1739 - chunk_idx,

1740 - storable_chunk_count,

1741 - alignment_chunk_count,

1742 - sliding_window_chunks,

1743 - is_eagle_group,

1744 - )

1745 - )

1746 - assert not unsuppliable, (

1747 - f"chunks {unsuppliable} are demanded by the load path but pruned by the "

1748 - f"store path, so a warm producer can never supply them"

1686 +

1687 + # Verify that the lookup path demands a reasonable set of chunks.

1688 + # The lookup path scans forward with a window of size sliding_window_chunks + 1

1689 + # (for EAGLE). With dense storage (retention_interval=None), all chunks should be

1690 + # available. This test validates that the lookup path doesn't demand chunks beyond

1691 + # the storable range.

1692 + assert all(0 <= chunk_idx < storable_chunk_count for chunk_idx in demanded), (

1693 + f"Lookup path demanded chunks {demanded} outside valid range "

1694 + f"[0, {storable_chunk_count})"

1695 )

1696

1697

@@ -2392,7 +2338,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):

2338 - Group 0: full attention (MLA-like), block_size=16

2339 - Group 1: SWA, block_size=4, sliding_window=8

2340

2395 - alignment_chunk_count = 16 / 4 = 4 SWA blocks per alignment segment.

2341 + alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment.

2342 sliding_window_size_in_chunks = ceil(8 / 4) = 2.

2343 Within each segment of 4 SWA blocks, only the trailing 2 are stored.

2344

@@ -2436,17 +2382,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool):

2382 kv_cache_groups=kv_cache_groups,

2383 )

2384

2439 - # Verify config: alignment_chunk_count computed correctly

2385 + # Verify config

2386 kv_group_configs = runner.connector_scheduler.config.kv_group_configs

2387 assert len(kv_group_configs) == 2

2442 - # Group 0: full attention -> no alignment skip

2443 - assert kv_group_configs[0].alignment_chunk_count is None

2388 + # Group 0: full attention

2389 assert kv_group_configs[0].sliding_window_size_in_chunks is None

2390 assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size

2446 - # Group 1: SWA -> alignment_chunk_count = 16/4 = 4, tail = 2

2447 - assert kv_group_configs[1].alignment_chunk_count == 4

2391 + # Group 1: SWA with sliding_window_size of 2 chunks

2392 assert kv_group_configs[1].sliding_window_size_in_chunks == 2

2393 assert kv_group_configs[1].tokens_per_chunk == swa_block_size

2394 + # alignment_tokens = full_attn_block_size = 16

2395 + assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size

2396

2397 # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks

2398 # (block_size=4). Decode 1 token to kick off processing (stores are

@@ -3782,3 +3728,225 @@ def test_chunked_local_attention_reports_its_chunk_window():

3728 assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8

3729 # Partial chunks round up, so the reachable tail is never understated.

3730 assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3

3731 +

3732 +

3733 +@pytest.mark.parametrize("blocks_per_chunk", [1, 2, 3])

3734 +def test_reachable_block_mask_chunk_to_block_index_conversion(

3735 + blocks_per_chunk: int,

3736 +):

3737 + """reachable_block_mask operates in KV-block coordinates, but the

3738 + offloading scheduler works in chunk coordinates. Verify that the

3739 + index conversion (chunk -> block -> chunk) produces the same

3740 + filtering result regardless of blocks_per_chunk.

3741 +

3742 + Uses a SWA spec with alignment_tokens large enough to create

3743 + unreachable blocks (alignment_tokens > sliding_window).

3744 + """

3745 + block_size = 4

3746 + sliding_window = 8

3747 + alignment_tokens = 16 * blocks_per_chunk

3748 +

3749 + swa_spec = SlidingWindowSpec(

3750 + block_size=block_size,

3751 + num_kv_heads=1,

3752 + head_size=1,

3753 + dtype=torch.float32,

3754 + sliding_window=sliding_window,

3755 + )

3756 +

3757 + num_chunks = 8

3758 + start_chunk_idx = 0

3759 +

3760 + manager_cls = KVCacheSpecRegistry.get_manager_class(swa_spec)

3761 + assert manager_cls is not None

3762 +

3763 + block_mask = manager_cls.reachable_block_mask(

3764 + start_block=start_chunk_idx * blocks_per_chunk,

3765 + end_block=num_chunks * blocks_per_chunk,

3766 + alignment_tokens=alignment_tokens,

3767 + kv_cache_spec=swa_spec,

3768 + use_eagle=False,

3769 + )

3770 +

3771 + if block_mask is None:

3772 + return

3773 +

3774 + assert len(block_mask) == num_chunks * blocks_per_chunk

3775 +

3776 + reachable_chunks = []

3777 + for chunk_idx in range(num_chunks):

3778 + if any(

3779 + block_mask[chunk_idx * blocks_per_chunk + b]

3780 + for b in range(blocks_per_chunk)

3781 + ):

3782 + reachable_chunks.append(chunk_idx)

3783 +

3784 + # The tail of each alignment segment should be reachable.

3785 + # sliding_window=8 with block_size=4 -> need=2 blocks.

3786 + # With blocks_per_chunk=1: alignment has 16 blocks, tail 2 reachable.

3787 + # With blocks_per_chunk=2: alignment has 16 blocks, tail 2 -> 1 chunk.

3788 + # With blocks_per_chunk=3: alignment has 16 blocks, tail 2 -> 1 chunk.

3789 + assert len(reachable_chunks) > 0

3790 + assert len(reachable_chunks) < num_chunks

3791 +

3792 +

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

3794 +def test_retention_interval_sparsifies_swa_stores(

3795 + request_runner, async_scheduling: bool

3796 +):

3797 + """retention_interval reduces the number of SWA chunks stored.

3798 +

3799 + Uses the same hybrid architecture as test_swa_alignment_skip but with

3800 + retention_interval=32 (read from kv_connector_extra_config). With

3801 + alignment_tokens=16 the default stores a 2-block SWA tail per 16-token

3802 + segment. retention_interval=32 widens segments to 32 tokens (8 SWA

3803 + blocks), so only the trailing 2 of each 8-block segment are stored.

3804 +

3805 + With 64 tokens (4 full-attn blocks, 16 SWA blocks):

3806 + - Group 0 (full attn): stores all 4 blocks

3807 + - Group 1 (SWA) without retention: 4 segments of 4 -> 8 stored

3808 + - Group 1 (SWA) with retention_interval=32: 2 segments of 8 -> 4 stored

3809 + """

3810 + full_attn_block_size = 16

3811 + swa_block_size = 4

3812 + sliding_window = 8

3813 + num_gpu_blocks = 200

3814 +

3815 + kv_cache_groups = [

3816 + KVCacheGroupSpec(

3817 + ["layer0"],

3818 + FullAttentionSpec(

3819 + block_size=full_attn_block_size,

3820 + num_kv_heads=1,

3821 + head_size=1,

3822 + dtype=torch.float32,

3823 + ),

3824 + ),

3825 + KVCacheGroupSpec(

3826 + ["layer1"],

3827 + SlidingWindowSpec(

3828 + block_size=swa_block_size,

3829 + num_kv_heads=1,

3830 + head_size=1,

3831 + dtype=torch.float32,

3832 + sliding_window=sliding_window,

3833 + ),

3834 + ),

3835 + ]

3836 +

3837 + runner = request_runner(

3838 + block_size=swa_block_size,

3839 + num_gpu_blocks=num_gpu_blocks,

3840 + async_scheduling=async_scheduling,

3841 + kv_cache_groups=kv_cache_groups,

3842 + retention_interval=32,

3843 + )

3844 +

3845 + assert runner.connector_scheduler.config.retention_interval == 32

3846 +

3847 + num_tokens = 64

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

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

3850 + generate_store_output(keys)

3851 + )

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

3853 +

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

3855 + generate_store_output(keys)

3856 + )

3857 + runner.run(

3858 + decoded_tokens=[EOS_TOKEN_ID],

3859 + # Group 0 (full attn, block_size=16): 4 blocks stored

3860 + # Group 1 (SWA, block_size=4): 16 blocks total, retention_interval=32

3861 + # Dense segment tails (per_segment=32/4=8, need=2):

3862 + # blocks 6, 7 and 14, 15

3863 + # Replay boundary (num_prompt_tokens-1=63, aligned to 48):

3864 + # end=48/4=12, tail blocks 10, 11

3865 + expected_stored=(

3866 + (0, 0),

3867 + (0, 1),

3868 + (0, 2),

3869 + (0, 3),

3870 + (1, 6),

3871 + (1, 7),

3872 + (1, 10),

3873 + (1, 11),

3874 + (1, 14),

3875 + (1, 15),

3876 + ),

3877 + )

3878 +

3879 +

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

3881 +def test_retention_interval_zero_stores_only_replay_boundary(

3882 + request_runner, async_scheduling: bool

3883 +):

3884 + """retention_interval=0 stores only the replay boundary SWA blocks.

3885 +

3886 + Same hybrid architecture. With retention_interval=0, no dense segment

3887 + tails are stored — only blocks reachable from the replay boundary

3888 + (num_prompt_tokens - 1).

3889 + """

3890 + full_attn_block_size = 16

3891 + swa_block_size = 4

3892 + sliding_window = 8

3893 + num_gpu_blocks = 200

3894 +

3895 + kv_cache_groups = [

3896 + KVCacheGroupSpec(

3897 + ["layer0"],

3898 + FullAttentionSpec(

3899 + block_size=full_attn_block_size,

3900 + num_kv_heads=1,

3901 + head_size=1,

3902 + dtype=torch.float32,

3903 + ),

3904 + ),

3905 + KVCacheGroupSpec(

3906 + ["layer1"],

3907 + SlidingWindowSpec(

3908 + block_size=swa_block_size,

3909 + num_kv_heads=1,

3910 + head_size=1,

3911 + dtype=torch.float32,

3912 + sliding_window=sliding_window,

3913 + ),

3914 + ),

3915 + ]

3916 +

3917 + runner = request_runner(

3918 + block_size=swa_block_size,

3919 + num_gpu_blocks=num_gpu_blocks,

3920 + async_scheduling=async_scheduling,

3921 + kv_cache_groups=kv_cache_groups,

3922 + retention_interval=0,

3923 + )

3924 +

3925 + assert runner.connector_scheduler.config.retention_interval == 0

3926 +

3927 + num_tokens = 64

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

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

3930 + generate_store_output(keys)

3931 + )

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

3933 +

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

3935 + generate_store_output(keys)

3936 + )

3937 + runner.run(

3938 + decoded_tokens=[EOS_TOKEN_ID],

3939 + # Group 0 (full attn): all 4 blocks stored

3940 + # Group 1 (SWA): retention_interval=0 -> only replay boundary

3941 + # replay boundary = num_prompt_tokens - 1 = 63

3942 + # aligned to alignment_tokens=16 -> 48, end=48/4=12

3943 + # tail = ceil(8/4) = 2 blocks -> blocks 10, 11

3944 + expected_stored=(

3945 + (0, 0),

3946 + (0, 1),

3947 + (0, 2),

3948 + (0, 3),

3949 + (1, 10),

3950 + (1, 11),

3951 + ),

3952 + )

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

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

180 kv_cache_groups: list[KVCacheGroupSpec] | None = None,

181 extra_config_overrides: dict[str, Any] | None = None,

182 worker_count: int = 1,

183 + retention_interval: int | None = None,

184 ):

185 assert blocks_per_chunk == 1 or kv_cache_groups is None, (

186 "blocks_per_chunk > 1 requires all groups to have the same "

@@ -200,6 +201,7 @@ class RequestRunner:

201 )

202 vllm_config.scheduler_config.async_scheduling = async_scheduling

203 vllm_config.parallel_config.world_size = worker_count

204 + vllm_config.cache_config.prefix_cache_retention_interval = retention_interval

205

206 extra_config: dict[str, Any] = {

207 "spec_name": "MockOffloadingSpec",

@@ -676,6 +678,7 @@ def request_runner():

678 kv_cache_groups=None,

679 extra_config_overrides=None,

680 worker_count=1,

681 + retention_interval=None,

682 ):

683 runner = RequestRunner(

684 block_size=block_size,

@@ -685,6 +688,7 @@ def request_runner():

688 kv_cache_groups=kv_cache_groups,

689 extra_config_overrides=extra_config_overrides,

690 worker_count=worker_count,

691 + retention_interval=retention_interval,

692 )

693 runners.append(runner)

694 return runner

tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py

+32

@@ -287,6 +287,38 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary():

287 assert masks[1] == [False, False, False, True, False, False, False, True]

288

289

290 +def test_store_mask_swa_does_not_double_scale_under_dcp():

291 + """The coordinator receives DCP-resolved specs (worker.py applies

292 + resolve_dcp_kv_cache_spec) and indexes chunks in that already-scaled block

293 + size. So dcp_world_size must NOT further scale reachable_block_mask; the

294 + mask must be identical for any dcp_world_size given the same specs.

295 +

296 + Regression: previously the coordinator forwarded its dcp_world_size, which

297 + double-scaled the SWA block size and produced a wrong mask when dcp > 1.

298 + """

299 + full = _full(32)

300 + swa = _swa(block_size=8, sliding_window=8)

301 + groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)]

302 + scheduler_block_size = lcm(32, 8)

303 +

304 + def make(dcp_world_size):

305 + return MooncakeStoreCoordinator(

306 + groups,

307 + scheduler_block_size=scheduler_block_size,

308 + hash_block_size=8,

309 + dcp_world_size=dcp_world_size,

310 + )

311 +

312 + mask_dcp1 = make(1).store_mask(64)

313 + mask_dcp2 = make(2).store_mask(64)

314 +

315 + # dcp must not change the mask; both match the correct single-block-size

316 + # result (chunks ending at 32 and 64 -> blocks 3 and 7).

317 + expected = [False, False, False, True, False, False, False, True]

318 + assert mask_dcp1[1] == expected

319 + assert mask_dcp2[1] == expected

320 +

321 +

322 def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm():

323 """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm

324 boundary should now span two SWA tail blocks."""

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -298,6 +298,10 @@ class MooncakeStoreCoordinator:

298 use_eagle=use_eagle,

299 retention_interval=retention_interval,

300 reachable_boundaries=reachable_boundaries,

301 + # ``spec`` is already DCP-resolved (worker.py applies

302 + # resolve_dcp_kv_cache_spec) and ``end_chunk`` is indexed in

303 + # that scaled block size, so the mask must not scale again.

304 + dcp_world_size=1,

305 )

306 if mask is not None:

307 assert len(mask) == end_chunk - start_chunk

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

+72 -54

@@ -30,6 +30,7 @@ from vllm.logger import init_logger

30 from vllm.utils.math_utils import cdiv, round_down

31 from vllm.v1.core.kv_cache_manager import KVCacheBlocks

32 from vllm.v1.core.sched.output import SchedulerOutput

33 +from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager

34 from vllm.v1.kv_cache_interface import (

35 ChunkedLocalAttentionSpec,

36 FullAttentionSpec,

@@ -38,6 +39,7 @@ from vllm.v1.kv_cache_interface import (

39 MambaSpec,

40 SlidingWindowSpec,

41 )

42 +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry

43 from vllm.v1.kv_offload.base import (

44 GPULoadStoreSpec,

45 Locality,

@@ -91,15 +93,12 @@ class GroupOffloadConfig(NamedTuple):

93 kv_event_group_spec: OffloadingEventGroupSpec

94 # None below means full attention

95 sliding_window_size_in_chunks: int | None

96 + kv_cache_spec: KVCacheSpec

97 + # Cached manager class for this group's KV cache spec, resolved at init time

98 + manager_cls: type[SingleTypeKVCacheManager]

99 # Partial-tail data for this group comes from the scheduler's CoW hand-off

100 # rather than the request block table.

101 requires_cow_source: bool = False

97 - # Number of this group's offloaded chunks per full-attention alignment

98 - # segment. Used to skip storing SWA chunks that can never serve a load

99 - # hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes

100 - # than the MLA full-attention group).

101 - # None for full-attention groups or when the optimization doesn't apply.

102 - alignment_chunk_count: int | None = None

102 # True for EAGLE/MTP draft-model attention groups. The trailing chunk

103 # of these groups is volatile and lacks a stable hash, so it must

104 # be excluded from store and load scheduling.

@@ -126,26 +125,6 @@ def get_sliding_window_size_in_chunks(

125 return None

126

127

129 -def is_store_reachable_swa_chunk(

130 - absolute_chunk_index: int,

131 - storable_chunk_count: int,

132 - alignment_chunk_count: int | None,

133 - sliding_window_chunks: int | None,

134 - is_eagle_group: bool,

135 -) -> bool:

136 - """Return whether an SWA chunk can participate in an external-cache hit."""

137 - if alignment_chunk_count is None:

138 - return True

139 - assert sliding_window_chunks is not None

140 - position_in_segment = absolute_chunk_index % alignment_chunk_count

141 - segment_start = absolute_chunk_index - position_in_segment

142 - actual_segment_length = min(

143 - alignment_chunk_count, storable_chunk_count - segment_start

144 - )

145 - reachable_tail = sliding_window_chunks + int(is_eagle_group)

146 - return position_in_segment >= actual_segment_length - reachable_tail

147 -

148 -

128 def resolve_mamba_align_size(

129 spec: "OffloadingSpec", kv_cache_config: KVCacheConfig

130 ) -> int | None:

@@ -176,6 +155,9 @@ class SchedulerOffloadConfig(NamedTuple):

155 num_workers: int

156 offload_prompt_only: bool

157 supports_partial_tail: bool

158 + alignment_tokens: int | None = None

159 + retention_interval: int | None = None

160 + dcp_world_size: int = 1

161

162 @classmethod

163 def from_spec(

@@ -199,23 +181,21 @@ class SchedulerOffloadConfig(NamedTuple):

181 full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk)

182

183 # Only apply the optimization if there's a single consistent

202 - # full-attention alignment size.

184 + # full-attention alignment size. For pure SWA/Mamba models (no

185 + # full-attention groups), fall back to the largest SWA/Mamba

186 + # tokens_per_chunk so retention_interval has a boundary

187 + # granularity to work with.

188 alignment_tokens: int | None = None

189 if len(full_attn_tokens_per_chunk) == 1:

190 alignment_tokens = full_attn_tokens_per_chunk.pop()

191 + elif not full_attn_tokens_per_chunk:

192 + all_tokens_per_chunk = [

193 + tpb * spec.blocks_per_chunk for tpb in spec.tokens_per_block

194 + ]

195 + if all_tokens_per_chunk:

196 + alignment_tokens = max(all_tokens_per_chunk)

197

207 - def _alignment_chunk_count(

208 - tokens_per_chunk: int,

209 - sliding_window_size_in_chunks: int | None,

210 - ) -> int | None:

211 - if alignment_tokens is None or sliding_window_size_in_chunks is None:

212 - return None

213 - if alignment_tokens <= tokens_per_chunk:

214 - return None

215 - per_segment = alignment_tokens // tokens_per_chunk

216 - if sliding_window_size_in_chunks >= per_segment:

217 - return None

218 - return per_segment

198 + retention_interval = vllm_config.cache_config.prefix_cache_retention_interval

199

200 eagle_groups = {

201 idx

@@ -245,6 +225,10 @@ class SchedulerOffloadConfig(NamedTuple):

225 sw = get_sliding_window_size_in_chunks(

226 kv_spec, tokens_per_block * spec.blocks_per_chunk

227 )

228 + manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec)

229 + assert manager_cls is not None, (

230 + f"No manager found for KV cache spec {type(kv_spec).__name__}"

231 + )

232 kv_group_configs_list.append(

233 GroupOffloadConfig(

234 group_idx=idx,

@@ -255,9 +239,8 @@ class SchedulerOffloadConfig(NamedTuple):

239 // spec.tokens_per_hash

240 ),

241 sliding_window_size_in_chunks=sw,

258 - alignment_chunk_count=_alignment_chunk_count(

259 - tokens_per_block * spec.blocks_per_chunk, sw

260 - ),

242 + kv_cache_spec=kv_spec,

243 + manager_cls=manager_cls,

244 kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group),

245 is_eagle_group=idx in eagle_groups,

246 requires_cow_source=(

@@ -289,6 +272,23 @@ class SchedulerOffloadConfig(NamedTuple):

272 and vllm_config.parallel_config.decode_context_parallel_size == 1

273 )

274

275 + if retention_interval is not None:

276 + if retention_interval < 0:

277 + raise ValueError(

278 + f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL "

279 + f"({retention_interval}) must be non-negative."

280 + )

281 + for config in kv_group_configs:

282 + if (

283 + config.sliding_window_size_in_chunks is not None

284 + and retention_interval % config.tokens_per_chunk != 0

285 + ):

286 + raise ValueError(

287 + f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL "

288 + f"({retention_interval}) must be a multiple of "

289 + f"tokens_per_chunk ({config.tokens_per_chunk})."

290 + )

291 +

292 return cls(

293 num_workers=vllm_config.parallel_config.world_size,

294 kv_group_configs=kv_group_configs,

@@ -296,6 +296,9 @@ class SchedulerOffloadConfig(NamedTuple):

296 tokens_per_hash=spec.tokens_per_hash,

297 offload_prompt_only=spec.offload_prompt_only,

298 supports_partial_tail=supports_partial_tail,

299 + alignment_tokens=alignment_tokens,

300 + retention_interval=retention_interval,

301 + dcp_world_size=vllm_config.parallel_config.decode_context_parallel_size,

302 )

303

304

@@ -1353,6 +1356,13 @@ class OffloadingConnectorScheduler:

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

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

1358 new_offload_keys: list[OffloadKey] = []

1359 +

1360 + reachable_boundaries: tuple[int, ...] = ()

1361 + if self.config.retention_interval is not None:

1362 + reachable_boundaries = (req.num_prompt_tokens - 1,)

1363 + if req.shared_prefix_boundary:

1364 + reachable_boundaries += (req.shared_prefix_boundary,)

1365 +

1366 for group_config, group_state in zip(

1367 self.config.kv_group_configs, req_status.group_states

1368 ):

@@ -1379,23 +1389,31 @@ class OffloadingConnectorScheduler:

1389 ]

1390 assert len(offload_keys) == len(offload_block_ids)

1391

1392 + # Use reachable_block_mask to filter unreachable chunks

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

1394 + # reachable_block_mask operates in KV-block coordinates,

1395 + # 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,

1403 + reachable_boundaries=reachable_boundaries,

1404 + dcp_world_size=self.config.dcp_world_size,

1405 + )

1406 +

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

1408 zip(offload_keys, offload_block_ids)

1409 ):

1410 if block_id == 0:

1411 continue

1387 - # Skip SWA chunks that can never serve a load hit:

1388 - # within each full-attention alignment segment, only the

1389 - # trailing chunks queried by _sliding_window_lookup are

1390 - # reachable. EAGLE/MTP requires one additional chunk that

1391 - # lookup later drops as its volatile draft tail.

1392 - abs_chunk_idx = start_chunk_idx + key_idx

1393 - if not is_store_reachable_swa_chunk(

1394 - abs_chunk_idx,

1395 - num_chunks,

1396 - group_config.alignment_chunk_count,

1397 - group_config.sliding_window_size_in_chunks,

1398 - group_config.is_eagle_group,

1412 + # A chunk is reachable if any of its constituent

1413 + # blocks is reachable.

1414 + if block_mask is not None and not any(

1415 + block_mask[key_idx * blocks_per_chunk + b]

1416 + for b in range(blocks_per_chunk)

1417 ):

1418 continue

1419 new_offload_keys.append(offload_key)

vllm/v1/core/single_type_kv_cache_manager.py

+13 -2

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

481 use_eagle=self.use_eagle,

482 retention_interval=retention_interval,

483 reachable_boundaries=reachable_boundaries,

484 + dcp_world_size=self.dcp_world_size,

485 )

486 self.block_pool.cache_full_blocks(

487 request=request,

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

505 use_eagle: bool,

506 retention_interval: int | None = None,

507 reachable_boundaries: Sequence[int] = (),

508 + dcp_world_size: int = 1,

509 ) -> list[bool] | None:

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

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

@@ -1028,14 +1030,22 @@ class SlidingWindowManager(SingleTypeKVCacheManager):

1030 use_eagle: bool,

1031 retention_interval: int | None = None,

1032 reachable_boundaries: Sequence[int] = (),

1033 + dcp_world_size: int = 1,

1034 ) -> list[bool] | None:

1035 assert isinstance(kv_cache_spec, SlidingWindowSpec)

1036 if alignment_tokens is None:

1037 # Fast path: when the coordinator imposes no alignment constraint.

1038 return None

1036 - assert alignment_tokens % kv_cache_spec.block_size == 0

1039 + block_size = kv_cache_spec.block_size * dcp_world_size

1040 + if alignment_tokens % block_size != 0:

1041 + # The mask is block-granular, so a sub-block alignment cannot be

1042 + # represented exactly. This happens for hybrid offloading, where

1043 + # ``alignment_tokens`` is the full-attention chunk size and need not

1044 + # be a multiple of this SWA group's (DCP-scaled) block size (e.g.

1045 + # Gemma). Fall back to dense: every block is reachable, which never

1046 + # drops a block that could serve a hit.

1047 + return None

1048

1038 - block_size = kv_cache_spec.block_size

1049 # Contiguous blocks a hit needs at a boundary (incl. the EAGLE peek).

1050 need = cls._contiguous_blocks_for_hit(

1051 window_size=kv_cache_spec.sliding_window,

@@ -1494,6 +1504,7 @@ class MambaManager(SingleTypeKVCacheManager):

1504 use_eagle: bool,

1505 retention_interval: int | None = None,

1506 reachable_boundaries: Sequence[int] = (),

1507 + dcp_world_size: int = 1,

1508 ) -> list[bool] | None:

1509 """Sparse Mamba state-snapshot retention.

1510