[KV Connector][Spec Decode] Lookahead block hashes for KV connectors

tests/v1/core/test_kv_cache_utils.py

@@ -1096,7 +1096,12 @@ def test_request_block_hasher_incremental_append_with_multiple_mm_features(
         pytest.param(
             {"disable_lookahead_block_hashes": True}, False, id="disabled-by-config"
         ),
-        pytest.param({"kv_connector": True}, False, id="kv-connector"),
+        pytest.param(
+            {"enable_prefix_caching": False, "connector_support": True},
+            True,
+            id="supporting-connector",
+        ),
+        pytest.param({"connector_support": False}, False, id="unsupported-connector"),
         pytest.param({"hisparse": True}, False, id="hisparse"),
     ],
 )
@@ -1108,6 +1113,7 @@ def test_lookahead_block_hashing_enabled(overrides: dict[str, Any], expected: bo
         "disable_lookahead_block_hashes",
         overrides.get("disable_lookahead_block_hashes", False),
     )
+    connector_support = overrides.get("connector_support")
     vllm_config = cast(
         VllmConfig,
         SimpleNamespace(
@@ -1118,14 +1124,19 @@ def test_lookahead_block_hashing_enabled(overrides: dict[str, Any], expected: bo
             num_prefill_lookahead_tokens=overrides.get(
                 "num_prefill_lookahead_tokens", 1
             ),
-            kv_transfer_config=object() if overrides.get("kv_connector") else None,
+            kv_transfer_config=None if connector_support is None else object(),
             attention_config=SimpleNamespace(
                 hisparse_config=object() if overrides.get("hisparse") else None
             ),
         ),
     )
+    connector = (
+        None
+        if connector_support is None
+        else SimpleNamespace(supports_lookahead_block_hashes=connector_support)
+    )
 
-    assert is_lookahead_block_hashing_enabled(vllm_config) == expected
+    assert is_lookahead_block_hashing_enabled(vllm_config, connector) == expected
 
 
 @pytest.mark.parametrize("hash_fn", [sha256, sha256_cbor])

tests/v1/core/test_scheduler.py

@@ -420,6 +420,52 @@ def test_finished_request_publishes_block_ending_at_last_sampled_token():
     assert num_tokens == 4
 
 
+def test_connector_finish_includes_partial_eagle_block(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    block_size = 16
+    init_none_hash(sha256)
+    scheduler = create_scheduler(
+        enable_prefix_caching=True,
+        block_size=block_size,
+        max_num_batched_tokens=64,
+        use_kv_connector=mock_kv(matched_tokens=0, is_async=False),
+    )
+    _enable_lookahead_block_hashes(scheduler)
+    request = Request(
+        request_id="request",
+        prompt_token_ids=list(range(33)),
+        sampling_params=SamplingParams(max_tokens=2),
+        pooling_params=None,
+        block_hasher=get_request_block_hasher(block_size, sha256, True),
+    )
+    scheduler.add_request(request)
+    scheduler_output = scheduler.schedule()
+    scheduler.update_from_output(
+        scheduler_output,
+        ModelRunnerOutput(
+            req_ids=[request.request_id],
+            req_id_to_index={request.request_id: 0},
+            sampled_token_ids=[[33]],
+        ),
+    )
+    assert request.num_computed_tokens == 33
+
+    captured_block_ids = None
+
+    def request_finished(_request, block_ids):
+        nonlocal captured_block_ids
+        captured_block_ids = block_ids
+        return False, None
+
+    assert scheduler.connector is not None
+    monkeypatch.setattr(scheduler.connector, "request_finished", request_finished)
+    scheduler._connector_finished(request)
+
+    assert captured_block_ids is not None
+    assert len(captured_block_ids) == 3
+
+
 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)]

tests/v1/e2e/spec_decode/mtp/qwen3_5/test_mtp.py

@@ -2,9 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import pytest
+import torch
 
+from tests.models.utils import check_logprobs_close
 from tests.utils import single_gpu_only
-from vllm import SamplingParams
+from vllm import LLM, SamplingParams, TokensPrompt
+from vllm.distributed import cleanup_dist_env_and_memory
 
 from .._correctness import check_mtp_correctness
 
@@ -37,3 +40,84 @@ def test_mtp_correctness(
         expected_accuracy_threshold,
         vllm_runner,
     )
