[Core][Spec Decode] Lookahead block hashes for EAGLE prefix caching (#61589)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Dakai An <dakaian108@gmail.com>
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Dakai An <dakaian108@gmail.com>
Co-authored-by: Dakai An <77474977+andakai@users.noreply.github.com>
Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

examples/features/kv_events/kv_events_subscriber.py

@@ -46,6 +46,7 @@ class BlockStored(KVCacheEvent):
     kv_cache_spec_sliding_window: int | None = None
     locality: str | None = None
     session_id: str | None = None
+    lookahead_token_ids: list[list[int]] | None = None
 
 
 class BlockRemoved(KVCacheEvent):

tests/distributed/test_kv_cache_events.py

@@ -190,15 +190,15 @@ def test_block_stored_locality_is_wire_compatible():
         kv_cache_spec_sliding_window=128,
     )
     legacy_payload = msgspec.msgpack.encode(legacy)
-    assert (
+    # More fields can widen the MessagePack map header, not change its contents.
+    assert msgspec.msgpack.decode(
         msgspec.msgpack.encode(
             _make_block_stored(
                 group_idx=2,
                 kv_cache_spec_sliding_window=128,
             )
         )
-        == legacy_payload
-    )
+    ) == msgspec.msgpack.decode(legacy_payload)
     assert msgspec.msgpack.decode(legacy_payload, type=BlockStored).locality is None
     new_payload = msgspec.msgpack.encode(_make_block_stored(locality="LOCAL"))
     assert msgspec.msgpack.decode(new_payload)["locality"] == "LOCAL"
@@ -229,7 +229,7 @@ def test_block_stored_session_id_is_wire_compatible():
         locality="LOCAL",
     )
     pre_session_payload = msgspec.msgpack.encode(pre_session)
-    assert (
+    assert msgspec.msgpack.decode(
         msgspec.msgpack.encode(
             _make_block_stored(
                 group_idx=2,
@@ -237,8 +237,7 @@ def test_block_stored_session_id_is_wire_compatible():
                 locality="LOCAL",
             )
         )
-        == pre_session_payload
-    )
+    ) == msgspec.msgpack.decode(pre_session_payload)
     assert (
         msgspec.msgpack.decode(pre_session_payload, type=BlockStored).session_id is None
     )

tests/models/kimi_k3/test_kda_metadata.py

@@ -13,7 +13,7 @@ from tests.v1.attention.utils import (
     create_common_attn_metadata,
     create_vllm_config,
 )
