[Core][Spec Decode] Lookahead block hashes for EAGLE prefix caching

tests/models/kimi_k3/test_kda_metadata.py

@@ -15,6 +15,7 @@
 )
 from vllm.config import SpeculativeConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.model_executor.layers.mamba import checkpoint as mamba_checkpoint
 from vllm.models.kimi_k3.nvidia.kda_metadata import (
     KimiK3KDAAttentionBackend,
     KimiK3KDAMetadata,
@@ -358,22 +359,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",
+        "lookahead_hashing",
+        "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, False, 16, 80),
+        (100, True, False, 16, 96),
+        (100, False, True, 16, 96),
+        (100, False, False, 8, None),
+        (96, False, False, 16, 80),
+        (96, True, False, 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,
+    lookahead_hashing: bool,
     prefix_match_unit: int,
     expected_offset: int | None,
 ) -> None:
+    """Lookahead hashes need no back-off from the replay boundary."""
+    monkeypatch.setattr(
+        mamba_checkpoint,
+        "is_lookahead_block_hashing_enabled",
+        lambda *_: lookahead_hashing,
+    )
     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

@@ -36,19 +36,24 @@ def _manager(
     num_blocks=8192,
     eagle_group=None,
     num_prefill_lookahead=0,
+    lookahead=False,
 ):
     init_none_hash(sha256)
     config = _make_hybrid_kv_cache_config(block_size, num_blocks, ["full", "mamba"])
+    groups = list(config.kv_cache_groups)
     if eagle_group is not None:
-        groups = list(config.kv_cache_groups)
         groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True)
-        config = replace(config, kv_cache_groups=groups)
+    elif lookahead:
+        # The scheduler clears ``use_eagle`` here, so flag the groups up front.
+        groups = [replace(group, is_eagle_group=True) for group in groups]
+    config = replace(config, kv_cache_groups=groups)
     return make_kv_cache_manager(
         kv_cache_config=config,
         max_model_len=1 << 20,
         enable_caching=True,
         hash_block_size=hash_block_size,
-        use_eagle=True,
+        use_eagle=not lookahead,
+        use_lookahead_block_hashes=lookahead,
         num_prefill_lookahead=num_prefill_lookahead,
         enable_mamba_shared_prefix_checkpoint=shared_prefix_checkpoint,
     )
@@ -69,8 +74,12 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True):
         max_num_scheduled_tokens=1 << 20,
         use_eagle=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,
+        # they exist only to compensate for the drop, which lookahead hashes
+        # make unnecessary.
+        use_eagle_block_drop=(
+            block_drop and not manager.block_pool.use_lookahead_block_hashes
+        ),
+        use_lookahead_block_hashes=manager.block_pool.use_lookahead_block_hashes,
         hash_block_size=hash_block_size,
         mamba_has_prefill_checkpoint_blocks=False,  # forced False under eagle
         mamba_partial_cache_hit=partial_hit,
@@ -121,6 +130,16 @@ def _prefill(manager, stub, request, *, external=0) -> list[int]:
     return ends
 
 
+def _request(request_id, tokens, stub):
+    return make_request(
+        request_id,
+        tokens,
+        stub.hash_block_size,
+        sha256,
+        use_lookahead_hashes=stub.use_lookahead_block_hashes,
+    )
+
+
 def _orphaned_full_attention_tail(manager, stub, prompt_len):
     """Produce the state a KV connector leaves behind, which creates a junction.
 
@@ -129,37 +148,34 @@ def _orphaned_full_attention_tail(manager, stub, prompt_len):
     producer starts past its own tail stop: full attention still registers a
     partial tail the Mamba group has no state for.
     """
-    producer = make_request(
-        "producer", PREFIX[:prompt_len], stub.hash_block_size, sha256
-    )
+    producer = _request("producer", PREFIX[:prompt_len], stub)
     _prefill(manager, stub, producer, external=prompt_len - 1)
 
 
-def _sibling_hit(manager, shared, suffix, hash_block_size):
-    request = make_request("sibling", PREFIX[:shared] + suffix, hash_block_size, sha256)
+def _sibling_hit(manager, stub, shared, suffix):
+    request = _request("sibling", PREFIX[:shared] + suffix, stub)
     return manager.get_computed_blocks(request)[1]
 
 
 # --------------------------------------------------------------------------
 # The two positions a sibling resumes at
 
 
-def test_sibling_resumes_from_the_observed_junction():
+@pytest.mark.parametrize("lookahead", [False, True])
+def test_sibling_resumes_from_the_observed_junction(lookahead):
     """The scheduler splits at the junction; the manager must cache there."""
     block_size, hash_block_size = 512, 32
-    manager = _manager(block_size, hash_block_size)
+    manager = _manager(block_size, hash_block_size, lookahead=lookahead)
     stub = _stub(manager, block_size, hash_block_size)
     assert stub.mamba_shared_prefix_checkpoint, "the feature must be armed"
     _orphaned_full_attention_tail(manager, stub, 2020)
 