+
+
+@single_gpu_only
+def test_qwen3_5_mtp_prefix_cache_reuses_last_safe_block(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    """Offloading uses the lookahead-proven MTP boundary without accuracy loss."""
+    aligned_page_size = 544
+    prompt_len = 2 * aligned_page_size + 1
+    sampling_params = SamplingParams(
+        temperature=0,
+        max_tokens=32,
+        logprobs=5,
+    )
+
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_MLA_DISABLE", "1")
+        llm = LLM(
+            model="Qwen/Qwen3.5-0.8B-Base",
+            trust_remote_code=True,
+            enable_prefix_caching=True,
+            max_model_len=2048,
+            speculative_config={
+                "method": "mtp",
+                "num_speculative_tokens": 1,
+                "max_model_len": 2048,
+            },
+            kv_transfer_config={
+                "kv_connector": "OffloadingConnector",
+                "kv_role": "kv_both",
+                "kv_connector_extra_config": {
+                    "cpu_bytes_to_use": 512 << 20,
+                },
+            },
+            limit_mm_per_prompt={"image": 0, "video": 0},
+        )
+        assert llm.llm_engine.vllm_config.cache_config.prefix_match_unit is None
+
+        tokenizer = llm.get_tokenizer()
+        source = (
+            "The following context is intentionally repeated to exercise prefix "
+            "caching. " * 300
+        ) + "Explain why deterministic inference should be reproducible."
+        prompt_token_ids = tokenizer.encode(source)[-prompt_len:]
+        assert len(prompt_token_ids) == prompt_len
+        prompt = TokensPrompt(prompt_token_ids=prompt_token_ids)
+
+        cold_output = llm.generate([prompt], sampling_params)[0]
+        warm_output = llm.generate([prompt], sampling_params)[0]
+
+        assert cold_output.num_cached_tokens == 0
+        max_safe_hit = 2 * aligned_page_size
+        # An unconditional one-page EAGLE drop would report only 544 tokens.
+        assert warm_output.num_cached_tokens == max_safe_hit
+
+        cold_completion = cold_output.outputs[0]
+        warm_completion = warm_output.outputs[0]
+        assert cold_completion.token_ids == warm_completion.token_ids
+        check_logprobs_close(
+            outputs_0_lst=[
+                (
+                    list(cold_completion.token_ids),
+                    cold_completion.text,
+                    cold_completion.logprobs,
+                )
+            ],
+            outputs_1_lst=[
+                (
+                    list(warm_completion.token_ids),
+                    warm_completion.text,
+                    warm_completion.logprobs,
+                )
+            ],
+            name_0="cold_mtp",
+            name_1="warm_mtp",
+            always_check_logprobs=True,
+        )
+
+        del llm
+        torch.accelerator.empty_cache()
+        cleanup_dist_env_and_memory()

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

@@ -327,6 +327,7 @@ run_test_for_device() {
 
   # Run test
   echo "Running spec decode acceptance test (kv_buffer_device=${kv_device}, backend=${ATTENTION_BACKEND})..."
+  PREFILL_PORT=${PREFILL_PORTS[0]} \
   DECODE_PORT=${DECODE_PORTS[0]} \
   SERVER_HOST=$SERVER_HOST \
   TEST_MODEL=$MODEL_NAME \

tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py

@@ -11,23 +11,27 @@
 
 Environment variables (set by spec_decode_acceptance_test.sh):
     TEST_MODEL   - target model name
+    PREFILL_PORT - port of the prefill vLLM server (for /metrics)
     DECODE_PORT  - port of the decode vLLM server (for /metrics)
     SD_METHOD    - "eagle3" (default) or "mtp"
 """
 
 import os
+import time
 from dataclasses import dataclass, field
 from types import SimpleNamespace
 from urllib.request import urlopen
 
 import openai
+import pytest
 import regex as re
 from transformers import AutoTokenizer
 
 from vllm.benchmarks.datasets import get_samples
 
 SERVER_HOST = os.environ.get("SERVER_HOST", "127.0.0.1")
 PROXY_BASE_URL = f"http://{SERVER_HOST}:8192/v1"
+PREFILL_PORT = os.environ.get("PREFILL_PORT", "8100")
 DECODE_PORT = os.environ.get("DECODE_PORT", "8200")
 MODEL_NAME = os.environ.get("TEST_MODEL", "meta-llama/Llama-3.1-8B-Instruct")
 SD_METHOD = os.environ.get("SD_METHOD", "eagle3").lower()
@@ -109,14 +113,85 @@ def _get_mt_bench_prompts() -> list[str]:
     return [sample.prompt for sample in samples]
 
 
-def _fetch_metric(metric_name: str) -> float:
-    """Fetch a single counter metric from the decode server's /metrics."""
-    url = f"http://{SERVER_HOST}:{DECODE_PORT}/metrics"
+def _fetch_metric(metric_name: str, port: str = DECODE_PORT) -> float:
+    """Fetch a counter, treating an unmaterialized zero series as zero."""
+    url = f"http://{SERVER_HOST}:{port}/metrics"
     body = urlopen(url).read().decode()
+    names = (
+        (metric_name,)
+        if metric_name.endswith("_total")
+        else (
+            metric_name,
+            metric_name + "_total",
+        )
+    )
     for line in body.split("\n"):
-        if line.startswith(metric_name + "{") or line.startswith(metric_name + " "):
+        if any(
+            line.startswith(name + "{") or line.startswith(name + " ") for name in names
+        ):
             return float(line.rsplit(" ", 1)[-1])
-    raise ValueError(f"Metric {metric_name} not found in decode /metrics")
+    return 0.0
+
+
+def _wait_for_metric_delta(
+    metric_name: str,
+    port: str,
+    baseline: float,
+    expected: float,
+    timeout: float = 10,
+) -> float:
+    deadline = time.monotonic() + timeout
+    while True:
+        delta = _fetch_metric(metric_name, port) - baseline
+        if delta >= expected or time.monotonic() >= deadline:
+            return delta
+        time.sleep(0.1)
+
+
+def test_mtp_pd_reuses_last_safe_prefix_block():
+    """Repeated MTP requests hit the lookahead-proven boundary on both nodes."""
+    if SD_METHOD != "mtp":
+        pytest.skip("lookahead-hash MTP prefix caching test")
+
+    aligned_page_size = 544
+    prompt_len = 2 * aligned_page_size + 1
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    source = (
+        "The following context is intentionally repeated to exercise prefix "
+        "caching across disaggregated prefill and decode nodes. "
+    ) * 300
+    prompt_token_ids = tokenizer.encode(source, add_special_tokens=False)[-prompt_len:]
+    assert len(prompt_token_ids) == prompt_len
+
+    client = openai.OpenAI(api_key="EMPTY", base_url=PROXY_BASE_URL)
+
+    def complete():
+        return client.completions.create(
+            model=MODEL_NAME,
+            prompt=prompt_token_ids,
+            max_tokens=32,
+            temperature=0.0,
+            seed=42,
+            extra_body={"add_special_tokens": False},
+        )
+
+    cold = complete()
+    prefill_hits = _fetch_metric("vllm:prefix_cache_hits", PREFILL_PORT)
+    decode_hits = _fetch_metric("vllm:prefix_cache_hits", DECODE_PORT)
+
+    warm = complete()
+    max_safe_hit = 2 * aligned_page_size
+    prefill_warm_hits = _wait_for_metric_delta(
+        "vllm:prefix_cache_hits", PREFILL_PORT, prefill_hits, max_safe_hit
+    )
+    decode_warm_hits = _wait_for_metric_delta(
+        "vllm:prefix_cache_hits", DECODE_PORT, decode_hits, max_safe_hit
+    )
+
+    # Legacy EAGLE dropping would leave only one 544-token page reusable.
+    assert prefill_warm_hits == max_safe_hit
+    assert decode_warm_hits == max_safe_hit
+    assert cold.choices[0].text == warm.choices[0].text
 
 
 def _fetch_per_position_acceptance() -> dict[int, float]:
@@ -150,6 +225,10 @@ def test_spec_decode_acceptance_length():
         f"Expected {DEFAULT_NUM_PROMPTS} prompts, got {len(prompts)}"
     )
 
+    drafts_before = _fetch_metric("vllm:spec_decode_num_drafts_total")
+    accepted_before = _fetch_metric("vllm:spec_decode_num_accepted_tokens_total")
+    per_position_before = _fetch_per_position_acceptance()
+
     client = openai.OpenAI(api_key="EMPTY", base_url=PROXY_BASE_URL)
     for i, prompt in enumerate(prompts):
         resp = client.completions.create(
@@ -168,8 +247,10 @@ def test_spec_decode_acceptance_length():
             print(f"  [{i}] {prompt[:60]}... -> {text}...")
 
     # ── Extract metrics from decode server ────────────────────────────
-    n_drafts = _fetch_metric("vllm:spec_decode_num_drafts_total")
-    n_accepted = _fetch_metric("vllm:spec_decode_num_accepted_tokens_total")
+    n_drafts = _fetch_metric("vllm:spec_decode_num_drafts_total") - drafts_before
+    n_accepted = (
+        _fetch_metric("vllm:spec_decode_num_accepted_tokens_total") - accepted_before
+    )
 
     assert n_drafts > 0, "No spec-decode drafts were generated"
 
@@ -193,7 +274,11 @@ def test_spec_decode_acceptance_length():
 
     # ── Assert per-position acceptance (EAGLE3) ───────────────────────
     if config.expected_acceptance_lengths_per_pos:
-        per_pos_counts = _fetch_per_position_acceptance()
+        per_position_after = _fetch_per_position_acceptance()
+        per_pos_counts = {
+            pos: count - per_position_before.get(pos, 0)
+            for pos, count in per_position_after.items()
+        }
         per_pos_rates = [
             per_pos_counts.get(i, 0) / n_drafts
             for i in range(len(config.expected_acceptance_lengths_per_pos))