-from vllm.config import SpeculativeConfig
+from vllm.config import SpeculativeConfig, VllmConfig
 from vllm.config.compilation import CUDAGraphMode
 from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
 from vllm.models.kimi_k3.nvidia.kda_metadata import (
@@ -223,6 +223,7 @@ def test_kda_recoverssm_startup_metadata_flow_without_model(monkeypatch):
         scheduler_config=SimpleNamespace(max_num_seqs=4),
         additional_config={},
         num_speculative_tokens=2,
+        num_hash_lookahead_tokens=0,
         use_v2_model_runner=True,
     )
     builder = KimiK3KDAMetadataBuilder(
@@ -363,22 +364,37 @@ def test_internal_checkpoint_metadata_targets_last_aligned_boundary(
 
 
 @pytest.mark.parametrize(
-    ("seq_len", "disable_eagle_block_drop", "prefix_match_unit", "expected_offset"),
+    (
+        "seq_len",
+        "disable_eagle_block_drop",
+        "num_hash_lookahead_tokens",
+        "prefix_match_unit",
+        "expected_offset",
+    ),
     [
-        (100, False, 16, 80),
-        (100, True, 16, 96),
-        (100, False, 8, None),
-        (96, False, 16, 80),
-        (96, True, 16, 80),
+        (100, False, 0, 16, 80),
+        (100, True, 0, 16, 96),
+        (100, False, 1, 16, 96),
+        (100, False, 0, 8, None),
+        (96, False, 0, 16, 80),
+        (96, True, 0, 16, 80),
     ],
 )
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_spec_internal_checkpoint_metadata_targets_replay_boundary(
+    monkeypatch: pytest.MonkeyPatch,
     seq_len: int,
     disable_eagle_block_drop: bool,
+    num_hash_lookahead_tokens: int,
     prefix_match_unit: int,
     expected_offset: int | None,
 ) -> None:
+    """Lookahead hashes need no back-off from the replay boundary."""
+    monkeypatch.setattr(
+        VllmConfig,
+        "num_hash_lookahead_tokens",
+        property(lambda _: num_hash_lookahead_tokens),
+    )
     device = torch.device("cuda")
     batch = BatchSpec(seq_lens=[seq_len], query_lens=[seq_len])
     common_attn_metadata = create_common_attn_metadata(

tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py

@@ -71,6 +71,7 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True):
         # The EAGLE adjustments key on the block-drop bit, not plain use_eagle:
         # they exist only to compensate for the drop.
         use_eagle_block_drop=block_drop,
+        num_hash_lookahead_tokens=0,
         hash_block_size=hash_block_size,
         mamba_has_prefill_checkpoint_blocks=False,  # forced False under eagle
         mamba_partial_cache_hit=partial_hit,

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -103,6 +103,7 @@ def make_full_mamba_manager(
     use_eagle: bool = False,
     num_speculative_blocks: int = 0,
     num_prefill_checkpoint_blocks: int = 0,
+    num_hash_lookahead_tokens: int = 0,
 ):
     mamba_group = KVCacheGroupSpec(
         ["mamba"],
@@ -146,6 +147,8 @@ def make_full_mamba_manager(
         scheduler_block_size=scheduler_block_size,
         hash_block_size=hash_block_size,
         use_eagle=use_eagle,
+        num_prefill_lookahead=num_hash_lookahead_tokens,
+        num_hash_lookahead_tokens=num_hash_lookahead_tokens,
     )
 
 
@@ -189,6 +192,7 @@ def test_dcp_fine_hit_retention_uses_hash_alignment(
         max_num_scheduled_tokens=8192,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=use_eagle,
+        num_hash_lookahead_tokens=0,
         mamba_partial_cache_hit=True,
         mamba_shared_prefix_checkpoint=False,
         mamba_has_prefill_checkpoint_blocks=bool(checkpoint_blocks),
@@ -244,6 +248,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
         max_num_scheduled_tokens=8192,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_hash_lookahead_tokens=0,
         hash_block_size=hash_block_size,
         dcp_world_size=dcp_world_size,
         scheduler_block_size=scheduler_block_size,
@@ -294,6 +299,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
         max_num_scheduled_tokens=token_budget,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_hash_lookahead_tokens=0,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -335,6 +341,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
             long_prefill_token_threshold=long_prefill_threshold
         ),
         use_eagle_block_drop=False,
+        num_hash_lookahead_tokens=0,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -1985,6 +1992,7 @@ def test_mamba_align_split_stops_below_eagle_proof_boundary():
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle=True,
         use_eagle_block_drop=True,
+        num_hash_lookahead_tokens=0,
         hash_block_size=hash_block_size,
         mamba_partial_cache_hit=True,
         mamba_shared_prefix_checkpoint=False,
@@ -2793,8 +2801,15 @@ def test_prompt_end_checkpoint_survives_sparse_retention_with_coarse_hash():
 
 
 @pytest.mark.parametrize(
-    "hash_block_size,mamba_block_size,use_eagle",
-    [(32, 32, False), (16, 32, False), (32, 32, True)],
+    "hash_block_size,mamba_block_size,use_eagle,lookahead",
+    [
+        (32, 32, False, 0),
+        (16, 32, False, 0),
+        (32, 32, True, 0),
+        (32, 32, False, 1),
+        (16, 32, False, 1),
+        (16, 32, False, 2),
+    ],
 )
 @pytest.mark.parametrize("retention_interval", [0, None])
 @pytest.mark.parametrize(
@@ -2806,6 +2821,7 @@ def test_checkpoint_reservation_matches_worker(
     hash_block_size,
     mamba_block_size,
     use_eagle,
+    lookahead,
     retention_interval,
     num_prompt_tokens,
     num_output_tokens,
@@ -2824,14 +2840,19 @@ def test_checkpoint_reservation_matches_worker(
         mamba_block_size=mamba_block_size,
         num_blocks=64,
         use_eagle=use_eagle,
-        num_speculative_blocks=3 if use_eagle else 0,
+        num_speculative_blocks=3 if use_eagle or lookahead else 0,
         num_prefill_checkpoint_blocks=1,
+        num_hash_lookahead_tokens=lookahead,
     )
     manager.coordinator.retention_interval = retention_interval
     mamba_manager = manager.coordinator.single_type_managers[1]
 
     request = make_request(
-        "req", list(range(num_prompt_tokens)), hash_block_size, sha256
+        "req",
+        list(range(num_prompt_tokens)),
+        hash_block_size,
+        sha256,
+        num_hash_lookahead_tokens=lookahead,
     )
     request.append_output_token_ids(list(range(1000, 1000 + num_output_tokens)))
     num_tokens = request.num_tokens
@@ -2849,6 +2870,7 @@ def test_checkpoint_reservation_matches_worker(
             checkpoint_alignment=16,
             drop_eagle_block=use_eagle,
             cache_hit_alignment_tokens=mamba_manager.cache_hit_alignment_tokens,
+            num_hash_lookahead_tokens=lookahead,
         )
         reserved = mamba_manager._checkpoints.get(request.request_id)
         if reserved is not None:

tests/v1/core/test_async_scheduler.py

@@ -323,6 +323,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
+    scheduler.num_spec_tokens = 0
     scheduler.make_stats = Mock(return_value=None)
     scheduler.max_model_len = 128
 

tests/v1/core/test_kv_cache_utils.py

@@ -6,7 +6,7 @@ import importlib
 from collections.abc import Callable
 from dataclasses import replace
 from types import SimpleNamespace
-from typing import Any
+from typing import Any, cast
 
 import pytest
 import torch
@@ -28,6 +28,7 @@ from vllm.config import (
 )
 from vllm.config.attention import HiSparseConfig
 from vllm.config.kv_events import KVEventsConfig
+from vllm.config.speculative import SpeculativeConfig
 from vllm.lora.request import LoRARequest
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
@@ -261,6 +262,7 @@ def make_request(
     mm_hashes: list[str] | None = None,
     cache_salt: str | None = None,
     prompt_embeds: torch.Tensor | None = None,
+    num_hash_lookahead_tokens: int = 0,
 ):
     mm_features = []
     if mm_positions is not None:
@@ -285,7 +287,9 @@ def make_request(
         pooling_params=None,
         lora_request=None,
         cache_salt=cache_salt,
-        block_hasher=get_request_block_hasher(block_size, hash_fn),
+        block_hasher=get_request_block_hasher(
+            block_size, hash_fn, num_hash_lookahead_tokens
+        ),
         prompt_embeds=prompt_embeds,
     )
 
@@ -1041,36 +1045,124 @@ def test_request_block_hasher(hash_fn):
 
 
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
-def test_request_block_hasher_incremental_append_with_multiple_mm_features(hash_fn):
+@pytest.mark.parametrize("lookahead", [False, True])
+def test_request_block_hasher_incremental_append_with_multiple_mm_features(
+    hash_fn, lookahead
+):
     mm_positions = [
         PlaceholderRange(offset=4, length=2),
         PlaceholderRange(offset=6, length=1),
     ]
+    num_tokens = 7 + lookahead
     incremental = make_request(
         request_id="incremental",
-        prompt_token_ids=list(range(7)),
+        prompt_token_ids=list(range(num_tokens)),
         block_size=4,
         hash_fn=hash_fn,
         mm_positions=mm_positions,
         mm_hashes=["A", "B"],
+        num_hash_lookahead_tokens=lookahead,
     )
-    incremental.append_output_token_ids(7)
+    incremental.append_output_token_ids(num_tokens)
     fresh = make_request(
         request_id="fresh",
-        prompt_token_ids=list(range(8)),
+        prompt_token_ids=list(range(num_tokens + 1)),
         block_size=4,
         hash_fn=hash_fn,
         mm_positions=mm_positions,
         mm_hashes=["A", "B"],
+        num_hash_lookahead_tokens=lookahead,
     )
 
+    mm_keys = (("mm", "A", 0), ("mm", "B", 2))
+    extra_keys = (mm_keys, 8) if lookahead else mm_keys
     expected_second_hash = hash_fn(
-        (incremental.block_hashes[0], (4, 5, 6, 7), (("mm", "A", 0), ("mm", "B", 2)))
+        (incremental.block_hashes[0], (4, 5, 6, 7), extra_keys)
     )
     assert incremental.block_hashes[1] == expected_second_hash
     assert incremental.block_hashes == fresh.block_hashes
 
 
+@pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
+@pytest.mark.parametrize(
+    ("num_lookahead", "first_mm_keys"),
+    [(1, (("mm", "A", 3),)), (2, (("mm", "A", 3), ("mm", "B", 5)))],
+)
+def test_lookahead_block_hash_keeps_mm_inputs_for_next_block(
+    hash_fn, num_lookahead, first_mm_keys
+):
+    """MM inputs in a block's lookahead tokens still key the next block."""
+    request = make_request(
+        request_id="0",
+        prompt_token_ids=list(range(8 + num_lookahead)),
+        block_size=4,
+        hash_fn=hash_fn,
+        mm_positions=[
+            PlaceholderRange(offset=3, length=2),
+            PlaceholderRange(offset=5, length=1),
+        ],
+        mm_hashes=["A", "B"],
+        num_hash_lookahead_tokens=num_lookahead,
+    )
+
+    first_hash = hash_fn(
+        (
+            kv_cache_utils.NONE_HASH,
+            (0, 1, 2, 3),
+            (first_mm_keys, *range(4, 4 + num_lookahead)),
+        )
+    )
+    second_mm_keys = (("mm", "A", -1), ("mm", "B", 1))
+    second_hash = hash_fn(
+        (first_hash, (4, 5, 6, 7), (second_mm_keys, *range(8, 8 + num_lookahead)))
+    )
+    assert request.block_hashes == [first_hash, second_hash]
+
+
+@pytest.mark.parametrize(
+    ("overrides", "expected"),
+    [
+        *(
+            pytest.param({"method": method}, 1, id=method)
+            for method in ("eagle", "eagle3", "mtp")
+        ),
+        *(
+            pytest.param({"method": method}, 0, id=method)
+            for method in ("dflash", "dspark", "ngram")
+        ),
+        pytest.param({"enable_prefix_caching": False}, 0, id="no-prefix-caching"),
+        pytest.param({"num_prefill_lookahead_tokens": 2}, 2, id="multi-module-mtp"),
+        pytest.param(
+            {"disable_lookahead_block_hashes": True}, 0, id="disabled-by-config"
+        ),
+        pytest.param({"kv_connector": True}, 0, id="kv-connector"),
+    ],
+)
+def test_num_hash_lookahead_tokens(overrides: dict[str, Any], expected: int):
+    speculative_config = object.__new__(SpeculativeConfig)
+    object.__setattr__(speculative_config, "method", overrides.get("method", "mtp"))
+    object.__setattr__(
+        speculative_config,
+        "disable_lookahead_block_hashes",
+        overrides.get("disable_lookahead_block_hashes", False),
+    )
+    vllm_config = cast(
+        VllmConfig,
+        SimpleNamespace(
+            cache_config=SimpleNamespace(
+                enable_prefix_caching=overrides.get("enable_prefix_caching", True)
+            ),
+            speculative_config=speculative_config,
+            num_prefill_lookahead_tokens=overrides.get(
+                "num_prefill_lookahead_tokens", 1
+            ),
+            kv_transfer_config=object() if overrides.get("kv_connector") else None,
+        ),
+    )
+
+    assert VllmConfig.num_hash_lookahead_tokens.fget(vllm_config) == expected
+
+
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])
 def test_hash_tokens_different_mm_input(hash_fn):
     request1 = make_request(

tests/v1/core/test_mamba_align_chunk_split.py

@@ -14,8 +14,11 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from tests.v1.core.test_prefix_caching import make_request
+from vllm.utils.hashing import sha256
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_utils import init_none_hash
 from vllm.v1.core.sched.scheduler import Scheduler
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -84,6 +87,7 @@ def test_mamba_prefill_checkpoint_valid(
 
 def _make_hybrid_kv_cache_manager(
     num_prefill_checkpoint_blocks: int = 0,
+    num_hash_lookahead_tokens: int = 0,
 ) -> KVCacheManager:
     config = KVCacheConfig(
         num_blocks=10000,
@@ -120,7 +124,9 @@ def _make_hybrid_kv_cache_manager(
         scheduler_block_size=MAMBA_BLOCK_SIZE,
         hash_block_size=ATTN_BLOCK_SIZE,
         enable_caching=True,
-        use_eagle=True,
+        use_eagle=not num_hash_lookahead_tokens,
+        num_prefill_lookahead=num_hash_lookahead_tokens,
+        num_hash_lookahead_tokens=num_hash_lookahead_tokens,
     )
 
 
@@ -132,6 +138,7 @@ def _split(
     partial_hit: bool = False,
     num_prefill_checkpoint_blocks: int = 0,
     max_num_scheduled_tokens: int = 16384,
+    num_hash_lookahead_tokens: int = 0,
 ) -> int:
     """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""
     if use_eagle_block_drop is None:
@@ -140,6 +147,7 @@ def _split(
         block_size=MAMBA_BLOCK_SIZE,
         cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),
         use_eagle_block_drop=use_eagle_block_drop,
+        num_hash_lookahead_tokens=num_hash_lookahead_tokens,
         max_num_scheduled_tokens=max_num_scheduled_tokens,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         # `prefix_match_unit` finer than the block size (#46384).
@@ -224,7 +232,7 @@ def test_disabling_eagle_block_drop_keeps_the_trailing_cache_boundary() -> None:
 
 
 def _run_chunked_prefill(
-    manager: KVCacheManager, request: Request, budgets: list[int]
+    manager: KVCacheManager, request: Request, budgets: list[int], **split_kwargs
 ) -> dict[int, int]:
     """Prefill `request`, one step per entry in `budgets`.
 
@@ -243,7 +251,9 @@ def _run_chunked_prefill(
         if computed >= request.num_tokens:
             break
         budget = budgets[step] if step < len(budgets) else request.num_tokens
-        num_new = _split(request, min(request.num_tokens - computed, budget))
+        num_new = _split(
+            request, min(request.num_tokens - computed, budget), **split_kwargs
+        )
         if num_new == 0:
             continue
         assert (
@@ -309,6 +319,42 @@ def test_fragmented_tail_chunk_does_not_poison_mamba_prefix_cache() -> None:
     assert _prefill(3602, budgets=[1600, 1000, 3602]) > 0
 
 
+@pytest.mark.parametrize("lookahead", [1, 2])
+def test_lookahead_split_caches_partial_tail_boundary(lookahead: int) -> None:
+    """Chunks stop where lookahead hashes put the prompt's partial tail, so the
+    manager caches that boundary and a resend hits it."""
+    init_none_hash(sha256)
+    prompt_len = 3601
+    boundary = (prompt_len - lookahead) // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE
+    manager = _make_hybrid_kv_cache_manager(num_hash_lookahead_tokens=lookahead)
+
+    def make(request_id: str) -> Request:
+        return make_request(
+            request_id,
+            list(range(prompt_len)),
+            ATTN_BLOCK_SIZE,
+            sha256,
+            num_hash_lookahead_tokens=lookahead,
+        )
+
+    producer = make("producer")
+    state_at = _run_chunked_prefill(
+        manager,
+        producer,
+        budgets=[],
+        use_eagle_block_drop=False,
+        partial_hit=True,
+        num_hash_lookahead_tokens=lookahead,
+    )
+    assert producer.num_computed_tokens == prompt_len
+    _count_cached_boundary_states(manager, producer, state_at)
+    manager.free(producer)
+    manager.new_step_starts()
+
+    _, num_computed, _ = manager.get_computed_blocks(make("resend"))
+    assert num_computed == boundary
+
+
 @pytest.mark.parametrize("first_chunk", [800, 900, 1599, 1601, 2000])
 def test_intermediate_chunk_ends_stay_block_aligned(first_chunk: int) -> None:
     """Every non-final prefill chunk must end on a mamba block boundary."""

tests/v1/core/test_prefix_caching.py

@@ -31,6 +31,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import (
     OffloadingConnector,
 )
 from vllm.lora.request import LoRARequest
+from vllm.model_executor.layers.mamba.checkpoint import (
+    compute_mamba_prefill_checkpoints,
+)
 from vllm.multimodal.inputs import (
     MultiModalFeatureSpec,
     MultiModalKwargsItem,
@@ -92,6 +95,7 @@ def make_request(
     prompt_logprobs: int | None = None,
     cache_salt: str | None = None,
     lora_request: LoRARequest | None = None,
+    num_hash_lookahead_tokens: int = 0,
     session_id: str | None = None,
 ):
     mm_features = []
@@ -117,7 +121,9 @@ def make_request(
         pooling_params=None,
         lora_request=lora_request,
         cache_salt=cache_salt,
-        block_hasher=get_request_block_hasher(block_size, hash_fn),
+        block_hasher=get_request_block_hasher(
+            block_size, hash_fn, num_hash_lookahead_tokens
+        ),
         session_id=session_id,
     )
 
@@ -2632,6 +2638,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
         max_num_scheduled_tokens=3 * block_size,
         scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
         use_eagle_block_drop=False,
+        num_hash_lookahead_tokens=0,
         hash_block_size=block_size,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -4194,67 +4201,275 @@ def test_emit_cached_block_events_zero_cached():
     assert pool.take_events() == []
 
 
-def test_eagle_enabled_removes_last_block():
-    """Verify Eagle does NOT remove blocks when request
-    length is divisible by block size."""
+@pytest.mark.parametrize(
+    ("lookahead", "num_prompt_tokens", "expected_blocks"),
+    [
+        # Legacy drop: pop the logits block, then drop one more.
+        (False, 48, 1),
+        (False, 37, 1),
+        # Lookahead hashes prove the boundary, so no extra block is dropped.
+        (True, 48, 2),
+        (True, 37, 2),
+    ],
+)
+def test_eagle_identical_prompt_hit(
+    lookahead: bool, num_prompt_tokens: int, expected_blocks: int
+):
     block_size = 16
     manager = make_kv_cache_manager(
         make_kv_cache_config(block_size, num_blocks=10),
         max_model_len=8192,
         enable_caching=True,
-        use_eagle=True,
+        use_eagle=not lookahead,
+        num_hash_lookahead_tokens=lookahead,
         hash_block_size=block_size,
     )
+    token_ids = [0] * num_prompt_tokens
 
-    # Request with 3 full blocks (48 tokens)
-    token_ids = [0] * (3 * block_size)
-    req = make_request("divisible_request", token_ids, block_size, sha256)
-
-    # Prime the cache
-    computed_blocks, _, _ = manager.get_computed_blocks(req)
-    manager.allocate_slots(
-        req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
+    req = make_request(
+        "prime", token_ids, block_size, sha256, num_hash_lookahead_tokens=lookahead
     )
+    assert manager.allocate_slots(req, len(token_ids)) is not None
     manager.free(req)
 
-    # New request with same tokens + Eagle enabled
-    req_eagle = make_request("eagle_divisible", token_ids, block_size, sha256)
+    req_eagle = make_request(
+        "eagle", token_ids, block_size, sha256, num_hash_lookahead_tokens=lookahead
+    )
     computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)
-
-    # Should retain 1 block:
-    # 1. Original 3 blocks → pop last hash → 2 matched blocks
-    # 2. drop last matched block → 1 remaining block
-    assert len(computed_blocks.blocks[0]) == 1
-    assert num_tokens == 1 * block_size  # 16 tokens
+    assert len(computed_blocks.blocks[0]) == expected_blocks
+    assert num_tokens == expected_blocks * block_size
 
 
-def test_eagle_with_partial_blocks():
-    """Test Eagle behavior with requests containing partial blocks."""
-    block_size = 16
+@pytest.mark.parametrize("lookahead", [1, 2])
+def test_kv_events_reconstruct_lookahead_block_hashes(lookahead: int):
+    block_size = 2
     manager = make_kv_cache_manager(
-        make_kv_cache_config(block_size, num_blocks=10),
+        make_kv_cache_config(block_size, num_blocks=20),
         max_model_len=8192,
         enable_caching=True,
-        use_eagle=True,
+        use_eagle=False,
+        num_hash_lookahead_tokens=lookahead,
         hash_block_size=block_size,
+        enable_kv_cache_events=True,
     )
-    # 2 full blocks + 5 tokens (non-divisible length)
-    token_ids = [0] * (2 * block_size + 5)
-    req = make_request("partial_block_test", token_ids, block_size, sha256)
-
-    # Prime the cache
-    computed_blocks, _, _ = manager.get_computed_blocks(req)
-    manager.allocate_slots(
-        req, len(token_ids), len(computed_blocks.blocks[0]) * 16, computed_blocks
+    token_ids = list(range(2 * block_size + lookahead))
+    request = make_request(
+        "request",
+        token_ids,
+        block_size,
+        sha256,
+        num_hash_lookahead_tokens=lookahead,
     )
-    manager.free(req)
+    assert manager.allocate_slots(request, request.num_tokens) is not None
 
-    # New request with Eagle enabled
-    req_eagle = make_request("partial_eagle", token_ids, block_size, sha256)
-    computed_blocks, num_tokens, _ = manager.get_computed_blocks(req_eagle)
-    # Original match: 2 full blocks → Eagle removes 1 → 1 remaining
-    assert len(computed_blocks.blocks[0]) == 1
-    assert num_tokens == 1 * block_size
+    (event,) = manager.take_events()
+    assert isinstance(event, BlockStored)
+    assert event.block_hashes == [
+        kv_cache_utils.maybe_convert_block_hash(block_hash)
+        for block_hash in request.block_hashes
+    ]
+    assert event.token_ids == [0, 1, 2, 3]
+    assert event.lookahead_token_ids == [
+        token_ids[end : end + lookahead] for end in (2, 4)
+    ]
+    reconstructed_hashes = []
+    parent_hash = event.parent_block_hash
+    assert event.extra_keys is not None
+    for block_start, extra_keys, lookahead_token_ids in zip(
+        range(0, len(event.token_ids), block_size),
+        event.extra_keys,
+        event.lookahead_token_ids,
+        strict=True,
+    ):
+        parent_hash = kv_cache_utils.hash_block_tokens(
+            sha256,
+            parent_hash,
+            event.token_ids[block_start : block_start + block_size],
+            (extra_keys, *lookahead_token_ids),
+        )
+        reconstructed_hashes.append(parent_hash)
+    assert reconstructed_hashes == request.block_hashes
+
+
+@pytest.mark.parametrize(
+    "lookahead,prompt_len,expected_boundary",
+    [
+        (0, 1089, 1088),
+        (1, 1089, 1088),
+        (2, 1089, 1024),
+        (2, 1090, 1088),
+        (3, 1089, 1024),
+        (3, 1091, 1088),
+    ],
+)
+@pytest.mark.parametrize("checkpoint_blocks", [0, 1])
+@pytest.mark.parametrize("full_block_size", [64, 128])
+def test_lookahead_checkpoint_has_available_hash(
+    lookahead, prompt_len, expected_boundary, checkpoint_blocks, full_block_size
+):
+    """Prefill checkpoints must be reusable at the boundary their hash proves."""
+    config = _make_hybrid_kv_cache_config(64, 100, ["full", "mamba", "full"])
+    config.kv_cache_groups[0].kv_cache_spec = replace(
+        config.kv_cache_groups[0].kv_cache_spec, block_size=full_block_size
+    )
+    config.kv_cache_groups[1].kv_cache_spec = replace(
+        config.kv_cache_groups[1].kv_cache_spec,
+        block_size=512,
+        num_prefill_checkpoint_blocks=checkpoint_blocks,
+        prefill_checkpoint_alignment=64 if checkpoint_blocks else None,
+    )
+    config.kv_cache_groups[2].is_eagle_group = bool(lookahead)
+    manager = make_kv_cache_manager(
+        config,
+        max_model_len=4096,
+        enable_caching=True,
+        hash_block_size=64,
+        num_prefill_lookahead=lookahead,
+        num_hash_lookahead_tokens=lookahead,
+    )
+    tokens = list(range(prompt_len))
+    request = make_request(
+        "producer", tokens, 64, sha256, num_hash_lookahead_tokens=lookahead
+    )
+    if not checkpoint_blocks:
+        assert manager.allocate_slots(request, expected_boundary) is not None
+        request.num_computed_tokens = expected_boundary
+        manager.new_step_starts()
+    assert manager.allocate_slots(request, len(tokens) - request.num_computed_tokens)
+    checkpoint_hash = request.block_hashes[expected_boundary // 64 - 1]
+    blocks = manager.block_pool.get_cached_block(checkpoint_hash, [1])
+    assert blocks is not None
+    assert blocks[0].block_hash_num_tokens == expected_boundary
+    if checkpoint_blocks:
+        offsets, cols = compute_mamba_prefill_checkpoints(
+            [prompt_len],
+            [prompt_len],
+            hash_block_size=64,
+            mamba_block_size=512,
+            checkpoint_alignment=64,
+            drop_eagle_block=False,
+            cache_hit_alignment_tokens=64,
+            num_hash_lookahead_tokens=lookahead,
+        )
+        assert offsets == [expected_boundary]
+        assert (
+            blocks[0].block_id == manager.get_block_ids(request.request_id)[1][cols[0]]
+        )
+
+    request.num_computed_tokens = len(tokens)
+    manager.free(request)
+    manager.new_step_starts()
+    replay = make_request(
+        "replay", tokens, 64, sha256, num_hash_lookahead_tokens=lookahead
+    )
+    _, hit_tokens, _ = manager.get_computed_blocks(replay)
+    assert hit_tokens == expected_boundary
+
+
+def test_eagle_hybrid_mamba_hits_partial_prompt_boundary():
+    hash_block_size = 16
+    mamba_block_size = 544
+    token_ids = list(range(1249))
+    config = _make_hybrid_kv_cache_config(hash_block_size, 200, ["full", "mamba"])
+    config.kv_cache_groups[1].kv_cache_spec = replace(
+        config.kv_cache_groups[1].kv_cache_spec, block_size=mamba_block_size
+    )
+    manager = make_kv_cache_manager(
+        config,
+        max_model_len=8192,
+        enable_caching=True,
+        use_eagle=False,
+        num_hash_lookahead_tokens=1,
+        hash_block_size=hash_block_size,
+    )
+
+    first = make_request(
+        "first",
+        token_ids,
+        hash_block_size,
+        sha256,
+        num_hash_lookahead_tokens=1,
+    )
+    assert manager.allocate_slots(first, 1248) is not None
+    first.num_computed_tokens = 1248
+    manager.new_step_starts()
+    assert manager.allocate_slots(first, 1) is not None
+    first.num_computed_tokens = 1249
+    manager.new_step_starts()
+    manager.free(first)
+
+    second = make_request(
+        "second",
+        token_ids,
+        hash_block_size,
+        sha256,
+        num_hash_lookahead_tokens=1,
+    )
+    _, num_computed, _ = manager.get_computed_blocks(second)
+
+    assert num_computed == 1248
+    # Per-group lookup must not add the legacy EAGLE drop margin.
+    _, per_group_hits = manager.coordinator.find_longest_cache_hit_per_group(
+        second.block_hashes, 2 * hash_block_size
+    )
+    assert per_group_hits[0] == 2 * hash_block_size
+
+
+@pytest.mark.parametrize("lookahead", [0, 1, 2])
+@pytest.mark.parametrize("diverge_offset", [0, 1, 2])
+def test_mamba_partial_hit_requires_matching_lookahead_tokens(
+    lookahead: int, diverge_offset: int
+):
+    """A partial Mamba hit at boundary b also needs tokens [b, b + lookahead)."""
+    hash_block_size = 16
+    mamba_block_size = 544
+    prompt_len = 1249
+    config = _make_hybrid_kv_cache_config(hash_block_size, 200, ["full", "mamba"])
+    config.kv_cache_groups[1].kv_cache_spec = replace(
+        config.kv_cache_groups[1].kv_cache_spec, block_size=mamba_block_size
+    )
+    manager = make_kv_cache_manager(
+        config,
+        max_model_len=8192,
+        enable_caching=True,
+        use_eagle=False,
+        num_hash_lookahead_tokens=lookahead,
+        hash_block_size=hash_block_size,
+    )
+    boundary = (prompt_len - max(1, lookahead)) // hash_block_size * hash_block_size
+    assert boundary % mamba_block_size != 0
+
+    token_ids = list(range(prompt_len))
+    producer = make_request(
+        "producer",
+        token_ids,
+        hash_block_size,
+        sha256,
+        num_hash_lookahead_tokens=lookahead,
+    )
+    for chunk_end in (boundary, prompt_len):
+        assert (
+            manager.allocate_slots(producer, chunk_end - producer.num_computed_tokens)
+            is not None
+        )
+        producer.num_computed_tokens = chunk_end
+        manager.new_step_starts()
+    manager.free(producer)
+
+    shared_len = boundary + diverge_offset
+    consumer_ids = token_ids[:shared_len] + [
+        100_000 + i for i in range(prompt_len + hash_block_size - shared_len)
+    ]
+    consumer = make_request(
+        "consumer",
+        consumer_ids,
+        hash_block_size,
+        sha256,
+        num_hash_lookahead_tokens=lookahead,
+    )
+    _, num_computed, _ = manager.get_computed_blocks(consumer)
+    # No Mamba state is retained below the prompt's partial boundary.
+    assert num_computed == (boundary if diverge_offset >= lookahead else 0)
 
 
 def test_eagle_with_sliding_window():
@@ -5181,6 +5396,51 @@ def test_hybrid_mamba_retention_mtp_resend_of_aligned_prompt():
     assert num_computed_tokens == 3 * block_size
 
 
+def test_hybrid_mamba_retention_lookahead_keeps_last_prompt_boundary():
+    """Sparse retention keeps the last boundary covered by lookahead hashes."""
+    block_size = 32
+    kv_cache_config = _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"])
+    kv_cache_config.kv_cache_groups[1].kv_cache_spec = replace(
+        kv_cache_config.kv_cache_groups[1].kv_cache_spec,
+        num_speculative_blocks=1,
+    )
+    manager = make_kv_cache_manager(
+        kv_cache_config=kv_cache_config,
+        max_model_len=8192,
+        enable_caching=True,
+        hash_block_size=block_size,
+        retention_interval=0,
+        use_eagle=False,
+        num_hash_lookahead_tokens=1,
+    )
+
+    # 129 tokens: the last boundary below the logits token is 128.
+    token_ids = [i for i in range(4) for _ in range(block_size)] + [4]
+    req0 = make_request("0", token_ids, block_size, sha256, num_hash_lookahead_tokens=1)
+    for chunk_end in (32, 64, 96, 128, 129):
+        blocks = manager.allocate_slots(
+            req0,
+            chunk_end - req0.num_computed_tokens,
+            num_lookahead_tokens=1,
+        )
+        assert blocks is not None
+        req0.num_computed_tokens = chunk_end
+
+    pool = manager.block_pool
+    for i in range(4):
+        cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
+        if i == 3:
+            assert cached is not None, "mamba state at 128 should be retained"
+        else:
+            assert cached is None, f"mamba hash {i} should not be cached"
+    manager.free(req0)
+
+    req1 = make_request("1", token_ids, block_size, sha256, num_hash_lookahead_tokens=1)
+    computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
+    assert num_computed_tokens == 4 * block_size
+    assert [len(blocks) for blocks in computed_blocks.blocks] == [4, 4]
+
+
 def test_block_lookup_cache_single_block_per_key():
     cache = BlockHashToBlockMap()
     key0 = BlockHashWithGroupId(b"hash0")

tests/v1/core/test_scheduler.py

@@ -16,6 +16,7 @@ from vllm.config import (
     KVTransferConfig,
     LoRAConfig,
     ModelConfig,
+    ParallelConfig,
     SchedulerConfig,
     SpeculativeConfig,
     VllmConfig,
@@ -450,6 +451,249 @@ def test_scheduler_stats_route_to_existing_output_client():
     assert len(engine_core_outputs[1].outputs) == 1
 
 
+def test_scheduler_publishes_lookahead_blocks_at_allocation():
+    """A lookahead block is published once its lookahead token exists: prompt
+    blocks at their own allocation, the block ending at the first sampled
+    token at the allocation after it is sampled."""
+    block_size = 2
+    init_none_hash(sha256)
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        block_size=block_size,
+        max_num_batched_tokens=16,
+    )
+
+    def make_request(request_id: str, prompt_token_ids: list[int]) -> Request:
+        return Request(
+            request_id=request_id,
+            prompt_token_ids=prompt_token_ids,
+            sampling_params=SamplingParams(max_tokens=3, ignore_eos=True),
+            pooling_params=None,
+            block_hasher=get_request_block_hasher(block_size, sha256, 1),
+        )
+
+    def num_cached_tokens(prompt_token_ids: list[int]) -> int:
+        probe = make_request("probe", prompt_token_ids)
+        return scheduler.kv_cache_manager.get_computed_blocks(probe)[1]
+
+    request = make_request("first", [0, 1, 2, 3])
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    # Block [2, 4) waits for its lookahead token, the first sampled one.
+    assert num_cached_tokens([0, 1, 2, 3, 5, 7]) == block_size
+
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[5]],
+        ),
+    )
+    scheduler.schedule()
+    assert num_cached_tokens([0, 1, 2, 3, 5, 7]) == 2 * block_size
+    assert num_cached_tokens([0, 1, 2, 3, 6, 7]) == block_size
+
+
+@pytest.mark.parametrize("async_scheduling", [False, True])
+@pytest.mark.parametrize("lookahead", [0, 1])
+@pytest.mark.parametrize("ending", ["prefill", "accepted", "stop", "rejected"])
+def test_finished_spec_request_caches_completed_blocks(
+    async_scheduling, lookahead, ending
+):
+    """Final output can make a computed block's hash available before freeing it."""
+    block_size = 4
+    init_none_hash(sha256)
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        block_size=block_size,
+        num_speculative_tokens=2,
+        speculative_method="ngram_gpu",
+        async_scheduling=async_scheduling,
+    )
+    hasher = get_request_block_hasher(block_size, sha256, lookahead)
+    prompt = list(range(4 if ending == "prefill" else 2))
+    request = Request(
+        request_id="producer",
+        prompt_token_ids=prompt,
+        sampling_params=SamplingParams(
+            max_tokens={"prefill": 1, "rejected": 2}.get(ending, 4),
+            ignore_eos=True,
+            stop_token_ids=[4] if ending == "stop" else None,
+        ),
+        pooling_params=None,
+        block_hasher=hasher,
+    )
+
+    def finish_step(output, tokens):
+        scheduler.update_from_output(
+            output,
+            ModelRunnerOutput(
+                req_ids=[request.request_id],
+                req_id_to_index={request.request_id: 0},
+                sampled_token_ids=[tokens],
+            ),
+        )
+
+    scheduler.add_request(request)
+    output = scheduler.schedule()
+    if ending == "prefill":
+        finish_step(output, [4])
+    else:
+        finish_step(output, [2])
+        scheduler.update_draft_token_ids(
+            DraftTokenIds(req_ids=[request.request_id], draft_token_ids=[[3, 4]])
+        )
+        output = scheduler.schedule()
+        finish_step(output, [8] if ending == "rejected" else [3, 4, 5])
+
+    assert request.is_finished()
+    assert request.num_output_placeholders == 0
+    probe = Request(
+        request_id="probe",
+        prompt_token_ids=list(request.all_token_ids) + [99] * block_size,
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        block_hasher=hasher,
+    )
+    expected = 0 if ending == "rejected" else block_size
+    assert scheduler.kv_cache_manager.get_computed_blocks(probe)[1] == expected
+
+
+@pytest.mark.parametrize("num_spec_tokens", [None, 2])
+def test_async_output_does_not_cache_inflight_tokens(num_spec_tokens):
+    """A returned token ID does not prove that its next-step KV is ready."""
+    init_none_hash(sha256)
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        block_size=1,
+        async_scheduling=True,
+        num_speculative_tokens=num_spec_tokens,
+        speculative_method="ngram_gpu" if num_spec_tokens else None,
+    )
+    hasher = get_request_block_hasher(1, sha256)
+    request = Request(
+        request_id="producer",
+        prompt_token_ids=[0, 1],
+        sampling_params=SamplingParams(max_tokens=16, ignore_eos=True),
+        pooling_params=None,
+        block_hasher=hasher,
+    )
+    scheduler.add_request(request)
+    first = scheduler.schedule()
+    second = scheduler.schedule()
+    assert second.num_scheduled_tokens[request.request_id] > 0
+    scheduler.update_from_output(
+        first,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[2]],
+        ),
+    )
+    probe = Request(
+        request_id="probe",
+        prompt_token_ids=[0, 1, 2, 3],
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        block_hasher=hasher,
+    )
+    assert scheduler.kv_cache_manager.get_computed_blocks(probe)[1] == 2
+
+
+@pytest.mark.parametrize("recovery", ["resident", "preempt", "reset"])
+def test_streaming_lookahead_hashes_follow_recomputed_kv(recovery):
+    """Recomputation uses the continuation's tokens without re-keying resident KV."""
+    block_size = 4
+    model_config = ModelConfig(
+        model="Qwen/Qwen3-Next-80B-A3B-Instruct",
+        max_model_len=128,
+        skip_tokenizer_init=True,
+    )
+    scheduler = create_scheduler(
+        model=model_config.model,
+        max_model_len=128,
+        max_num_batched_tokens=128,
+        enable_prefix_caching=True,
+        block_size=block_size,
+        speculative_config=SpeculativeConfig(
+            method="mtp",
+            num_speculative_tokens=1,
+            target_model_config=model_config,
+            target_parallel_config=ParallelConfig(),
+        ),
+    )
+    assert scheduler.num_hash_lookahead_tokens == 1
+    init_none_hash(sha256)
+    hasher = get_request_block_hasher(block_size, sha256, 1)
+
+    def request(request_id, tokens, max_tokens=16, resumable=False):
+        return Request(
+            request_id=request_id,
+            prompt_token_ids=list(tokens),
+            sampling_params=SamplingParams(max_tokens=max_tokens, ignore_eos=True),
+            pooling_params=None,
+            block_hasher=hasher,
+            resumable=resumable,
+        )
+
+    session = request("session", [1, 2, 3], max_tokens=2, resumable=True)
+    scheduler.add_request(session)
+
+    def step(token):
+        output = scheduler.schedule()
+        scheduler.update_from_output(
+            output,
+            ModelRunnerOutput(
+                req_ids=[session.request_id],
+                req_id_to_index={session.request_id: 0},
+                sampled_token_ids=[[token]],
+            ),
+        )
+
+    step(4)
+    step(5)
+    assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+    manager = scheduler.kv_cache_manager
+    old_block = manager.get_blocks(session.request_id).blocks[0][0]
+    old_key = old_block.block_hash
+    old_hash = session.block_hashes[0]
+    assert old_key is not None
+    old_probe = request("old", [1, 2, 3, 4, 5, 100])
+    assert manager.get_computed_blocks(old_probe)[1] == block_size
+
+    scheduler.add_request(request("session", [42, 43, 44], resumable=True))
+    assert list(session.all_token_ids) == [1, 2, 3, 4, 42, 43, 44]
+    step(99)
+    assert session.block_hashes[0] == old_hash
+    assert old_block.block_hash == old_key
+    current = request("current", session.all_token_ids)
+    assert current.block_hashes[0] != old_hash
+    assert manager.get_computed_blocks(current)[1] == 0
+    if recovery == "resident":
+        return
+
+    if recovery == "reset":
+        assert scheduler.reset_prefix_cache(reset_running_requests=True)
+    else:
+        scheduler.running.remove(session)
+        scheduler._preempt_request(session, 0.0)
+    assert session.num_computed_tokens == 0
+    output = scheduler.schedule()
+    assert session.block_hashes == current.block_hashes
+    assert output.num_scheduled_tokens[session.request_id] == session.num_tokens
+    new_blocks, hit_tokens, _ = manager.get_computed_blocks(current)
+    assert hit_tokens == block_size
+    old_blocks, old_hit_tokens, _ = manager.get_computed_blocks(old_probe)
+    if recovery == "reset":
+        assert old_hit_tokens == 0
+    else:
+        assert old_hit_tokens == block_size
+        assert old_blocks.blocks[0][0] is old_block
+        assert new_blocks.blocks[0][0] is not old_block
+        assert old_block.block_hash == old_key
+
+
 def test_schedule_multimodal_requests():
     scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
     mm_positions = [[PlaceholderRange(offset=i, length=100)] for i in range(10)]
@@ -4743,6 +4987,7 @@ def test_abort_request_when_structured_output_fsm_cannot_advance():
     scheduler.return_sampling_mask = False
     scheduler.recompute_kv_load_failures = False
     scheduler.defer_block_free = False
+    scheduler.num_spec_tokens = 0
     scheduler.make_stats = Mock(return_value=None)
     scheduler.max_model_len = 128
 

tests/v1/core/utils.py

@@ -91,6 +91,7 @@ def create_scheduler(
     diffusion_canvas_length: int | None = None,
     scheduler_cls: type[Scheduler] | None = None,
     lora_config: LoRAConfig | None = None,
+    speculative_config: SpeculativeConfig | None = None,
 ) -> Scheduler | AsyncScheduler:
     """Create scheduler under test.
 
@@ -174,8 +175,8 @@ def create_scheduler(
             kv_connector_extra_config={"shared_storage_path": "local_storage"},
         )
 
-    speculative_config: SpeculativeConfig | None = None
     if num_speculative_tokens is not None:
+        assert speculative_config is None
         spec_kwargs: dict = dict(
             model="ngram", num_speculative_tokens=num_speculative_tokens
         )

tests/v1/e2e/general/test_mamba_prefix_cache.py

@@ -786,17 +786,20 @@ def get_mamba_prefix_cache_step_configs(
                 ),
             ],
         ),
+        # Lookahead hashes reuse the last aligned boundary below the logits
+        # token, so prompts with a partial tail no longer stop one Mamba
+        # page early.
         "prompt_2_block_size_10": TestConfig(
             num_prompt_tokens=B * 2 + 10,
             num_generated_tokens=10,
             num_accepted_tokens=4,
             step_actions=[
-                StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
-                StepAction(B, B + 10, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
+                StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(B * 2, 10, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
                 StepAction(
                     B * 2 + 10,
                     4,
-                    [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
+                    [0, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
                     (-1, -1),
                     (-1, -1),
                 ),
@@ -823,12 +826,18 @@ def get_mamba_prefix_cache_step_configs(
             num_generated_tokens=10,
             num_accepted_tokens=4,
             step_actions=[
-                StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
-                StepAction(B * 2, B + 10, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
+                StepAction(0, B * 3, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
+                StepAction(
+                    B * 3,
+                    10,
+                    [0, 0, 1, 1, 1, 1, 1],
+                    (2, 3),
+                    (-1, -1),
+                ),
                 StepAction(
                     B * 3 + 10,
                     4,
-                    [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
+                    [0, 0, 1, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
                     (-1, -1),
                     (-1, -1),
                 ),
@@ -875,18 +884,18 @@ def get_mamba_prefix_cache_step_configs(
                 StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
                 StepAction(
                     B * 5,
-                    B * 4,
-                    [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
-                    (4, 8),
+                    B * 5,
+                    [0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1],
+                    (4, 9),
                     (-1, -1),
                 ),
                 StepAction(
-                    B * 9,
-                    B + 10,
-                    [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1]
+                    B * 10,
+                    10,
+                    [0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 1, 1, 1, 1]
                     if a
-                    else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
-                    (8, 10),
+                    else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
+                    (9, 10),
                     (-1, -1),
                 ),
             ],
@@ -900,7 +909,7 @@ def get_internal_checkpoint_step_configs() -> dict[str, TestConfig]:
     the prompt's last block boundary itself, so the final chunk runs whole.
 
     The checkpoint goes in the second-to-last block the chunk reaches, keyed to
-    the boundary: column 1 holds the state at B for a 2B + 10 prompt.
+    the boundary: column 1 holds the state at 2B for a 2B + 10 prompt.
     """
     B = BLOCK_SIZE
     return {
@@ -1328,13 +1337,13 @@ def _run_mamba_prefix_cache_mrv2(
         assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
 
         if internal_checkpoints:
-            # The checkpoint a 2B + 10 prefill exports at B (column 1) is the
-            # final state of a prefill that stops at B (column 0). Distinct
+            # The checkpoint a 2B + 10 prefill exports at 2B (column 1) is the
+            # final state of a prefill that stops at 2B (column 1). Distinct
             # tokens make a shifted conv window visible.
             tokens = [(i * 7919) % 30000 + 1 for i in range(BLOCK_SIZE * 2 + 10)]
             scheduler = engine.llm_engine.engine_core.engine_core.scheduler
             states = []
-            for num_tokens, col in ((BLOCK_SIZE, 0), (len(tokens), 1)):
+            for num_tokens, col in ((BLOCK_SIZE * 2, 1), (len(tokens), 1)):
                 captured["snapshot_col"] = col
                 _ = engine.generate(
                     [TokensPrompt(prompt_token_ids=tokens[:num_tokens])],

tests/v1/kv_connector/unit/test_handshake_pp_aggregation.py

@@ -101,6 +101,7 @@ def _run_engine_core_handshake(
         ec_transfer_config=None,
         is_mm_encoder_only=False,
         max_concurrent_batches=1,
+        num_hash_lookahead_tokens=0,
         model_config=SimpleNamespace(
             runner_type="generate",
             is_diffusion=False,

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -1567,6 +1567,7 @@ def test_failed_load_rezeroes_unwritten_skipped_blocks():
 
     scheduler = object.__new__(Scheduler)
     scheduler.connector = MagicMock()
+    scheduler.num_hash_lookahead_tokens = 0
     scheduler.needs_kv_cache_zeroing = True
     scheduler.kv_cache_manager = _make_fake_kv_cache_manager()
     scheduler.kv_cache_manager.cache_blocks = MagicMock()  # type: ignore[method-assign]

vllm/config/speculative.py

@@ -475,6 +475,12 @@ class SpeculativeConfig:
     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."""
+    disable_lookahead_block_hashes: bool = False
+    """Disable lookahead block hashes for EAGLE-like speculative methods and
+    fall back to dropping the trailing prefix-cache block. Lookahead hashes
+    also cover the tokens the drafter reads after each block, so engines
+    sharing a KV store and external consumers of KV cache events must use the
+    same setting."""
     draft_lm_head_quantization: Literal["fp8", "nvfp4"] | None = None
     """Quantize the drafter's copy of a target lm_head it shares (EAGLE and
     MTP). "fp8" uses per-row weight and per-token activation scales on the

vllm/config/vllm.py

@@ -674,6 +674,23 @@ class VllmConfig:
         # Eagle-family drafters read only the immediate next token.
         return 1
 
+    @property
+    def num_hash_lookahead_tokens(self) -> int:
+        """Tokens after each block that its hash covers; 0 keeps the legacy hashes."""
+        speculative_config = self.speculative_config
+        if (
+            speculative_config is None
+            or not speculative_config.use_eagle_block_drop()
+            or speculative_config.disable_lookahead_block_hashes
+            or not self.cache_config.enable_prefix_caching
+        ):
+            return 0
+        # KV connectors, including HiSparse's, keep the legacy hashes.
+        if self.kv_transfer_config is not None:
+            return 0
+        # The drafter's KV at a block's last position reads this far ahead.
+        return self.num_prefill_lookahead_tokens
+
     @property
     def uniform_decode_query_len(self) -> int:
         """Query length of every request in a uniform decode batch.

vllm/distributed/kv_events.py

@@ -92,6 +92,10 @@ class BlockStored(KVCacheEvent):
     ownership of the underlying block, which may be shared across sessions.
     """
     session_id: str | None = None
+    """Tokens beyond each block used in its hash, one list per block_hashes entry.
+    None when lookahead hashing is disabled.
+    """
+    lookahead_token_ids: list[list[int]] | None = None
 
     def __hash__(self) -> int:
         return hash(
@@ -109,6 +113,9 @@ class BlockStored(KVCacheEvent):
                 self.locality,
                 self.ownership,
                 self.session_id,
+                tuple(map(tuple, self.lookahead_token_ids))
+                if self.lookahead_token_ids is not None
+                else None,
             )
         )
 

vllm/model_executor/layers/mamba/checkpoint.py

@@ -26,6 +26,7 @@ def compute_mamba_prefill_checkpoints(
     checkpoint_alignment: int | None,
     drop_eagle_block: bool,
     cache_hit_alignment_tokens: int,
+    num_hash_lookahead_tokens: int = 0,
 ) -> tuple[list[int], list[int]]:
     """Per-row internal prefill checkpoint offsets and cache block columns.
 
@@ -44,7 +45,10 @@ def compute_mamba_prefill_checkpoints(
     for seq_len, query_len in zip(seq_lens, query_lens):
         query_start = seq_len - query_len
         position = get_mamba_prefill_checkpoint_position(
-            seq_len, cache_hit_alignment_tokens, drop_eagle_block=drop_eagle_block
+            seq_len,
+            cache_hit_alignment_tokens,
+            drop_eagle_block=drop_eagle_block,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
         )
         valid = is_mamba_prefill_checkpoint_valid(
             query_start=query_start,
@@ -99,6 +103,12 @@ class MambaPrefillCheckpointBuilder:
     def __init__(self, vllm_config: VllmConfig, kv_cache_spec: MambaSpec) -> None:
         self.vllm_config = vllm_config
         self.kv_cache_spec = kv_cache_spec
+        speculative_config = vllm_config.speculative_config
+        self.drop_eagle_block = (
+            speculative_config is not None
+            and speculative_config.use_eagle_block_drop()
+            and not vllm_config.num_hash_lookahead_tokens
+        )
 
     def build(
         self,
@@ -120,18 +130,15 @@ class MambaPrefillCheckpointBuilder:
         query_lens = [all_query_lens[row] for row in request_rows]
         seq_lens = m.seq_lens_cpu_upper_bound.tolist()
         block_size = self.kv_cache_spec.block_size
-        speculative_config = self.vllm_config.speculative_config
-        drop_eagle_block = (
-            speculative_config is not None and speculative_config.use_eagle_block_drop()
-        )
         checkpoint_offsets, checkpoint_cols = compute_mamba_prefill_checkpoints(
             [seq_lens[row] for row in request_rows],
             query_lens,
             hash_block_size=hash_block_size,
             mamba_block_size=block_size,
             checkpoint_alignment=self.kv_cache_spec.prefill_checkpoint_alignment,
-            drop_eagle_block=drop_eagle_block,
+            drop_eagle_block=self.drop_eagle_block,
             cache_hit_alignment_tokens=cache_hit_alignment_tokens,
+            num_hash_lookahead_tokens=self.vllm_config.num_hash_lookahead_tokens,
         )
         if not any(checkpoint_offsets):
             return None

vllm/v1/core/block_pool.py

@@ -20,6 +20,7 @@ from vllm.v1.core.kv_cache_utils import (
     FreeKVCacheBlockQueue,
     KVCacheBlock,
     generate_block_hash_extra_keys,
+    generate_lookahead_block_hash_extra_keys,
     get_block_hash,
     get_group_id,
     make_block_hash_with_group_id,
@@ -150,6 +151,7 @@ class BlockPool:
             actual block size can be a multiple of hash_block_size.
         enable_kv_cache_events: Whether to enable kv cache events.
         metrics_collector: Optional metrics collector for tracking block residency.
+        num_hash_lookahead_tokens: Tokens after each block covered by its hash.
         medium: Storage medium reported in KV cache events.
 
     """
@@ -161,6 +163,7 @@ class BlockPool:
         hash_block_size: int,
         enable_kv_cache_events: bool = False,
         metrics_collector: KVCacheMetricsCollector | None = None,
+        num_hash_lookahead_tokens: int = 0,
         medium: str = MEDIUM_GPU,
     ):
         assert isinstance(num_gpu_blocks, int) and num_gpu_blocks > 0
@@ -168,6 +171,7 @@ class BlockPool:
         self.medium = medium
         self.enable_caching = enable_caching
         self.hash_block_size = hash_block_size
+        self.num_hash_lookahead_tokens = num_hash_lookahead_tokens
         # All kv-cache blocks.
         self.blocks: list[KVCacheBlock] = [
             KVCacheBlock(idx, pool=self) for idx in range(num_gpu_blocks)
@@ -317,6 +321,9 @@ class BlockPool:
             # features, or cache_salt only for the first block).
             # Skip null/masked-out blocks to match the length of new_hashes.
             extra_keys_list: list[tuple[Any, ...] | None] = []
+            lookahead_token_ids: list[list[int]] | None = (
+                [] if self.num_hash_lookahead_tokens else None
+            )
             curr_mm_idx = 0
             for i in range(num_cached_blocks, num_full_blocks):
                 if blocks[i].is_null:
@@ -325,10 +332,14 @@ class BlockPool:
                     continue
                 block_start = i * block_size
                 block_end = block_start + block_size
-                extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
+                extra_keys, curr_mm_idx = self._generate_event_extra_keys(
                     request, block_start, block_end, curr_mm_idx
                 )
                 extra_keys_list.append(extra_keys)
+                if lookahead_token_ids is not None:
+                    lookahead_token_ids.append(
+                        self._lookahead_token_ids(request, block_end)
+                    )
 
             self.kv_event_queue.append(
                 self._build_block_stored_event(
@@ -340,9 +351,25 @@ class BlockPool:
                     block_size=block_size,
                     kv_cache_group_id=kv_cache_group_id,
                     extra_keys_list=extra_keys_list,
+                    lookahead_token_ids=lookahead_token_ids,
                 )
             )
 
+    def _generate_event_extra_keys(
+        self, request: Request, start: int, end: int, start_mm_idx: int
+    ) -> tuple[tuple[Any, ...] | None, int]:
+        # Event extra keys cover the same window as the block hash.
+        if self.num_hash_lookahead_tokens:
+            return generate_lookahead_block_hash_extra_keys(
+                request, start, end, end + self.num_hash_lookahead_tokens, start_mm_idx
+            )
+        return generate_block_hash_extra_keys(request, start, end, start_mm_idx)
+
+    def _lookahead_token_ids(self, request: Request, block_end: int) -> list[int]:
+        return request.all_token_ids[
+            block_end : block_end + self.num_hash_lookahead_tokens
+        ]
+
     def _build_block_stored_event(
         self,
         request: Request,
@@ -353,6 +380,7 @@ class BlockPool:
         block_size: int,
         kv_cache_group_id: int,
         extra_keys_list: list[tuple[Any, ...] | None],
+        lookahead_token_ids: list[list[int]] | None,
     ) -> BlockStored:
         """Build a ``BlockStored`` KV event for ``request``.
 
@@ -371,6 +399,7 @@ class BlockPool:
             extra_keys=to_event_extra_keys(extra_keys_list),
             group_idx=kv_cache_group_id,
             session_id=request.session_id,
+            lookahead_token_ids=lookahead_token_ids,
         )
 
     def emit_cached_block_events(
@@ -403,15 +432,22 @@ class BlockPool:
         # Collect external hashes and extra_keys for cached blocks.
         cached_hashes: list[ExternalBlockHash] = []
         extra_keys_list: list[tuple[Any, ...] | None] = []
+        lookahead_token_ids: list[list[int]] | None = (
+            [] if self.num_hash_lookahead_tokens else None
+        )
         curr_mm_idx = 0
         for i in range(num_cached_blocks):
             block_start = i * block_size
             block_end = block_start + block_size
             cached_hashes.append(maybe_convert_block_hash(block_hashes[i]))
-            extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
+            extra_keys, curr_mm_idx = self._generate_event_extra_keys(
                 request, block_start, block_end, curr_mm_idx
             )
             extra_keys_list.append(extra_keys)
+            if lookahead_token_ids is not None:
+                lookahead_token_ids.append(
+                    self._lookahead_token_ids(request, block_end)
+                )
 
         if not cached_hashes:
             return
@@ -443,6 +479,7 @@ class BlockPool:
                 block_size=block_size,
                 kv_cache_group_id=kv_cache_group_id,
                 extra_keys_list=extra_keys_list,
+                lookahead_token_ids=lookahead_token_ids,
             )
         )
 
@@ -539,7 +576,7 @@ class BlockPool:
                     curr_mm_idx, _ = get_mm_features_in_window(
                         mm_features, block_start, block_end
                     )
-            extra_keys, _ = generate_block_hash_extra_keys(
+            extra_keys, _ = self._generate_event_extra_keys(
                 request, block_start, block_end, curr_mm_idx
             )
             self.kv_event_queue.append(
@@ -558,6 +595,9 @@ class BlockPool:
                     extra_keys=to_event_extra_keys([extra_keys]),
                     group_idx=kv_cache_group_id,
                     session_id=request.session_id,
+                    lookahead_token_ids=[self._lookahead_token_ids(request, block_end)]
+                    if self.num_hash_lookahead_tokens
+                    else None,
                 )
             )
         return block_hash_with_group_id

vllm/v1/core/kv_cache_coordinator.py

@@ -105,6 +105,7 @@ class KVCacheCoordinator(ABC):
         hash_block_size: int,
         metrics_collector: KVCacheMetricsCollector | None = None,
         num_prefill_lookahead: int = 0,
+        num_hash_lookahead_tokens: int = 0,
     ):
         self.kv_cache_config = kv_cache_config
         # The scheduling granularity (LCM of all group block sizes), must be a multiple
@@ -122,6 +123,7 @@ class KVCacheCoordinator(ABC):
             hash_block_size=hash_block_size,
             enable_kv_cache_events=enable_kv_cache_events,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
         )
 
         # KV cache group indices that get the EAGLE last-block drop.
@@ -131,6 +133,8 @@ class KVCacheCoordinator(ABC):
         # Conservatively fall back to flag all groups when no group is flagged.
         if use_eagle and not self.eagle_group_ids:
             self.eagle_group_ids = set(range(len(kv_cache_config.kv_cache_groups)))
+        # Lookahead hashes clear ``use_eagle``: groups stay flagged, none drops.
+        self.drop_eagle_group_ids = self.eagle_group_ids if use_eagle else set()
 
         # During chunked prefill with EAGLE, the single next prefill lookahead
         # token past the chunk boundary is combined with the final hidden state
@@ -364,7 +368,7 @@ class KVCacheCoordinator(ABC):
         boundaries = replay_boundaries(
             request.num_prompt_tokens,
             self.scheduler_block_size,
-            use_eagle=bool(self.eagle_group_ids),
+            use_eagle=bool(self.drop_eagle_group_ids),
         )
         for group_id, manager in enumerate(self.single_type_managers):
             if not manager.enable_caching:
@@ -492,6 +496,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
         scheduler_block_size: int,
         hash_block_size: int,
         metrics_collector: KVCacheMetricsCollector | None = None,
+        num_hash_lookahead_tokens: int = 0,
         num_prefill_lookahead: int = 0,
     ):
         super().__init__(
@@ -506,6 +511,7 @@ class KVCacheCoordinatorNoPrefixCache(KVCacheCoordinator):
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             num_prefill_lookahead=num_prefill_lookahead,
         )
         self.num_single_type_manager = len(self.single_type_managers)
@@ -540,6 +546,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
         scheduler_block_size: int,
         hash_block_size: int,
         metrics_collector: KVCacheMetricsCollector | None = None,
+        num_hash_lookahead_tokens: int = 0,
         num_prefill_lookahead: int = 0,
     ):
         super().__init__(
@@ -554,6 +561,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             num_prefill_lookahead=num_prefill_lookahead,
         )
         self.kv_cache_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec
@@ -569,7 +577,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
             "UnitaryKVCacheCoordinator assumes only one kv cache group"
         )
         # Single group; useless but just set ``use_eagle`` for consistency regardless.
-        self.single_type_managers[0].use_eagle = 0 in self.eagle_group_ids
+        self.single_type_managers[0].use_eagle = 0 in self.drop_eagle_group_ids
 
     def find_longest_cache_hit(
         self,
@@ -582,7 +590,7 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
             kv_cache_group_ids=[0],
             block_pool=self.block_pool,
             kv_cache_spec=self.kv_cache_spec,
-            drop_eagle_block=0 in self.eagle_group_ids,
+            drop_eagle_block=0 in self.drop_eagle_group_ids,
             alignment_tokens=self.block_size,
             dcp_world_size=self.dcp_world_size,
             pcp_world_size=self.pcp_world_size,
@@ -624,6 +632,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
         scheduler_block_size: int,
         hash_block_size: int,
         metrics_collector: KVCacheMetricsCollector | None = None,
+        num_hash_lookahead_tokens: int = 0,
         num_prefill_lookahead: int = 0,
         allow_partial_hash_hits: bool = True,
     ):
@@ -639,6 +648,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             num_prefill_lookahead=num_prefill_lookahead,
         )
         # hash_block_size: the block size used to compute block hashes.
@@ -720,7 +730,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
                 continue
             manager_cls = self.single_type_managers[i].__class__
             spec = g.kv_cache_spec
-            use_eagle = i in self.eagle_group_ids
+            use_eagle = i in self.drop_eagle_group_ids
 
             # Try to find an existing group with the same spec
             for idx, group in enumerate(self.attention_groups):
@@ -989,6 +999,7 @@ def get_kv_cache_coordinator(
     scheduler_block_size: int,
     hash_block_size: int,
     metrics_collector: KVCacheMetricsCollector | None = None,
+    num_hash_lookahead_tokens: int = 0,
     num_prefill_lookahead: int = 0,
     allow_partial_hash_hits: bool = True,
 ) -> KVCacheCoordinator:
@@ -1004,6 +1015,7 @@ def get_kv_cache_coordinator(
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             num_prefill_lookahead=num_prefill_lookahead,
         )
     if len(kv_cache_config.kv_cache_groups) == 1:
@@ -1019,6 +1031,7 @@ def get_kv_cache_coordinator(
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             metrics_collector=metrics_collector,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             num_prefill_lookahead=num_prefill_lookahead,
         )
     return HybridKVCacheCoordinator(
@@ -1033,6 +1046,7 @@ def get_kv_cache_coordinator(
         scheduler_block_size=scheduler_block_size,
         hash_block_size=hash_block_size,
         metrics_collector=metrics_collector,
+        num_hash_lookahead_tokens=num_hash_lookahead_tokens,
         num_prefill_lookahead=num_prefill_lookahead,
         allow_partial_hash_hits=allow_partial_hash_hits,
     )

vllm/v1/core/kv_cache_manager.py

@@ -138,6 +138,7 @@ class KVCacheManager:
         max_in_flight_tokens: int | None = None,
         enable_caching: bool = True,
         use_eagle: bool = False,
+        num_hash_lookahead_tokens: int = 0,
         num_prefill_lookahead: int = 0,
         log_stats: bool = False,
         enable_kv_cache_events: bool = False,
@@ -169,6 +170,7 @@ class KVCacheManager:
             max_model_len=self.max_model_len,
             max_in_flight_tokens=max_in_flight_tokens,
             use_eagle=self.use_eagle,
+            num_hash_lookahead_tokens=num_hash_lookahead_tokens,
             enable_caching=self.enable_caching,
             enable_kv_cache_events=enable_kv_cache_events,
             dcp_world_size=dcp_world_size,

vllm/v1/core/kv_cache_utils.py

@@ -953,6 +953,7 @@ def resolve_cache_hit_alignment_tokens(
 def get_request_block_hasher(
     hash_block_size: int,
     caching_hash_fn: Callable[[Any], bytes],
+    num_hash_lookahead_tokens: int = 0,
 ) -> Callable[[Request], list[BlockHash]]:
     """Returns a function which computes the list of un-computed block hashes
     of a request.
@@ -961,11 +962,17 @@ def get_request_block_hasher(
     full prefix, so each hash uniquely fingerprints the prefix ending at its
     boundary. Coarser group block sizes and partial-cache boundaries reuse
     these hashes directly (see ``BlockHashListWithBlockSize``).
+
+    With ``num_hash_lookahead_tokens`` (L) > 0, each hash also covers the L
+    tokens after its block, which the drafter's KV at the block's last position
+    depends on, so it is emitted only once those tokens exist. The block tokens
+    stay the hashed tokens; the extra keys become ``(E, *tokens)``, where ``E``
+    covers the block and those tokens.
     """
 
     def request_block_hasher(request: Request) -> list[BlockHash]:
         start_token_idx = len(request.block_hashes) * hash_block_size
-        num_tokens = request.num_tokens
+        num_tokens = request.num_tokens - num_hash_lookahead_tokens
 
         if start_token_idx + hash_block_size > num_tokens:
             # Early stop when there no new full blocks created.
@@ -993,9 +1000,26 @@ def get_request_block_hasher(
                 break
 
             # MM and LoRA requests need extra keys for block-hash computation.
-            extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
-                request, start_token_idx, end_token_idx, curr_mm_idx
-            )
+            if num_hash_lookahead_tokens:
+                hash_end_token_idx = end_token_idx + num_hash_lookahead_tokens
+                extra_keys, curr_mm_idx = generate_lookahead_block_hash_extra_keys(
+                    request,
+                    start_token_idx,
+                    end_token_idx,
+                    hash_end_token_idx,
+                    curr_mm_idx,
+                )
+                if num_hash_lookahead_tokens == 1:
+                    extra_keys = (extra_keys, request.all_token_ids[end_token_idx])
+                else:
+                    extra_keys = (
+                        extra_keys,
+                        *request.all_token_ids[end_token_idx:hash_end_token_idx],
+                    )
+            else:
+                extra_keys, curr_mm_idx = generate_block_hash_extra_keys(
+                    request, start_token_idx, end_token_idx, curr_mm_idx
+                )
 
             # Compute the hash of the current block
             block_tokens = request.all_token_ids[start_token_idx:end_token_idx]
@@ -1012,6 +1036,28 @@ def get_request_block_hasher(
     return request_block_hasher
 
 
+def generate_lookahead_block_hash_extra_keys(
+    request: Request,
+    start_token_idx: int,
+    end_token_idx: int,
+    hash_end_token_idx: int,
+    start_mm_idx: int,
+) -> tuple[tuple[Any, ...] | None, int]:
+    """Generate extra keys for the block and its lookahead tokens.
+
+    Rewind the MM index to retain inputs overlapping the next block.
+    """
+    extra_keys, mm_idx = generate_block_hash_extra_keys(
+        request, start_token_idx, hash_end_token_idx, start_mm_idx
+    )
+    while mm_idx > 0:
+        mm_position = request.mm_features[mm_idx - 1].mm_position
+        if mm_position.offset + mm_position.length <= end_token_idx:
+            break
+        mm_idx -= 1
+    return extra_keys, mm_idx
+
+
 def get_max_concurrency_for_kv_cache_config(
     vllm_config: VllmConfig, kv_cache_config: KVCacheConfig
 ) -> float:

vllm/v1/core/sched/async_scheduler.py

@@ -6,7 +6,6 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler
-from vllm.v1.request import Request, RequestStatus
 
 logger = init_logger(__name__)
 
@@ -55,26 +54,3 @@ class AsyncScheduler(Scheduler):
                 request.next_decode_eligible_step = (
                     self.current_step + self.decode_stagger
                 )
-
-    def _update_request_with_output(
-        self, request: Request, new_token_ids: list[int], is_stale: bool = False
-    ) -> tuple[list[int], bool]:
-        status_before_update = request.status
-        # Tokens trimmed at a stop resolve their placeholders too.
-        num_new_tokens = len(new_token_ids)
-        new_token_ids, stopped = super()._update_request_with_output(
-            request, new_token_ids, is_stale
-        )
-
-        # Placeholders were zeroed at preemption; a stale delivery must not
-        # decrement them (it would underflow).
-        if not is_stale:
-            request.num_output_placeholders -= num_new_tokens
-            assert request.num_output_placeholders >= 0
-
-        # Cache the new tokens. Preempted requests should be skipped.
-        if status_before_update == RequestStatus.RUNNING:
-            self.kv_cache_manager.cache_blocks(
-                request, request.num_computed_tokens - request.num_output_placeholders
-            )
-        return new_token_ids, stopped

vllm/v1/core/sched/scheduler.py

@@ -331,6 +331,11 @@ class Scheduler(SchedulerInterface):
                     "acceptance rates."
                 )
 
+        self.num_hash_lookahead_tokens = vllm_config.num_hash_lookahead_tokens
+        if self.num_hash_lookahead_tokens:
+            # Lookahead hashes prove the boundary token, so EAGLE needs no drop.
+            self.use_eagle_block_drop = False
+
         # Create the KV cache manager.
         if hash_block_size is None:
             hash_block_size = block_size
@@ -341,6 +346,7 @@ class Scheduler(SchedulerInterface):
             max_in_flight_tokens=vllm_config.max_in_flight_tokens,
             enable_caching=self.cache_config.enable_prefix_caching,
             use_eagle=self.use_eagle_block_drop,
+            num_hash_lookahead_tokens=self.num_hash_lookahead_tokens,
             num_prefill_lookahead=self.num_prefill_lookahead,
             log_stats=self.log_stats,
             enable_kv_cache_events=self.enable_kv_cache_events,
@@ -561,6 +567,7 @@ class Scheduler(SchedulerInterface):
             request.num_tokens,
             block_size,
             self.use_eagle_block_drop,
+            num_hash_lookahead_tokens=self.num_hash_lookahead_tokens,
         )
 
         end = start + num_new_tokens
@@ -568,6 +575,7 @@ class Scheduler(SchedulerInterface):
             prefill_end,
             ckpt_block_size,
             drop_eagle_block=self.use_eagle_block_drop,
+            num_hash_lookahead_tokens=self.num_hash_lookahead_tokens,
         )
         use_internal_checkpoint = (
             self.mamba_has_prefill_checkpoint_blocks
@@ -603,6 +611,7 @@ class Scheduler(SchedulerInterface):
                 request.num_prompt_tokens,
                 self.hash_block_size,
                 self.use_eagle_block_drop,
+                num_hash_lookahead_tokens=self.num_hash_lookahead_tokens,
             )
             if self.mamba_partial_cache_hit and not use_internal_checkpoint
             else 0
@@ -1086,6 +1095,11 @@ class Scheduler(SchedulerInterface):
 
                 # Get already-cached tokens.
                 if request.num_computed_tokens == 0:
+                    if request.resumable and self.num_hash_lookahead_tokens:
+                        # Resident KV can retain a discarded streaming token's
+                        # hash. Recomputing KV must use the current token sequence.
+                        request.block_hashes.clear()
+                        request.update_block_hashes()
                     did_prefix_cache_lookup = True
                     max_cache_hit_length = (
                         self.aux_output_connector.max_cache_hit(request)
@@ -2353,9 +2367,19 @@ class Scheduler(SchedulerInterface):
 
             # Check for stop and update request status.
             if new_token_ids:
+                has_output_placeholders = request.num_output_placeholders > 0
                 new_token_ids, stopped = self._update_request_with_output(
                     request, new_token_ids, is_stale=output_is_stale
                 )
+                if (
+                    not output_is_stale
+                    and status_before_stop == RequestStatus.RUNNING
+                    and (has_output_placeholders or self.num_spec_tokens > 0)
+                ):
+                    self.kv_cache_manager.cache_blocks(
+                        request,
+                        request.num_computed_tokens - request.num_in_flight_tokens,
+                    )
             elif request.pooling_params and pooler_output is not None:
                 # Pooling stops as soon as there is output.
                 request.status = RequestStatus.FINISHED_STOPPED
@@ -2687,6 +2711,7 @@ class Scheduler(SchedulerInterface):
     def _update_request_with_output(
         self, request: Request, new_token_ids: list[int], is_stale: bool = False
     ) -> tuple[list[int], bool]:
+        num_new_tokens = len(new_token_ids)
         # Append generated tokens and check for stop. Note that if
         # a request is still being prefilled, we expect the model runner
         # to return empty token ids for the request.
@@ -2704,6 +2729,11 @@ class Scheduler(SchedulerInterface):
                     request.num_computed_tokens -= len(new_token_ids) - num_new
                 del new_token_ids[num_new:]  # Trim new tokens if needed.
                 break
+
+        if not is_stale and request.num_output_placeholders > 0:
+            # Resolve all returned tokens, including those trimmed at a stop.
+            request.num_output_placeholders -= num_new_tokens
+            assert request.num_output_placeholders >= 0
         return new_token_ids, stopped
 
     def _free_encoder_inputs(self, request: Request) -> None:

vllm/v1/core/single_type_kv_cache_manager.py

@@ -499,7 +499,8 @@ class SingleTypeKVCacheManager(ABC):
         if not self.kv_cache_spec.prefix_cacheable:
             return
         num_cached_blocks = self.num_cached_block.get(request.request_id, 0)
-        num_full_blocks = num_tokens // self.block_size
+        num_hashed_tokens = len(request.block_hashes) * self.block_pool.hash_block_size
+        num_full_blocks = min(num_tokens, num_hashed_tokens) // self.block_size
 
         if num_cached_blocks >= num_full_blocks:
             return
@@ -874,7 +875,10 @@ class FullAttentionManager(SingleTypeKVCacheManager):
             if self.use_eagle
             else request.num_prompt_tokens - 1
         )
-        boundary_tokens = token_limit // hash_block_size * hash_block_size
+        boundary_tokens = (
+            min(token_limit // hash_block_size, len(request.block_hashes))
+            * hash_block_size
+        )
         if boundary_tokens == 0 or boundary_tokens > num_tokens:
             return
         if boundary_tokens % self.block_size == 0:
@@ -1758,6 +1762,7 @@ class MambaManager(SingleTypeKVCacheManager):
                 num_tokens,
                 self.cache_hit_alignment_tokens,
                 self.drop_eagle_checkpoint_block,
+                num_hash_lookahead_tokens=self.block_pool.num_hash_lookahead_tokens,
             )
             if not self._needs_internal_checkpoint(
                 request_id,
@@ -2043,6 +2048,7 @@ class MambaManager(SingleTypeKVCacheManager):
                     request.num_prompt_tokens,
                     self.cache_hit_alignment_tokens,
                     self.drop_eagle_checkpoint_block,
+                    num_hash_lookahead_tokens=self.block_pool.num_hash_lookahead_tokens,
                 )
             )
             is_retained_checkpoint = False
@@ -2082,6 +2088,7 @@ class MambaManager(SingleTypeKVCacheManager):
             request.num_prompt_tokens,
             hash_block_size,
             self.drop_eagle_checkpoint_block,
+            num_hash_lookahead_tokens=self.block_pool.num_hash_lookahead_tokens,
         )
         # The junction is the other position a sibling resumes at: where one was
         # observed to stop, and where the scheduler already ends a chunk. Bounded

vllm/v1/engine/core.py

@@ -221,7 +221,9 @@ class EngineCore:
             init_none_hash(caching_hash_fn)
 
             self.request_block_hasher = get_request_block_hasher(
-                hash_block_size, caching_hash_fn
+                hash_block_size,
+                caching_hash_fn,
+                vllm_config.num_hash_lookahead_tokens,
             )
 
         self.step_fn = (

vllm/v1/kv_cache_interface.py

@@ -1170,11 +1170,16 @@ def get_mamba_prefill_checkpoint_position(
     num_tokens: int,
     hash_block_size: int,
     drop_eagle_block: bool,
+    num_hash_lookahead_tokens: int = 0,
 ) -> int:
     """Return the reusable Mamba checkpoint boundary for a prefill."""
-    # Without EAGLE, leave one token to recompute on resend.
+    # Leave tokens for logits recomputation and the checkpoint's hash.
     # EAGLE's block drop already leaves tokens to recompute.
-    proof_limit = num_tokens if drop_eagle_block else num_tokens - 1
+    proof_limit = (
+        num_tokens
+        if drop_eagle_block
+        else num_tokens - max(1, num_hash_lookahead_tokens)
+    )
     checkpoint_position = proof_limit // hash_block_size * hash_block_size
     if drop_eagle_block:
         checkpoint_position -= hash_block_size