-    consumer = make_request(
-        "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
-    )
+    consumer = _request("consumer", PREFIX[:2016] + [-1] * 584, stub)
     junction = manager.get_computed_blocks(consumer)[2]
     assert junction, "expected a junction against the orphaned full-attention tail"
     _prefill(manager, stub, consumer)
 
-    hit = _sibling_hit(manager, 2016, [-2] * 584, hash_block_size)
+    hit = _sibling_hit(manager, stub, 2016, [-2] * 584)
     assert hit == junction, (
         f"the chunk stopped at {junction} but nothing was cached there: "
         f"sibling resumes at {hit}"
@@ -179,18 +195,16 @@ def test_sibling_resumes_below_the_block_grid_when_the_prefix_ends_early():
     stub = _stub(manager, block_size, hash_block_size)
 
     shared = 24
-    owner = make_request("owner", PREFIX[:shared] + [-1] * 16, hash_block_size, sha256)
+    owner = _request("owner", PREFIX[:shared] + [-1] * 16, stub)
     _prefill(manager, stub, owner)
 
     # The first follower observes the junction (16 shared, dropped to 12) and
     # registers state there; the second one gets to resume from it.
-    follower = make_request(
-        "follower", PREFIX[:shared] + [-2] * 16, hash_block_size, sha256
-    )
+    follower = _request("follower", PREFIX[:shared] + [-2] * 16, stub)
     _prefill(manager, stub, follower)
 
     resume = shared // block_size * block_size - hash_block_size
-    hit = _sibling_hit(manager, shared, [-3] * 16, hash_block_size)
+    hit = _sibling_hit(manager, stub, shared, [-3] * 16)
     assert hit == resume, f"expected the resume point at {resume}, got {hit}"
 
 
@@ -205,14 +219,11 @@ def test_mamba_prefix_cache_drops_hash_block(mtp_draft_group_identified):
     prompt_a = PREFIX[:1600]
     hit_after_mtp_block_drop = len(prompt_a) - hash_block_size
 
-    _prefill(
-        manager,
-        _stub(manager, block_size, hash_block_size),
-        make_request("A", prompt_a, hash_block_size, sha256),
-    )
+    stub = _stub(manager, block_size, hash_block_size)
+    _prefill(manager, stub, _request("A", prompt_a, stub))
 
     tokens_after_a = [-1] * 500
-    hit = _sibling_hit(manager, len(prompt_a), tokens_after_a, hash_block_size)
+    hit = _sibling_hit(manager, stub, len(prompt_a), tokens_after_a)
     assert hit == hit_after_mtp_block_drop
 
 
@@ -232,7 +243,8 @@ def _armed_configs():
                 yield block_size, hash_block_size
 
 
-def test_scheduler_never_stops_where_the_manager_refuses():
+@pytest.mark.parametrize("lookahead", [False, True])
+def test_scheduler_never_stops_where_the_manager_refuses(lookahead):
     """A junction stop must always be a position the manager caches.
 
     "The scheduler schedules a number of tokens it thinks can be cached, but the
@@ -246,15 +258,13 @@ def test_scheduler_never_stops_where_the_manager_refuses():
             3 * block_size + hash_block_size + 1,
             4 * block_size - 1,
         ):
-            manager = _manager(block_size, hash_block_size)
+            manager = _manager(block_size, hash_block_size, lookahead=lookahead)
             stub = _stub(manager, block_size, hash_block_size)
             _orphaned_full_attention_tail(manager, stub, prompt_len)
 
             shared = prompt_len // hash_block_size * hash_block_size
             suffix = [-1] * (block_size + 3)
-            consumer = make_request(
-                "consumer", PREFIX[:shared] + suffix, hash_block_size, sha256
-            )
+            consumer = _request("consumer", PREFIX[:shared] + suffix, stub)
             junction = manager.get_computed_blocks(consumer)[2]
             if not junction:
                 continue
@@ -266,9 +276,7 @@ def test_scheduler_never_stops_where_the_manager_refuses():
                 skipped.append(case)
                 continue
             evaluated += 1
-            hit = _sibling_hit(
-                manager, shared, [-2] * (block_size + 3), hash_block_size
-            )
+            hit = _sibling_hit(manager, stub, shared, [-2] * (block_size + 3))
             if hit < junction:
                 refused.append(case + (hit,))
 
@@ -295,7 +303,7 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
     manager = _manager(block_size, hash_block_size)
     stub = _stub(manager, block_size, hash_block_size)
 
-    request = make_request("r", PREFIX[:2000], hash_block_size, sha256)
+    request = _request("r", PREFIX[:2000], stub)
     for _ in range(2000):
         request.append_output_token_ids(1)
     assert request.num_prompt_tokens == 2000 and request.num_tokens == 4000
@@ -329,9 +337,10 @@ def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
     )
 
 
+@pytest.mark.parametrize("lookahead", [False, True])
 @pytest.mark.parametrize("eagle_group", [None, 0])
 @pytest.mark.parametrize("num_prefill_lookahead", [0, 2, 33])
-def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead):
+def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead, lookahead):
     """Turning the flag on may add check-points; it must never move one away.
 
     Two axes decide whether the manager can honour the junction stop the
@@ -352,14 +361,13 @@ def sibling_hit(shared_prefix_checkpoint):
             shared_prefix_checkpoint=shared_prefix_checkpoint,
             eagle_group=eagle_group,
             num_prefill_lookahead=num_prefill_lookahead,
+            lookahead=lookahead,
         )
         stub = _stub(manager, block_size, hash_block_size)
         _orphaned_full_attention_tail(manager, stub, 2020)
