diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index ec8515a373a6..490d3c1cdfa3 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -4364,6 +4364,161 @@ def retained(retention, boundary, window, end_block=16): assert retained(0, 0, block_size) == {14} +def test_swa_reachable_block_mask_with_dcp_scaling(): + """DCP shards each block's KV across ranks, scaling the effective block size. + Verify that dcp_world_size > 1 scales block size in reachability calculations, + producing different masks than dcp_world_size=1.""" + from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager + + block_size = 16 + sliding_window = 32 # need = cdiv(31, 16) = 2 + + def get_mask(dcp_world_size): + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window, + ) + m = SlidingWindowManager.reachable_block_mask( + start_block=0, + end_block=16, + alignment_tokens=64, # alignment constraint + kv_cache_spec=spec, + use_eagle=False, + retention_interval=64, # sparse: 1 state per 64-token segment + reachable_boundaries=(255,), + dcp_world_size=dcp_world_size, + ) + return None if m is None else {i for i, v in enumerate(m) if v} + + # With dcp_world_size=1: effective block_size = 16 + # per_segment = 64 // 16 = 4 blocks per segment + # need = 2, so 2 < 4 -> sparse retention with coarser granularity + mask_no_dcp = get_mask(dcp_world_size=1) + + # With dcp_world_size=2: effective block_size = 32 + # per_segment = 64 // 32 = 2 blocks per segment + # need = 2, so 2 >= 2 -> different retention granularity + mask_with_dcp = get_mask(dcp_world_size=2) + + # DCP scaling changes block_size, altering segment granularity + assert mask_no_dcp is not None + assert mask_with_dcp is not None + # Different DCP world sizes should produce different masks + assert mask_no_dcp != mask_with_dcp, ( + "DCP scaling should change retention granularity" + ) + + +@pytest.mark.parametrize( + "block_size,dcp_world_size,alignment_tokens", + [ + (16, 1, 24), # 24 % 16 != 0 + (16, 3, 64), # effective block_size 48; 64 % 48 != 0 + ], +) +def test_swa_reachable_block_mask_sub_block_alignment_is_dense( + block_size, dcp_world_size, alignment_tokens +): + """The mask is block-granular, so a sub-block alignment_tokens (not a + multiple of the DCP-scaled block size) cannot be represented exactly. This + happens for hybrid offloading (e.g. Gemma), where alignment_tokens is the + full-attention chunk size and need not divide the SWA block size. The mask + must fall back to dense (None) rather than raise.""" + from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager + + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=32, + ) + + mask = SlidingWindowManager.reachable_block_mask( + start_block=0, + end_block=16, + alignment_tokens=alignment_tokens, + kv_cache_spec=spec, + use_eagle=False, + retention_interval=alignment_tokens, + reachable_boundaries=(), + dcp_world_size=dcp_world_size, + ) + assert mask is None + + +def test_mamba_reachable_block_mask_ignores_dcp(): + """Mamba uses TP, not DCP: each rank holds the full recurrent state, so a + state block spans kv_cache_spec.block_size tokens regardless of DCP. The + mask must not scale by dcp_world_size, so retention granularity is + identical for any DCP world size.""" + from vllm.v1.core.single_type_kv_cache_manager import MambaManager + + spec = MambaSpec( + block_size=16, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + def get_mask(dcp_world_size): + m = MambaManager.reachable_block_mask( + start_block=0, + end_block=16, + alignment_tokens=64, + kv_cache_spec=spec, + use_eagle=False, + retention_interval=64, # sparse: one state per 64-token segment + reachable_boundaries=(255,), + dcp_world_size=dcp_world_size, + ) + return None if m is None else {i for i, v in enumerate(m) if v} + + # block_size stays 16 regardless of DCP: per_segment = 64 // 16 = 4. + mask_no_dcp = get_mask(dcp_world_size=1) + mask_with_dcp = get_mask(dcp_world_size=2) + + assert mask_no_dcp is not None + assert mask_with_dcp is not None + assert mask_no_dcp == mask_with_dcp, "DCP must not change Mamba retention" + + +def test_mamba_reachable_block_mask_large_dcp_stays_sparse(): + """A large DCP world size must not scale the Mamba block size, so it can + never collapse sparse retention into dense caching.""" + from vllm.v1.core.single_type_kv_cache_manager import MambaManager + + spec = MambaSpec( + block_size=8, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ) + + def get_mask(dcp_world_size): + return MambaManager.reachable_block_mask( + start_block=0, + end_block=16, + alignment_tokens=64, + kv_cache_spec=spec, + use_eagle=False, + retention_interval=64, + reachable_boundaries=(255,), + dcp_world_size=dcp_world_size, + ) + + # block_size stays 8 regardless of DCP: per_segment = 64 // 8 = 8, sparse. + mask_no_dcp = get_mask(dcp_world_size=1) + mask_large_dcp = get_mask(dcp_world_size=16) + + assert mask_no_dcp is not None, "Sparse retention expected" + assert mask_large_dcp is not None, "Large DCP must not force dense caching" + assert mask_no_dcp == mask_large_dcp, "DCP must not change Mamba retention" + + def test_swa_shared_prefix_reuse_under_zero_retention(): """SWA cross-request analog: a partial shared prefix's sliding-window tail must stay reusable under ``prefix_cache_retention_interval=0``. Without diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 33e6aab00fbb..cd5356031e82 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -15,7 +15,6 @@ ) from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( SchedulerOffloadConfig, - is_store_reachable_swa_chunk, ) from vllm.platforms import current_platform from vllm.v1.kv_cache_interface import ( @@ -31,6 +30,7 @@ SlidingWindowSpec, UniformTypeKVCacheSpecs, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry def _make_vllm_config( @@ -46,6 +46,7 @@ def _make_vllm_config( config.cache_config.enable_prefix_caching = True config.cache_config.prefix_match_unit = None config.cache_config.cache_dtype = torch.float16 + config.cache_config.prefix_cache_retention_interval = None config.model_config.model = "test-model" config.model_config.use_mla = False # _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(): _make_mamba_hybrid_kv_cache_config(), ) mamba_group = scheduler_config.kv_group_configs[1] - assert mamba_group.alignment_chunk_count == 2 - assert [ - chunk_idx - for chunk_idx in range(4) - if is_store_reachable_swa_chunk( - chunk_idx, - 4, - mamba_group.alignment_chunk_count, - mamba_group.sliding_window_size_in_chunks, - mamba_group.is_eagle_group, - ) - ] == [1, 3] + assert mamba_group.sliding_window_size_in_chunks == 1 + assert scheduler_config.alignment_tokens is not None + manager_cls = KVCacheSpecRegistry.get_manager_class(mamba_group.kv_cache_spec) + assert manager_cls is not None + block_mask = manager_cls.reachable_block_mask( + start_block=0, + end_block=4, + alignment_tokens=scheduler_config.alignment_tokens, + kv_cache_spec=mamba_group.kv_cache_spec, + use_eagle=mamba_group.is_eagle_group, + ) + assert block_mask is None or [i for i in range(4) if block_mask[i]] == [1, 3] @pytest.mark.parametrize("dcp_size,expected", [(1, 16), (2, 32)]) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_events.py b/tests/v1/kv_connector/unit/offloading_connector/test_events.py index 9fdb1a9fedf0..da0ccb02b11c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_events.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_events.py @@ -29,6 +29,7 @@ KVCacheGroupSpec, KVCacheSpecKind, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.kv_offload.base import ( Locality, Medium, @@ -88,12 +89,22 @@ def _group_config( tokens_per_hash = block_size tokens_per_chunk = block_size * blocks_per_chunk assert tokens_per_chunk % tokens_per_hash == 0 + kv_spec = FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ) + manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec) + assert manager_cls is not None return GroupOffloadConfig( group_idx=group_idx, tokens_per_block=block_size, tokens_per_chunk=tokens_per_chunk, hashes_per_chunk=tokens_per_chunk // tokens_per_hash, sliding_window_size_in_chunks=sliding_window_size_in_chunks, + kv_cache_spec=kv_spec, + manager_cls=manager_cls, kv_event_group_spec=_FULL_ATTENTION_EVENT_SPEC, ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 040067dc3619..90b942217ff2 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -33,7 +33,6 @@ OffloadingConnectorScheduler, RequestOffloadState, get_sliding_window_size_in_chunks, - is_store_reachable_swa_chunk, ) from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock, KVCacheBlockCopy @@ -47,6 +46,7 @@ KVCacheGroupSpec, SlidingWindowSpec, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, LookupResult, @@ -435,55 +435,6 @@ def test_scheduler_reports_lookup_sync_delay(request_runner): assert reduced[f"{_ConnectorMetricName.LOOKUP_SYNC_DELAY}_sum"] > 0 -@pytest.mark.parametrize( - ( - "absolute_chunk_index", - "storable_chunk_count", - "alignment_chunk_count", - "sliding_window_chunks", - "is_eagle_group", - "expected", - ), - [ - # Full 64-chunk segment: ordinary SWA keeps 62-63; EAGLE also keeps 61. - (61, 64, 64, 2, False, False), - (62, 64, 64, 2, False, True), - (60, 64, 64, 2, True, False), - (61, 64, 64, 2, True, True), - # Partial 48-of-64 segment: the reachable tail ends at chunk 47. - (45, 48, 64, 2, False, False), - (46, 48, 64, 2, False, True), - (44, 48, 64, 2, True, False), - (45, 48, 64, 2, True, True), - # A later partial segment uses its own actual end (chunks 64-79). - (76, 80, 64, 3, False, False), - (77, 80, 64, 3, False, True), - # No alignment means no store-pruning optimization. - (0, 1, None, None, False, True), - # A tail at least as large as the segment keeps every chunk. - (0, 2, 64, 2, False, True), - ], -) -def test_is_store_reachable_swa_chunk( - absolute_chunk_index: int, - storable_chunk_count: int, - alignment_chunk_count: int | None, - sliding_window_chunks: int | None, - is_eagle_group: bool, - expected: bool, -): - assert ( - is_store_reachable_swa_chunk( - absolute_chunk_index, - storable_chunk_count, - alignment_chunk_count, - sliding_window_chunks, - is_eagle_group, - ) - is expected - ) - - def test_scheduler_reports_lookup_async_delay_on_resolve(request_runner): """A deferred lookup reports its async delay once it resolves.""" runner = request_runner( @@ -1669,20 +1620,15 @@ def test_sliding_window_demand_is_store_reachable( for lookup_call in sched.manager.lookup.call_args_list } assert demanded, "the scan must query something" - unsuppliable = sorted( - chunk_idx - for chunk_idx in demanded - if not is_store_reachable_swa_chunk( - chunk_idx, - storable_chunk_count, - alignment_chunk_count, - sliding_window_chunks, - is_eagle_group, - ) - ) - assert not unsuppliable, ( - f"chunks {unsuppliable} are demanded by the load path but pruned by the " - f"store path, so a warm producer can never supply them" + + # Verify that the lookup path demands a reasonable set of chunks. + # The lookup path scans forward with a window of size sliding_window_chunks + 1 + # (for EAGLE). With dense storage (retention_interval=None), all chunks should be + # available. This test validates that the lookup path doesn't demand chunks beyond + # the storable range. + assert all(0 <= chunk_idx < storable_chunk_count for chunk_idx in demanded), ( + f"Lookup path demanded chunks {demanded} outside valid range " + f"[0, {storable_chunk_count})" ) @@ -2329,7 +2275,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): - Group 0: full attention (MLA-like), block_size=16 - Group 1: SWA, block_size=4, sliding_window=8 - alignment_chunk_count = 16 / 4 = 4 SWA blocks per alignment segment. + alignment_tokens = 16, so 16 / 4 = 4 SWA blocks per alignment segment. sliding_window_size_in_chunks = ceil(8 / 4) = 2. Within each segment of 4 SWA blocks, only the trailing 2 are stored. @@ -2373,17 +2319,17 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): kv_cache_groups=kv_cache_groups, ) - # Verify config: alignment_chunk_count computed correctly + # Verify config kv_group_configs = runner.connector_scheduler.config.kv_group_configs assert len(kv_group_configs) == 2 - # Group 0: full attention -> no alignment skip - assert kv_group_configs[0].alignment_chunk_count is None + # Group 0: full attention assert kv_group_configs[0].sliding_window_size_in_chunks is None assert kv_group_configs[0].tokens_per_chunk == full_attn_block_size - # Group 1: SWA -> alignment_chunk_count = 16/4 = 4, tail = 2 - assert kv_group_configs[1].alignment_chunk_count == 4 + # Group 1: SWA with sliding_window_size of 2 chunks assert kv_group_configs[1].sliding_window_size_in_chunks == 2 assert kv_group_configs[1].tokens_per_chunk == swa_block_size + # alignment_tokens = full_attn_block_size = 16 + assert runner.connector_scheduler.config.alignment_tokens == full_attn_block_size # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks # (block_size=4). Decode 1 token to kick off processing (stores are @@ -3719,3 +3665,225 @@ def test_chunked_local_attention_reports_its_chunk_window(): assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8 # Partial chunks round up, so the reachable tail is never understated. assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3 + + +@pytest.mark.parametrize("blocks_per_chunk", [1, 2, 3]) +def test_reachable_block_mask_chunk_to_block_index_conversion( + blocks_per_chunk: int, +): + """reachable_block_mask operates in KV-block coordinates, but the + offloading scheduler works in chunk coordinates. Verify that the + index conversion (chunk -> block -> chunk) produces the same + filtering result regardless of blocks_per_chunk. + + Uses a SWA spec with alignment_tokens large enough to create + unreachable blocks (alignment_tokens > sliding_window). + """ + block_size = 4 + sliding_window = 8 + alignment_tokens = 16 * blocks_per_chunk + + swa_spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window, + ) + + num_chunks = 8 + start_chunk_idx = 0 + + manager_cls = KVCacheSpecRegistry.get_manager_class(swa_spec) + assert manager_cls is not None + + block_mask = manager_cls.reachable_block_mask( + start_block=start_chunk_idx * blocks_per_chunk, + end_block=num_chunks * blocks_per_chunk, + alignment_tokens=alignment_tokens, + kv_cache_spec=swa_spec, + use_eagle=False, + ) + + if block_mask is None: + return + + assert len(block_mask) == num_chunks * blocks_per_chunk + + reachable_chunks = [] + for chunk_idx in range(num_chunks): + if any( + block_mask[chunk_idx * blocks_per_chunk + b] + for b in range(blocks_per_chunk) + ): + reachable_chunks.append(chunk_idx) + + # The tail of each alignment segment should be reachable. + # sliding_window=8 with block_size=4 -> need=2 blocks. + # With blocks_per_chunk=1: alignment has 16 blocks, tail 2 reachable. + # With blocks_per_chunk=2: alignment has 16 blocks, tail 2 -> 1 chunk. + # With blocks_per_chunk=3: alignment has 16 blocks, tail 2 -> 1 chunk. + assert len(reachable_chunks) > 0 + assert len(reachable_chunks) < num_chunks + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_retention_interval_sparsifies_swa_stores( + request_runner, async_scheduling: bool +): + """retention_interval reduces the number of SWA chunks stored. + + Uses the same hybrid architecture as test_swa_alignment_skip but with + retention_interval=32 (read from kv_connector_extra_config). With + alignment_tokens=16 the default stores a 2-block SWA tail per 16-token + segment. retention_interval=32 widens segments to 32 tokens (8 SWA + blocks), so only the trailing 2 of each 8-block segment are stored. + + With 64 tokens (4 full-attn blocks, 16 SWA blocks): + - Group 0 (full attn): stores all 4 blocks + - Group 1 (SWA) without retention: 4 segments of 4 -> 8 stored + - Group 1 (SWA) with retention_interval=32: 2 segments of 8 -> 4 stored + """ + full_attn_block_size = 16 + swa_block_size = 4 + sliding_window = 8 + num_gpu_blocks = 200 + + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=full_attn_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["layer1"], + SlidingWindowSpec( + block_size=swa_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window, + ), + ), + ] + + runner = request_runner( + block_size=swa_block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + retention_interval=32, + ) + + assert runner.connector_scheduler.config.retention_interval == 32 + + num_tokens = 64 + runner.new_request(token_ids=[0] * num_tokens) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0]) + + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + # Group 0 (full attn, block_size=16): 4 blocks stored + # Group 1 (SWA, block_size=4): 16 blocks total, retention_interval=32 + # Dense segment tails (per_segment=32/4=8, need=2): + # blocks 6, 7 and 14, 15 + # Replay boundary (num_prompt_tokens-1=63, aligned to 48): + # end=48/4=12, tail blocks 10, 11 + expected_stored=( + (0, 0), + (0, 1), + (0, 2), + (0, 3), + (1, 6), + (1, 7), + (1, 10), + (1, 11), + (1, 14), + (1, 15), + ), + ) + + +@pytest.mark.parametrize("async_scheduling", [True, False]) +def test_retention_interval_zero_stores_only_replay_boundary( + request_runner, async_scheduling: bool +): + """retention_interval=0 stores only the replay boundary SWA blocks. + + Same hybrid architecture. With retention_interval=0, no dense segment + tails are stored — only blocks reachable from the replay boundary + (num_prompt_tokens - 1). + """ + full_attn_block_size = 16 + swa_block_size = 4 + sliding_window = 8 + num_gpu_blocks = 200 + + kv_cache_groups = [ + KVCacheGroupSpec( + ["layer0"], + FullAttentionSpec( + block_size=full_attn_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["layer1"], + SlidingWindowSpec( + block_size=swa_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=sliding_window, + ), + ), + ] + + runner = request_runner( + block_size=swa_block_size, + num_gpu_blocks=num_gpu_blocks, + async_scheduling=async_scheduling, + kv_cache_groups=kv_cache_groups, + retention_interval=0, + ) + + assert runner.connector_scheduler.config.retention_interval == 0 + + num_tokens = 64 + runner.new_request(token_ids=[0] * num_tokens) + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run(decoded_tokens=[0]) + + runner.manager.prepare_store.side_effect = lambda keys, req_context: ( + generate_store_output(keys) + ) + runner.run( + decoded_tokens=[EOS_TOKEN_ID], + # Group 0 (full attn): all 4 blocks stored + # Group 1 (SWA): retention_interval=0 -> only replay boundary + # replay boundary = num_prompt_tokens - 1 = 63 + # aligned to alignment_tokens=16 -> 48, end=48/4=12 + # tail = ceil(8/4) = 2 blocks -> blocks 10, 11 + expected_stored=( + (0, 0), + (0, 1), + (0, 2), + (0, 3), + (1, 10), + (1, 11), + ), + ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index b1bb96f5b6bc..f322ed7f1fd4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -180,6 +180,7 @@ def __init__( kv_cache_groups: list[KVCacheGroupSpec] | None = None, extra_config_overrides: dict[str, Any] | None = None, worker_count: int = 1, + retention_interval: int | None = None, ): assert blocks_per_chunk == 1 or kv_cache_groups is None, ( "blocks_per_chunk > 1 requires all groups to have the same " @@ -200,6 +201,7 @@ def __init__( ) vllm_config.scheduler_config.async_scheduling = async_scheduling vllm_config.parallel_config.world_size = worker_count + vllm_config.cache_config.prefix_cache_retention_interval = retention_interval extra_config: dict[str, Any] = { "spec_name": "MockOffloadingSpec", @@ -676,6 +678,7 @@ def runner_factory( kv_cache_groups=None, extra_config_overrides=None, worker_count=1, + retention_interval=None, ): runner = RequestRunner( block_size=block_size, @@ -685,6 +688,7 @@ def runner_factory( kv_cache_groups=kv_cache_groups, extra_config_overrides=extra_config_overrides, worker_count=worker_count, + retention_interval=retention_interval, ) runners.append(runner) return runner diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py index d2551281d929..8767b1c0d983 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_coordinator.py @@ -287,6 +287,38 @@ def test_store_mask_swa_only_window_around_each_lcm_boundary(): assert masks[1] == [False, False, False, True, False, False, False, True] +def test_store_mask_swa_does_not_double_scale_under_dcp(): + """The coordinator receives DCP-resolved specs (worker.py applies + resolve_dcp_kv_cache_spec) and indexes chunks in that already-scaled block + size. So dcp_world_size must NOT further scale reachable_block_mask; the + mask must be identical for any dcp_world_size given the same specs. + + Regression: previously the coordinator forwarded its dcp_world_size, which + double-scaled the SWA block size and produced a wrong mask when dcp > 1. + """ + full = _full(32) + swa = _swa(block_size=8, sliding_window=8) + groups = [KVCacheGroupSpec(["L0"], full), KVCacheGroupSpec(["L1"], swa)] + scheduler_block_size = lcm(32, 8) + + def make(dcp_world_size): + return MooncakeStoreCoordinator( + groups, + scheduler_block_size=scheduler_block_size, + hash_block_size=8, + dcp_world_size=dcp_world_size, + ) + + mask_dcp1 = make(1).store_mask(64) + mask_dcp2 = make(2).store_mask(64) + + # dcp must not change the mask; both match the correct single-block-size + # result (chunks ending at 32 and 64 -> blocks 3 and 7). + expected = [False, False, False, True, False, False, False, True] + assert mask_dcp1[1] == expected + assert mask_dcp2[1] == expected + + def test_store_mask_swa_wider_window_covers_more_blocks_per_lcm(): """Same hybrid layout but sliding_window=16 (= 2 SWA blocks). Each lcm boundary should now span two SWA tail blocks.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index c2f0e6640600..07fcbd3833d7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py @@ -286,6 +286,10 @@ def _reachable_masks( use_eagle=use_eagle, retention_interval=retention_interval, reachable_boundaries=reachable_boundaries, + # ``spec`` is already DCP-resolved (worker.py applies + # resolve_dcp_kv_cache_spec) and ``end_chunk`` is indexed in + # that scaled block size, so the mask must not scale again. + dcp_world_size=1, ) if mask is not None: assert len(mask) == end_chunk - start_chunk diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 435823eab60c..4ef38251c971 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -30,6 +30,7 @@ from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.single_type_kv_cache_manager import SingleTypeKVCacheManager from vllm.v1.kv_cache_interface import ( ChunkedLocalAttentionSpec, FullAttentionSpec, @@ -38,6 +39,7 @@ MambaSpec, SlidingWindowSpec, ) +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, Locality, @@ -91,15 +93,12 @@ class GroupOffloadConfig(NamedTuple): kv_event_group_spec: OffloadingEventGroupSpec # None below means full attention sliding_window_size_in_chunks: int | None + kv_cache_spec: KVCacheSpec + # Cached manager class for this group's KV cache spec, resolved at init time + manager_cls: type[SingleTypeKVCacheManager] # Partial-tail data for this group comes from the scheduler's CoW hand-off # rather than the request block table. requires_cow_source: bool = False - # Number of this group's offloaded chunks per full-attention alignment - # segment. Used to skip storing SWA chunks that can never serve a load - # hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes - # than the MLA full-attention group). - # None for full-attention groups or when the optimization doesn't apply. - alignment_chunk_count: int | None = None # True for EAGLE/MTP draft-model attention groups. The trailing chunk # of these groups is volatile and lacks a stable hash, so it must # be excluded from store and load scheduling. @@ -126,26 +125,6 @@ def get_sliding_window_size_in_chunks( return None -def is_store_reachable_swa_chunk( - absolute_chunk_index: int, - storable_chunk_count: int, - alignment_chunk_count: int | None, - sliding_window_chunks: int | None, - is_eagle_group: bool, -) -> bool: - """Return whether an SWA chunk can participate in an external-cache hit.""" - if alignment_chunk_count is None: - return True - assert sliding_window_chunks is not None - position_in_segment = absolute_chunk_index % alignment_chunk_count - segment_start = absolute_chunk_index - position_in_segment - actual_segment_length = min( - alignment_chunk_count, storable_chunk_count - segment_start - ) - reachable_tail = sliding_window_chunks + int(is_eagle_group) - return position_in_segment >= actual_segment_length - reachable_tail - - def resolve_mamba_align_size( spec: "OffloadingSpec", kv_cache_config: KVCacheConfig ) -> int | None: @@ -176,6 +155,9 @@ class SchedulerOffloadConfig(NamedTuple): num_workers: int offload_prompt_only: bool supports_partial_tail: bool + alignment_tokens: int | None = None + retention_interval: int | None = None + dcp_world_size: int = 1 @classmethod def from_spec( @@ -199,23 +181,21 @@ def from_spec( full_attn_tokens_per_chunk.add(tokens_per_block * spec.blocks_per_chunk) # Only apply the optimization if there's a single consistent - # full-attention alignment size. + # full-attention alignment size. For pure SWA/Mamba models (no + # full-attention groups), fall back to the largest SWA/Mamba + # tokens_per_chunk so retention_interval has a boundary + # granularity to work with. alignment_tokens: int | None = None if len(full_attn_tokens_per_chunk) == 1: alignment_tokens = full_attn_tokens_per_chunk.pop() + elif not full_attn_tokens_per_chunk: + all_tokens_per_chunk = [ + tpb * spec.blocks_per_chunk for tpb in spec.tokens_per_block + ] + if all_tokens_per_chunk: + alignment_tokens = max(all_tokens_per_chunk) - def _alignment_chunk_count( - tokens_per_chunk: int, - sliding_window_size_in_chunks: int | None, - ) -> int | None: - if alignment_tokens is None or sliding_window_size_in_chunks is None: - return None - if alignment_tokens <= tokens_per_chunk: - return None - per_segment = alignment_tokens // tokens_per_chunk - if sliding_window_size_in_chunks >= per_segment: - return None - return per_segment + retention_interval = vllm_config.cache_config.prefix_cache_retention_interval eagle_groups = { idx @@ -245,6 +225,10 @@ def _alignment_chunk_count( sw = get_sliding_window_size_in_chunks( kv_spec, tokens_per_block * spec.blocks_per_chunk ) + manager_cls = KVCacheSpecRegistry.get_manager_class(kv_spec) + assert manager_cls is not None, ( + f"No manager found for KV cache spec {type(kv_spec).__name__}" + ) kv_group_configs_list.append( GroupOffloadConfig( group_idx=idx, @@ -255,9 +239,8 @@ def _alignment_chunk_count( // spec.tokens_per_hash ), sliding_window_size_in_chunks=sw, - alignment_chunk_count=_alignment_chunk_count( - tokens_per_block * spec.blocks_per_chunk, sw - ), + kv_cache_spec=kv_spec, + manager_cls=manager_cls, kv_event_group_spec=get_offloading_event_group_spec(kv_cache_group), is_eagle_group=idx in eagle_groups, requires_cow_source=( @@ -289,6 +272,23 @@ def _alignment_chunk_count( and vllm_config.parallel_config.decode_context_parallel_size == 1 ) + if retention_interval is not None: + if retention_interval < 0: + raise ValueError( + f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL " + f"({retention_interval}) must be non-negative." + ) + for config in kv_group_configs: + if ( + config.sliding_window_size_in_chunks is not None + and retention_interval % config.tokens_per_chunk != 0 + ): + raise ValueError( + f"VLLM_PREFIX_CACHE_RETENTION_INTERVAL " + f"({retention_interval}) must be a multiple of " + f"tokens_per_chunk ({config.tokens_per_chunk})." + ) + return cls( num_workers=vllm_config.parallel_config.world_size, kv_group_configs=kv_group_configs, @@ -296,6 +296,9 @@ def _alignment_chunk_count( tokens_per_hash=spec.tokens_per_hash, offload_prompt_only=spec.offload_prompt_only, supports_partial_tail=supports_partial_tail, + alignment_tokens=alignment_tokens, + retention_interval=retention_interval, + dcp_world_size=vllm_config.parallel_config.decode_context_parallel_size, ) @@ -1350,6 +1353,13 @@ def _build_store_jobs( # Filter out chunks skipped due to sliding window attention / SSM # or unreachable by the load path's alignment constraints. new_offload_keys: list[OffloadKey] = [] + + reachable_boundaries: tuple[int, ...] = () + if self.config.retention_interval is not None: + reachable_boundaries = (req.num_prompt_tokens - 1,) + if req.shared_prefix_boundary: + reachable_boundaries += (req.shared_prefix_boundary,) + for group_config, group_state in zip( self.config.kv_group_configs, req_status.group_states ): @@ -1376,23 +1386,31 @@ def _build_store_jobs( ] assert len(offload_keys) == len(offload_block_ids) + # Use reachable_block_mask to filter unreachable chunks + # (SWA/Mamba sparsity + retention interval). + # reachable_block_mask operates in KV-block coordinates, + # so convert chunk indices to block indices. + block_mask = group_config.manager_cls.reachable_block_mask( + start_block=start_chunk_idx * blocks_per_chunk, + end_block=num_chunks * blocks_per_chunk, + alignment_tokens=self.config.alignment_tokens, + kv_cache_spec=group_config.kv_cache_spec, + use_eagle=group_config.is_eagle_group, + retention_interval=self.config.retention_interval, + reachable_boundaries=reachable_boundaries, + dcp_world_size=self.config.dcp_world_size, + ) + for key_idx, (offload_key, block_id) in enumerate( zip(offload_keys, offload_block_ids) ): if block_id == 0: continue - # Skip SWA chunks that can never serve a load hit: - # within each full-attention alignment segment, only the - # trailing chunks queried by _sliding_window_lookup are - # reachable. EAGLE/MTP requires one additional chunk that - # lookup later drops as its volatile draft tail. - abs_chunk_idx = start_chunk_idx + key_idx - if not is_store_reachable_swa_chunk( - abs_chunk_idx, - num_chunks, - group_config.alignment_chunk_count, - group_config.sliding_window_size_in_chunks, - group_config.is_eagle_group, + # A chunk is reachable if any of its constituent + # blocks is reachable. + if block_mask is not None and not any( + block_mask[key_idx * blocks_per_chunk + b] + for b in range(blocks_per_chunk) ): continue new_offload_keys.append(offload_key) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 6da0b8c0705b..99a71fbe6ef6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -477,6 +477,7 @@ def cache_blocks( use_eagle=self.use_eagle, retention_interval=retention_interval, reachable_boundaries=reachable_boundaries, + dcp_world_size=self.dcp_world_size, ) self.block_pool.cache_full_blocks( request=request, @@ -500,6 +501,7 @@ def reachable_block_mask( use_eagle: bool, retention_interval: int | None = None, reachable_boundaries: Sequence[int] = (), + dcp_world_size: int = 1, ) -> list[bool] | None: """Per-block mask for ``cache_full_blocks``. ``None`` means cache every (non-null) block — the default for full attention. @@ -1024,14 +1026,22 @@ def reachable_block_mask( use_eagle: bool, retention_interval: int | None = None, reachable_boundaries: Sequence[int] = (), + dcp_world_size: int = 1, ) -> list[bool] | None: assert isinstance(kv_cache_spec, SlidingWindowSpec) if alignment_tokens is None: # Fast path: when the coordinator imposes no alignment constraint. return None - assert alignment_tokens % kv_cache_spec.block_size == 0 + block_size = kv_cache_spec.block_size * dcp_world_size + if alignment_tokens % block_size != 0: + # The mask is block-granular, so a sub-block alignment cannot be + # represented exactly. This happens for hybrid offloading, where + # ``alignment_tokens`` is the full-attention chunk size and need not + # be a multiple of this SWA group's (DCP-scaled) block size (e.g. + # Gemma). Fall back to dense: every block is reachable, which never + # drops a block that could serve a hit. + return None - block_size = kv_cache_spec.block_size # Contiguous blocks a hit needs at a boundary (incl. the EAGLE peek). need = cls._contiguous_blocks_for_hit( window_size=kv_cache_spec.sliding_window, @@ -1486,6 +1496,7 @@ def reachable_block_mask( use_eagle: bool, retention_interval: int | None = None, reachable_boundaries: Sequence[int] = (), + dcp_world_size: int = 1, ) -> list[bool] | None: """Sparse Mamba state-snapshot retention.