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 bc5190db00f5..d89c7a2478a8 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -18,6 +18,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, + is_store_reachable_swa_chunk, ) from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -123,6 +124,55 @@ 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( 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 5e98c1266e20..acc7fe33cbf9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -111,6 +111,26 @@ 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: @@ -974,9 +994,6 @@ def _build_store_jobs( ] assert len(offload_keys) == len(offload_block_ids) - alignment_chunk_count = group_config.alignment_chunk_count - tail = group_config.sliding_window_size_in_chunks - for key_idx, (offload_key, block_id) in enumerate( zip(offload_keys, offload_block_ids) ): @@ -984,15 +1001,18 @@ def _build_store_jobs( continue # Skip SWA chunks that can never serve a load hit: # within each full-attention alignment segment, only the - # trailing `tail` chunks are reachable by - # _sliding_window_lookup. For DeepSeek V4 with 100K - # tokens this reduces SWA stores by ~78%. - if alignment_chunk_count is not None: - assert tail is not None - abs_chunk_idx = start_chunk_idx + key_idx - pos_in_segment = abs_chunk_idx % alignment_chunk_count - if pos_in_segment < alignment_chunk_count - tail: - continue + # 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, + ): + continue new_offload_keys.append(offload_key) if not new_offload_keys: