[Paged KV] Enable prefix caching on the unified paged path

tests/test_paged_deterministic.py

@@ -118,12 +118,25 @@ def _set_env():
 
 
 @pytest.fixture(scope="module")
-def vllm_outputs():
-    """Run vLLM offline inference once for all prompts.
-
-    Uses max_num_seqs=1 to avoid batch-invariance non-determinism on Metal.
+def _paged_llm():
+    """Single LLM shared across both test classes.
+
+    A second module-scope LLM cannot fit alongside this one — Metal
+    memory held by the first LLM is not released by Python gc, so the
+    second initialisation hits ``kv_budget=0`` and aborts.  Both the
+    cache-off baseline test and the prefix-cache-hit e2e test reuse
+    this instance.
+
+    ``enable_prefix_caching=True`` is safe for the baseline test because
+    the first ``generate`` pass walks an empty cache and produces output
+    identical to a cache-off run.  Subsequent passes hit the cache.
     """
-    llm = LLM(model=MODEL_NAME, max_model_len=512, max_num_seqs=1)
+    llm = LLM(
+        model=MODEL_NAME,
+        max_model_len=512,
+        max_num_seqs=1,
+        enable_prefix_caching=True,
+    )
 
     # Verify paged KV path is active when requested
     if os.environ.get("VLLM_METAL_USE_PAGED_ATTENTION", "0") == "1":
@@ -138,8 +151,14 @@ def vllm_outputs():
         attn = runner.model.model.layers[0].self_attn
         assert isinstance(attn, MetalKernelPagedAttentionWrapper)
 
+    return llm
+
+
+@pytest.fixture(scope="module")
+def vllm_outputs(_paged_llm):
+    """First pass — empty cache, output equals cache-off baseline."""
     sp = SamplingParams(temperature=0, max_tokens=MAX_TOKENS)
-    outputs = llm.generate(PROMPTS, sp)
+    outputs = _paged_llm.generate(PROMPTS, sp)
     return {o.prompt: o for o in outputs}
 
 
@@ -177,3 +196,63 @@ def test_generate_matches_golden(self, vllm_outputs, prompt):
             f"Expected (mlx): {mlx_expected}\n"
             f"Expected (pgd): {paged_expected}"
         )
+
+
+@pytest.fixture(scope="module")
+def vllm_prefix_cached_outputs(_paged_llm, vllm_outputs):
+    """Second pass — cache populated by the ``vllm_outputs`` priming pass.
+
+    ``vllm_outputs`` is declared as a dependency only so pytest orders the
+    priming pass first; its return value is intentionally unused.
+
+    Upstream invariant: a ``generate`` call repeating the same prompt
+    triggers the prefix-cache lookup (block hashes match), so the second
+    pass walks the ``start_pos > 0`` path inside the model_runner
+    (issue #182).  We do not assert ``num_computed_tokens > 0`` directly
+    — that contract belongs to upstream's own tests.  This fixture's job
+    is to confirm vllm-metal's cache-hit code path produces correct output.
+    """
+    del vllm_outputs  # ordering-only dependency
+    if os.environ.get("VLLM_METAL_USE_PAGED_ATTENTION", "0") != "1":
+        pytest.skip("Prefix caching e2e test only meaningful on the paged path")
+    sp = SamplingParams(temperature=0, max_tokens=MAX_TOKENS)
+    outputs = _paged_llm.generate(PROMPTS, sp)
+    return {o.prompt: o for o in outputs}
+
+
+class TestPagedPrefixCacheCorrectness:
+    """End-to-end correctness of paged prefix caching (issue #182)."""
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize("prompt", PROMPTS)
+    def test_prefix_cached_matches_golden(self, vllm_prefix_cached_outputs, prompt):
+        output = vllm_prefix_cached_outputs[prompt]
+        token_ids = list(output.outputs[0].token_ids)
+        text = output.outputs[0].text
+
+        mlx_expected = GOLDEN_MLX[prompt]
+        paged_expected = GOLDEN_PAGED[prompt]
+
+        mlx_match = token_ids == mlx_expected
+        paged_match = token_ids == paged_expected
+        print(
+            f"VLLM_METAL_USE_PAGED_ATTENTION: {os.environ.get('VLLM_METAL_USE_PAGED_ATTENTION')}"
+        )
+        print(f"\n  prompt: {prompt!r}  (prefix-cached)")
+        print(f"  output: {text!r}")
+        print(f"  ids:    {token_ids}")
+        if mlx_match:
+            print("  result: MATCHED mlx-cache golden")
+        elif paged_match:
+            print("  result: MATCHED paged-cache golden")
+        else:
+            print("  result: NO MATCH")
+            print(f"  expected (mlx):   {mlx_expected}")
+            print(f"  expected (paged): {paged_expected}")
+
+        assert mlx_match or paged_match, (
+            f"Output for {prompt!r} (prefix-cached) matched neither golden set.\n"
+            f"Got:            {token_ids}\n"
+            f"Expected (mlx): {mlx_expected}\n"
+            f"Expected (pgd): {paged_expected}"
+        )

tools/README.md

@@ -65,16 +65,18 @@ python -m tools.benchmark.attention_benchmark --mode varlen --q-lens 1,4,16,64 -
 
 ## Prefix Caching Benchmark
 
-Measures TTFT with shared-prefix workloads using `prefix_repetition` dataset.
-Establishes a baseline before prefix caching is implemented (#159).
+Measures TTFT / TPOT / E2EL with shared-prefix workloads using the
+upstream `prefix_repetition` dataset.  Compare cache-off baseline vs
+cache-on by toggling `--enable-prefix-caching` / `--no-enable-prefix-caching`.
 
 **1. Start the server:**
 
 ```bash
 # Adjust MEMORY_FRACTION based on available RAM (lower if OOM).
 VLLM_METAL_USE_PAGED_ATTENTION=1 VLLM_METAL_MEMORY_FRACTION=0.7 \
   vllm serve Qwen/Qwen3-0.6B \
-    --port 8000 --max-model-len 2048 --max-num-seqs 8
+    --port 8000 --max-model-len 2048 --max-num-seqs 8 \
+    --enable-prefix-caching
 ```
 
 **2. Run the benchmark:**
@@ -93,8 +95,8 @@ vllm bench serve \
   --request-rate inf \
   --percentile-metrics ttft,tpot,e2el \
   --metric-percentiles 50,99 \
-  --save-result --label baseline
+  --save-result --label cache-on
 ```
 
-Key metric is **TTFT** — with prefix caching enabled, requests sharing
-the same prefix should show lower TTFT on cache hits.
+For a cache-off baseline, restart the server with
+`--no-enable-prefix-caching` and re-run with `--label baseline`.

vllm_metal/platform.py

@@ -275,15 +275,6 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                     scheduler_config.max_num_batched_tokens,
                 )
 
-        if config.use_paged_attention and getattr(
-            cache_config, "enable_prefix_caching", False
-        ):
-            # The unified paged path does not yet safely support vLLM core
-            # prefix-cache hits for new requests. Disable the feature at the
-            # platform layer until that path is fully supported.
-            cache_config.enable_prefix_caching = False
-            logger.info("Metal: disabled prefix caching")
-
         # Configure cache — ensure block_size is at least the Metal kernel
         # minimum.  With chunked prefill enabled, upstream may default to
         # block_size=1 for fine-grained scheduling, but our Metal paged