[RFC] Prefix caching on paged KV cache path

[ricky-chaoju]

Problem

The paged KV path recomputes KV from scratch on every prefill, even when requests share the same prefix (system prompts, multi-turn history). This hurts TTFT.

What the upstream scheduler already does

The upstream vLLM scheduler handles block hashing, cache hit detection, and block reuse:

1. Request.update_block_hashes() computes chained block hashes on request creation
2. kv_cache_manager.get_computed_blocks(request) finds the longest prefix cache hit
3. NewRequestData.num_computed_tokens carries the hit count to the model runner
4. block_ids includes both cached and newly allocated blocks
5. enable_prefix_caching defaults to True, nothing in vllm-metal forces it off

No hash/LRU code needed on our side — the scheduler handles all of that.

What vllm-metal's paged path does today

The model runner reads num_computed_tokens (model_runner.py:1949), but only uses it for intermediate chunk detection. For complete prefills, it always processes all tokens from position 0.

Why this is blocked on kernel_v2

Our prefill uses mx.fast.scaled_dot_product_attention(Q, K, V) — it takes Q/K/V as input tensors and can't read cached K/V from paged blocks. If we skip the prefix tokens, SDPA won't have the prefix K/V and attention output will be wrong.

On GPU, flash attention reads cached K/V directly from paged blocks via block tables. We need kernel_v2 (varlen attention, #148 Stage 2) to do the same on Metal.

Once kernel_v2 lands:

1. When num_computed_tokens > 0, only prefill token_ids[num_computed_tokens:]
2. Add offset to prepare_prefill() so slot_mapping starts from the right position
3. Set RoPE offset = num_computed_tokens
4. Kernel_v2 reads cached prefix K/V from paged blocks via block tables

RFC @WindChimeRan @LxYuan0420

[WindChimeRan]

@ricky-chaoju

"implement hash_block_tokens() as the first step": I think that's already done by the upstream vllm scheduler.

Is it possible to get the cache hit info from the vllm scheduler and then wire it to our paged kv?

Don't wanna re-create old wheels and cause possible code divergences from the upstream.

[ricky-chaoju]

@ricky-chaoju

"implement hash_block_tokens() as the first step": I think that's already done by the upstream vllm scheduler.

Is it possible to get the cache hit info from the vllm scheduler and then wire it to our paged kv?

Don't wanna re-create old wheels and cause possible code divergences from the upstream.

You're right — the scheduler already handles it. num_computed_tokens is already passed to the model runner; we just ignore it for fully completed prefills.

The main blocker right now is that our prefill SDPA path can't read cached K/V from paged blocks the way flash attention does on GPU. That capability will come with kernel_v2.

I've updated the issue to reflect this.

[WindChimeRan]

LGTM @ricky-chaoju Thanks! I think kernel_v2 will be ready next week or so.

[WindChimeRan]

@ricky-chaoju

Can we prepare a benchmark script for this before kernel_v2 lands? That way we have a clean baseline to compare against.

The previous prefix cache PRs (#80, #82) were merged without latency/throughput benchmarks. It'd be good to not repeat that here.

Proposal:

• Use vllm bench serve with a shared-prefix dataset (e.g., N requests sharing a long system prompt + short unique user queries). Upstream vLLM has dataset support in vllm.benchmarks.datasets worth checking for reference (and relevant vllm issue & PR).
• Key metric is TTFT — that's what prefix caching directly improves.
• put the scripts in tools, then in the comments we can have what dataset to download via hf download, and vllm server command .etc

[WindChimeRan]

prefix_repetition from vllm bench serve might be a good fit.

[ricky-chaoju]

Opened #182

[LxYuan0420]

Sorry, was busy with work and missed this.

Agree that we should rely on the cache-hit / reuse information already coming from upstream vLLM, rather than building a second metal-local prefix-cache mechanism.

My suggestion is to keep the next step narrow so we do not waste work:

1. docs-only TTFT baseline first, i.e. just document the exact vllm serve + vllm bench serve commands and use that as the baseline, without adding a new wrapper script yet
2. wire the paged prefill path so it only prefills the uncached suffix
3. add correctness tests for full-prefill vs cached-prefix equivalence
4. rerun the same benchmark and compare before/after TTFT