From 481839ad9e5ebf87aecb54fa5c9d986bd5ea4b81 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Tue, 1 Sep 2026 19:12:36 +0800 Subject: [PATCH] [Feature][Spec] Support disabling trailing prefix-cache block dropping (#53388) Signed-off-by: Ziming Huang Signed-off-by: Ziming Huang <48115868+ZeldaHuang@users.noreply.github.com> Co-authored-by: OpenAI Codex --- tests/test_config.py | 17 ++++++++ .../test_partial_prefix_cache_hits.py | 6 +-- tests/v1/core/test_kv_cache_utils.py | 1 + tests/v1/core/test_mamba_align_chunk_split.py | 15 ++++++- tests/v1/core/test_prefix_caching.py | 2 +- .../core/test_single_type_kv_cache_manager.py | 41 +++++++++++++++++++ vllm/config/speculative.py | 9 ++++ .../kv_connector/v1/mooncake/store/worker.py | 9 ++-- .../kv_connector/v1/offloading/scheduler.py | 6 +-- vllm/v1/core/kv_cache_utils.py | 2 +- vllm/v1/core/sched/scheduler.py | 12 +++++- vllm/v1/core/single_type_kv_cache_manager.py | 7 ++-- vllm/v1/simple_kv_offload/manager.py | 6 ++- 13 files changed, 112 insertions(+), 21 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 1a36b0ec92..aa0c838e3f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -2152,6 +2152,23 @@ def test_draft_sample_method_probabilistic_is_accepted(): assert speculative_config.draft_sample_method == "probabilistic" +@pytest.mark.parametrize("disable_eagle_block_drop", [False, True]) +def test_eagle_block_drop_can_be_disabled_without_disabling_eagle( + disable_eagle_block_drop: bool, +): + # Start from an ngram config to avoid loading model metadata: these predicates + # depend only on the speculative method and the new switch. + speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=3, + disable_eagle_block_drop=disable_eagle_block_drop, + ) + speculative_config.method = "eagle3" + + assert speculative_config.use_eagle() + assert speculative_config.use_eagle_block_drop() is not disable_eagle_block_drop + + def test_draft_sample_method_gumbel_is_rejected(): with pytest.raises(ValidationError): SpeculativeConfig( diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 1d5a139eb8..2dcc0e1184 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -170,7 +170,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): cache_config=SimpleNamespace(block_size=block_size), max_num_scheduled_tokens=8192, scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), - use_eagle=False, + use_eagle_block_drop=False, hash_block_size=hash_block_size, dcp_world_size=dcp_world_size, scheduler_block_size=scheduler_block_size, @@ -218,7 +218,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget(): cache_config=SimpleNamespace(block_size=block_size), max_num_scheduled_tokens=token_budget, scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), - use_eagle=False, + use_eagle_block_drop=False, hash_block_size=32, mamba_partial_cache_hit=False, mamba_has_prefill_checkpoint_blocks=False, @@ -257,7 +257,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold(): scheduler_config=SimpleNamespace( long_prefill_token_threshold=long_prefill_threshold ), - use_eagle=False, + use_eagle_block_drop=False, hash_block_size=32, mamba_partial_cache_hit=False, mamba_has_prefill_checkpoint_blocks=False, diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 69ccfafc68..c5159e0d16 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3296,6 +3296,7 @@ def _spec_decode_grouping_config(method="dspark", model_type=None): speculative_config=SimpleNamespace( method=method, use_eagle=lambda: True, + use_eagle_block_drop=lambda: True, ), ) diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py index bc809fc3b5..2d7d77f1ce 100644 --- a/tests/v1/core/test_mamba_align_chunk_split.py +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -80,13 +80,16 @@ def _split( request: Request, num_new_tokens: int, use_eagle: bool = True, + use_eagle_block_drop: bool | None = None, partial_hit: bool = False, num_prefill_checkpoint_blocks: int = 0, ) -> int: """Call the real `Scheduler._mamba_block_aligned_split` on a stub self.""" + if use_eagle_block_drop is None: + use_eagle_block_drop = use_eagle stub = SimpleNamespace( cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE), - use_eagle=use_eagle, + use_eagle_block_drop=use_eagle_block_drop, max_num_scheduled_tokens=16384, scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), # `prefix_match_unit` finer than the block size (#46384). @@ -128,6 +131,16 @@ def test_internal_checkpoint_split( assert all(not block.is_null for block in blocks) # checkpoint + running state +def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None: + (request,) = create_requests(1, num_tokens=3602, block_size=ATTN_BLOCK_SIZE) + + with_drop = _split(request, request.num_tokens, use_eagle_block_drop=True) + without_drop = _split(request, request.num_tokens, use_eagle_block_drop=False) + + assert with_drop == MAMBA_BLOCK_SIZE + assert without_drop == 2 * MAMBA_BLOCK_SIZE + + def _run_chunked_prefill( manager: KVCacheManager, request: Request, budgets: list[int] ) -> dict[int, int]: diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index c42ec12738..ec8515a373 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1163,7 +1163,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): cache_config=SimpleNamespace(block_size=block_size), max_num_scheduled_tokens=3 * block_size, scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), - use_eagle=False, + use_eagle_block_drop=False, hash_block_size=block_size, mamba_partial_cache_hit=False, mamba_has_prefill_checkpoint_blocks=False, diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 57885c9cb9..4be95aba52 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -399,6 +399,47 @@ def test_sliding_window_possible_cached_prefix(): ) +def test_sliding_window_cache_hit_with_finer_hash_alignment(): + """Sliding-window lookup uses full blocks with finer hybrid-cache hashes.""" + hash_block_size = 2 + block_size = 4 + sliding_window_spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=8, + ) + block_pool = BlockPool( + num_gpu_blocks=100, + enable_caching=True, + hash_block_size=hash_block_size, + ) + manager = get_sliding_window_manager(sliding_window_spec, block_pool) + block_hashes = [BlockHash(str(i).encode()) for i in range(4)] + + # Each physical block uses the last chained hash in its block-size view. + for block_hash, block in zip( + (block_hashes[1], block_hashes[3]), block_pool.blocks[10:12] + ): + block_pool.cached_block_hash_to_block.insert( + make_block_hash_with_group_id(block_hash, 0), block + ) + + computed_blocks, hit_length = manager.find_longest_cache_hit( + block_hashes=block_hashes, + max_length=8, + kv_cache_group_ids=[0], + block_pool=block_pool, + kv_cache_spec=sliding_window_spec, + drop_eagle_block=False, + alignment_tokens=hash_block_size, + ) + + assert hit_length == 8 + assert computed_blocks[0] == block_pool.blocks[10:12] + + def test_chunked_local_attention_remove_skipped_blocks(): attention_spec = ChunkedLocalAttentionSpec( block_size=2, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index fde84a2622..92cf7c2ec3 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -436,6 +436,11 @@ class SpeculativeConfig: speculative input batches can contain sequences of different lengths, which may only be supported by certain attention backends. This currently only affects the EAGLE method of speculation.""" + disable_eagle_block_drop: bool = False + """Disable dropping the trailing prefix-cache block for EAGLE-like + speculative methods. This is an experimental option for measuring the + acceptance-rate impact of reusing that block. It does not disable the + speculative drafter itself.""" use_local_argmax_reduction: bool = False """Use vocab-parallel local argmax instead of all-gathering full logits for draft token generation. Reduces communication from O(vocab_size) to @@ -1846,6 +1851,10 @@ class SpeculativeConfig: # TODO(ben): Refactor this so the naming is clearer return self.method in ("eagle", "eagle3", "mtp", "dflash", "dspark") + def use_eagle_block_drop(self) -> bool: + """Whether volatile trailing cache blocks should be discarded.""" + return self.use_eagle() and not self.disable_eagle_block_drop + def use_dflash(self) -> bool: return self.method == "dflash" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index 2331d4da63..95c892c06f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py @@ -1593,16 +1593,17 @@ class MooncakeStoreWorker: for group in kv_cache_config.transfer_groups ] spec_cfg = getattr(vllm_config, "speculative_config", None) - use_eagle = bool( - spec_cfg.use_eagle() - if spec_cfg is not None and callable(getattr(spec_cfg, "use_eagle", None)) + use_eagle_block_drop = bool( + spec_cfg.use_eagle_block_drop() + if spec_cfg is not None + and callable(getattr(spec_cfg, "use_eagle_block_drop", None)) else False ) self.coord = MooncakeStoreCoordinator( self._kv_cache_groups, scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, - use_eagle=use_eagle, + use_eagle=use_eagle_block_drop, retention_interval=kv_cache_config.prefix_cache_retention_interval, dcp_world_size=self.dcp_size, ) 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 330f8270b3..402bd55dc5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -223,11 +223,11 @@ class SchedulerOffloadConfig(NamedTuple): if g.is_eagle_group } - use_eagle = ( + use_eagle_block_drop = ( vllm_config.speculative_config is not None - and vllm_config.speculative_config.use_eagle() + and vllm_config.speculative_config.use_eagle_block_drop() ) - if use_eagle and not eagle_groups: + if use_eagle_block_drop and not eagle_groups: eagle_groups = set(range(len(kv_cache_config.kv_cache_groups))) if eagle_groups: diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 0b57ce7e96..1bbe002c97 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1841,7 +1841,7 @@ def _annotate_eagle_groups( use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group. """ spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): + if spec_config is None or not spec_config.use_eagle_block_drop(): return for group in kv_cache_groups: diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 83ed59799e..b0fd3d25bf 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -255,6 +255,7 @@ class Scheduler(SchedulerInterface): ) speculative_config = vllm_config.speculative_config self.use_eagle = False + self.use_eagle_block_drop = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = vllm_config.num_lookahead_tokens # Positions past the computed tokens that the drafter reads mid-prefill. @@ -279,6 +280,13 @@ class Scheduler(SchedulerInterface): if speculative_config.use_multi_module_mtp() else 1 ) + self.use_eagle_block_drop = speculative_config.use_eagle_block_drop() + if self.use_eagle and not self.use_eagle_block_drop: + logger.warning( + "EAGLE trailing prefix-cache block dropping is disabled. " + "This is experimental and may affect speculative-token " + "acceptance rates." + ) # Create the KV cache manager. if hash_block_size is None: @@ -289,7 +297,7 @@ class Scheduler(SchedulerInterface): max_model_len=self.max_model_len, max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, - use_eagle=self.use_eagle, + use_eagle=self.use_eagle_block_drop, num_prefill_lookahead=self.num_prefill_lookahead, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, @@ -414,7 +422,7 @@ class Scheduler(SchedulerInterface): # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle: + if self.use_eagle_block_drop: last_cache_position = max(last_cache_position - block_size, 0) end = start + num_new_tokens diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 443017012f..2091be31c6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -925,10 +925,9 @@ class SlidingWindowManager(SingleTypeKVCacheManager): ) assert dcp_world_size == 1, "DCP not support sliding window attn now." assert pcp_world_size == 1, "PCP not support sliding window attn now." - # Fine-grained partial hits are not supported for sliding window now - assert alignment_tokens % kv_cache_spec.block_size == 0, ( - "SlidingWindowManager does not support fine-grained (partial) cache hits" - ) + # Sliding-window cache hits must stay at the group's physical block + # granularity. resolve_block_hashes() converts finer-grained hashes to + # that view when the hybrid-cache alignment is smaller than block_size. block_hashes = resolve_block_hashes( block_hashes, block_pool.hash_block_size, diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 0a78c5e9c6..cc3a1fcc36 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -125,12 +125,14 @@ class SimpleCPUOffloadScheduler: ) spec_config = vllm_config.speculative_config - use_eagle = spec_config is not None and spec_config.use_eagle() + use_eagle_block_drop = ( + spec_config is not None and spec_config.use_eagle_block_drop() + ) self.cpu_coordinator: KVCacheCoordinator = get_kv_cache_coordinator( kv_cache_config=self.cpu_kv_cache_config, max_model_len=vllm_config.model_config.max_model_len, max_in_flight_tokens=vllm_config.max_in_flight_tokens, - use_eagle=use_eagle, + use_eagle=use_eagle_block_drop, enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=dcp_world_size,