-        consumer = make_request(
-            "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
-        )
+        consumer = _request("consumer", PREFIX[:2016] + [-1] * 584, stub)
         _prefill(manager, stub, consumer)
-        return _sibling_hit(manager, 2016, [-2] * 584, hash_block_size)
+        return _sibling_hit(manager, stub, 2016, [-2] * 584)
 
     on, off = sibling_hit(True), sibling_hit(False)
     assert on >= off, f"enabling the flag cut the sibling's hit from {off} to {on}"
@@ -372,15 +380,51 @@ def test_disabled_by_default():
     stub = _stub(manager, block_size, hash_block_size)
     _orphaned_full_attention_tail(manager, stub, 2020)
 
-    consumer = make_request(
-        "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
-    )
+    consumer = _request("consumer", PREFIX[:2016] + [-1] * 584, stub)
     junction = manager.get_computed_blocks(consumer)[2]
     assert junction, "expected a junction"
     _prefill(manager, stub, consumer)
 
-    assert _sibling_hit(manager, 2016, [-2] * 584, hash_block_size) < junction
+    assert _sibling_hit(manager, stub, 2016, [-2] * 584) < junction
 
-    request = make_request("r", PREFIX[:10000], hash_block_size, sha256)
+    request = _request("r", PREFIX[:10000], stub)
     request.shared_prefix_boundary = 2000
     assert Scheduler._mamba_block_aligned_split(stub, request, 8192) == 1536
+
+
+@pytest.mark.parametrize("lookahead", [False, True])
+def test_async_prompt_tail_state_is_copied_before_the_next_chunk(lookahead):
+    """Async scheduling allocates the next chunk before the previous chunk's
+    output arrives. The next chunk's forward overwrites the partial prompt-tail
+    Mamba state in place, so its copy must be queued by that allocation."""
+    block_size, hash_block_size = 16, 4
+    manager = _manager(block_size, hash_block_size, lookahead=lookahead)
+    stub = _stub(manager, block_size, hash_block_size)
+    mamba = manager.coordinator.single_type_managers[1]
+    request = _request("r", PREFIX[:30], stub)
+    manager.get_computed_blocks(request)
+
+    ends: list[int] = []
+    copied_sources: list[set[int]] = []
+    while request.num_computed_tokens < request.num_tokens:
+        start = request.num_computed_tokens
+        num_new = Scheduler._mamba_block_aligned_split(
+            stub, request, request.num_tokens - start, 0, 0
+        )
+        assert manager.allocate_slots(request, num_new, has_scheduled_reqs=False)
+        request.num_computed_tokens = start + num_new
+        ends.append(request.num_computed_tokens)
+        copies, retained = manager.take_kv_cache_block_copies()
+        copied_sources.append({copy.src_block_id for copy in copies})
+        if retained:
+            manager.block_pool.free_blocks(retained)
+        manager.new_step_starts()
+
+    tail = next(e for e in ends if e % block_size and e < request.num_prompt_tokens)
+    step = ends.index(tail)
+    assert step + 1 < len(ends), f"no chunk follows the tail: {ends}"
+    tail_block = mamba.req_to_blocks["r"][tail // block_size]
+    assert tail_block.block_id in copied_sources[step + 1], (
+        f"chunks {ends}: the state at {tail} lives in block {tail_block.block_id} "
+        f"but the next chunk's copies were {copied_sources[step + 1]}"
+    )

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

@@ -244,6 +244,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,
+        use_lookahead_block_hashes=False,
         hash_block_size=hash_block_size,
         dcp_world_size=dcp_world_size,
         scheduler_block_size=scheduler_block_size,
@@ -294,6 +295,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,
+        use_lookahead_block_hashes=False,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -335,6 +337,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
             long_prefill_token_threshold=long_prefill_threshold
         ),
         use_eagle_block_drop=False,
+        use_lookahead_block_hashes=False,
         hash_block_size=32,
         mamba_partial_cache_hit=False,
         mamba_shared_prefix_checkpoint=False,
@@ -1985,6 +1988,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,
+        use_lookahead_block_hashes=False,
         hash_block_size=hash_block_size,
         mamba_partial_cache_hit=True,
         mamba_shared_prefix_checkpoint=False,