[2/3] kv: cross-request prefix caching

[Gaurav-Shah05]

What does this PR do?

Part of #210, stacked on #279 . Reuses KV pages across requests that start the same way: chained SHA-256 page keys under a fingerprint root, computed on the preprocess worke. On by default for any model that declares a token stream for a KV resource (Bagel's text LLM and Orpheus here), inert for everything else.

A hit

Engine._prepare_inputs asks the node's resources how much of the request is cached, takes the minimum, tells each resource, and trims that many tokens before the step is declared; bucket selection, attention, positions and commit see a shorter request. The sampler is told too, so its repetition-penalty mask still covers the skipped tokens. Keys come only from the preprocess worker; anything a client sends under those names is discarded.

Storage

The arena's own pages under a second owner, no copies. A page is offered to the index by the commit that fills it on the keyed walk, prompt or generated. An allocation or reload that comes up short evicts leaves from the index before failing, skipping any a live request still reads.

What is cached

Bagel text chat and Orpheus prompts, generated tokens included. The model names the walk that writes the keyed text, so that walk alone is probed and its commits alone extend the chain; another walk writing the same stream ends it, as does a write the keys do not cover.

Not cached: anything after an image, since there is no digest stream yet, which the follow up adds; and Bagel's guided image prefill, whose inputs the default cut refuses because requires_cfg travels with the request rather than the walk. The follow up does not change that second one, which needs a Bagel split_inputs that cuts the cfg_img span and the fork. A windowed stream keys only its protected prefix. TP is #281.

Declaring a model

A model opts in with prefix_key_streams(), naming the tensor that holds the token ids and the walks that write and decode it, plus checkpoint_path(), since the weights are part of the root and a model that names none needs a salt to use the cache at all. It is about four lines each for Bagel and Orpheus, and a model that declares nothing is untouched. A declaration the node cannot honour is refused at load. I have updated docs/adding_models.rst w/ the contract.

How was it tested?

178 new tests. pytest test/modular -q at d172d293 gives 1 failed, 1112 passed, 20 skipped (the failure is pre-existing on main). Cache off is bit-identical to main at 424a112c within one allocation. A hit changes the output without changing what it says: of six two-turn chats, three answered exactly as the uncached run and three paraphrased it, and on a text-to-image about 1% of pixels move by more than 32 of 255 with the scene unchanged, so prefix_cache=False is there when an output has to be reproducible. Turn-two TTFT went from 44.9 ms to 26.4 ms. An image-first edit followed by the same prompt as text-to-image matches nothing and its image is identical to the uncached run, and 32 same-prefix requests on a 256-page pool all complete.

Checklist

• ruff check . passes
• Added or updated tests / docs where relevant

[merceod]

Three issues I found:

1. The implementation has wrong-position keying. Keys assume text_inputs[0] is stream position 0. A Bagel edit sent as image,text writes the VAE and ViT walks into main first, and their pages get indexed under the text's keys. A later T2I with cfg off and the same prompt matched "2 of 2 pages" and produced a garbage image. Please try to reproduce yourself too and see the output yourself because I might be mistaken. Potential fix: only let commits the engine offered to the probe advance the chain, or declare the keyed walk on PrefixStream.

2. Noticed an unconverted hit after a refused batch admit. The stream keeps pages with stored_len 0, the retry goes again, admit skips the conversion, the trimmed step runs from position 0, and its pages get filed under the prefix keys so later requests read garbage. A one-line guard might fix it. Also, your PR #281 fixes it by probing at ingest.

3. The CUDA-graph index buffers overflow. The attention wrappers size paged_kv_indices by max_num_pages which is a bound that is broken by sharing. A small-pool hitFlashInfer's plan check and failed two whole batches. Also, the default pool has the same problem once batch size times pages per request passes 2048.

Note (should-fix): the leaves heap grows one entry per and nothing drains it until the pool fills (10k lookups left 10k entries).

Adding the following part later (a few details to go with the points above):

On the first one, to reproduce it the edit has to be the first writer of that prompt. If a t2i or a text-first edit with the same text ran earlier, the correct text pages are already in the index and "first writer wins" refuses the VAE pages, so nothing looks wrong. It should be a fresh prompt of at least 128 tokens, send the edit with input_modalities=image,text and a picture, then a t2i with the same prompt and cfg_text_scale=1.0, cfg_img_scale=1.0 (with cfg on the t2i is never tested). The log then says matched N of its N pages and the image is garbage. Unit version: ingest keys for a 40 token prompt on a 16 page manager, commit a 50 token step on main first, and the index has two entries. The check at manager.py line 453 only looks backward, it passes because 0 is not past covered_len.

On the second, the guard is the eligibility test in resolve_cached_prefix (manager.py line 399): add or stream.page_indices there, or convert anyway at admit (manager.py line744) and release the private pages.

On the third, it's the _paged_kv_indices_buf = torch.zeros(max_num_pages, ...) in wrappers.py at 84 and 229. The error is flashinfer's "The size of indices should be less than or equal to the allocated buffer", raised from the decode preplan. Sizing these as capture batch size times ceil(max_seq_len / page_size) should fix it, it's only int32s.

On the heap, it's the push in lookup (prefix_index.py line 46), one entry per hit on a leaf. Lazy re-push keeps it bounded so don't push on lookup, and in _pop_leaf when an entry's stamp is stale but the page is still a leaf, push it back with the current stamp instead of dropping it.

[merceod]

Every lookup that ends on a leaf pushes a new heap entry, and nothing pops them until the pool is full. A system prompt gets re-stamped on every request, so on a server whose pool never fills this list grows by one tuple per hit (10000 lookups in a test left 10000 entries), and the first eviction then walks through all of them under the lock. Lazy re-push would keep it bounded (so don't push here, and in _pop_leaf when an entry's stamp is stale but the page is still a leaf, push it back with the current stamp instead of dropping it).

[Gaurav-Shah05]

Good catch! I dropped the push in lookup, and _pop_leaf now puts a stale entry back when the page is still a leaf rather than losing it. Ten thousand reads of one leaf leave one entry

[merceod]

This seems fragile. It works for Bagel because requires_cfg is False for chat and True for t2i, but a model that puts 0 or an empty container here would and one that puts anything else would never, and nothing in the model says so. Letting the model name the keyed walk (say, on PrefixStream) makes this clear, and it also becomes the store-side signal for the position-zero problem in _index_filled_pages.

[Gaurav-Shah05]

The walk is named on PrefixStream now, and that's the half that fixes the indexing. I kept this check beside it though, because requires_cfg is per request: the same walk is sliceable for a chat and not for a guided t2i, so the model can't say up front which one it is.

[merceod]

Small layering thing, the engine picking out KVManager by type. If enable_prefix_cache were a Resource with a no-op default the engine wouldn't need to know which resource kind keeps an index.

[Gaurav-Shah05]

Fair, the engine shouldn't know which resource kind keeps an index. It's a Resource method with a no-op default now.

[merceod]

If the stream already held pages here (a refused admit leaves some), they're overwritten and never released. An assert that it's empty or a release would do.

[Gaurav-Shah05]

Right, those were leaked. It releases them first now, in the same commit as the admit conversion.

[merceod]

Nothing checks the page isn't already indexed under another key. Shouldn't happen today (I think), but if it ever does the old _by_key entry just hangs, and an assert here would be good.

[Gaurav-Shah05]

Added, with the page and both keys in the message.

[merceod]

process_prompt still gets the client's original kwargs here, prefix_keys included if they sent some. The stripped model_kwargs is right above and passing that instead would keep the "only this worker keys a prompt" idea alive all the way through.

[Gaurav-Shah05]

You're right, that one slipped through. It gets the stripped copy now.

[Gaurav-Shah05]

@merceod I've pushed the fixes and updated the docs.

Position-zero keying: the model names the walk that writes the keyed text, and a commit from any other walk ends that stream's chain. Index buffers are sized per row, so the 256-page pool that failed 29 of 32 requests completes all of them.

Two I did differently from what you suggested. The buffer bound is num_rows x (max_num_pages - 1), since BAGEL gives a whole image block one position and a multi-image request can hold more KV than max_seq_len. And I kept the opaque-inputs check beside the walk name, since requires_cfg travels with the request, so the same walk is sliceable for a chat and not for a guided t2i.

Let me know if these make sense to you and if you'd like me to change anything.

[merceod]

Approving this since everything lgtm, but had one important note.

@Gaurav-Shah05

I ran a 96 page pool with 32 requests sharing a 59 page prefix. They all were allowed/admitted, since, with the prefix cached, each prefill only needed itstail page. Then every decode step wanted one more page, nothing was evictable because every page in the index was live, and with no cpu offload configured the worker just held all of them at the 50ms backoff until the clients timed out. It recovered fine once they did. The hold path is main brnach's, but a prefill no longer reserves the prompt's pages, so it no longer bounds concurrency by what the pool can hold. The sizing rule from your RFC fixes it I think, so maybe a small comment in serving.rst for now, and a headroom check at admission or a preemption path when a decode step finds nothing to evict would be a good follow up. Also that path logs a WARNING per retry per batch, I got 19k lines in six minutes, so rate limiting it would help.

[Gaurav-Shah05]

Good catch! It's worse than a sizing note: prompt reservation was doing admission control by accident. Before this PR a request reserved its whole prompt, so concurrency was bounded by pool over prompt pages; with a hit the prefill reserves a tail page and that bound is gone. The cache is on by default, so a pool sized for the old behaviour over-admits on upgrade.

Worth knowing, the victim path already exists: _try_offload_cold_request would have offloaded one of 32 requests, but cpu_offload_pages defaults to 0 and only three shipped configs set it, none of them Bagel's or Orpheus's. Configuring host pages is the mitigation available today, and the note I added in the docs says so.

The new commits add the sizing rule in serving.rst, the rate limit on that warning, and two lines that would have made this visible: one at load naming which nodes are keyed and warning when a cached node has no host pages, one the first time requests reach a keyed node with no keys.

There were additional fixes I added as well. Reading the manager for #281 I found admit_retrieve clearing the chain's pending tokens even when the published record is the rank's own, so generated tokens stop being keyed after the first step wherever KV is published. And the last one moves the chain state off CacheStream into its own object, so it's more readable.

@merceod Do you want prefix_cache defaulting off until the headroom check or preemption exists?

I'm running a suite of benchmarks and will report the numbers here.

[Gaurav-Shah05]

@merceod Okay I have also filed the follow-ups. #307 covers the stall you found, #308 covers TP, and #309 covers guided prefill. None of them block this PR. The stall has a workaround and a warning at load, TP stays off above one rank, and admission needs the benchmark numbers which I will add soon.

[merceod]

Good catch! It's worse than a sizing note: prompt reservation was doing admission control by accident. Before this PR a request reserved its whole prompt, so concurrency was bounded by pool over prompt pages; with a hit the prefill reserves a tail page and that bound is gone. The cache is on by default, so a pool sized for the old behaviour over-admits on upgrade.

Worth knowing, the victim path already exists: _try_offload_cold_request would have offloaded one of 32 requests, but cpu_offload_pages defaults to 0 and only three shipped configs set it, none of them Bagel's or Orpheus's. Configuring host pages is the mitigation available today, and the note I added in the docs says so.

The new commits add the sizing rule in serving.rst, the rate limit on that warning, and two lines that would have made this visible: one at load naming which nodes are keyed and warning when a cached node has no host pages, one the first time requests reach a keyed node with no keys.

There were additional fixes I added as well. Reading the manager for #281 I found admit_retrieve clearing the chain's pending tokens even when the published record is the rank's own, so generated tokens stop being keyed after the first step wherever KV is published. And the last one moves the chain state off CacheStream into its own object, so it's more readable.

@merceod Do you want prefix_cache defaulting off until the headroom check or preemption exists?

I'm running a suite of benchmarks and will report the numbers here.

@Gaurav-Shah05 Keep it on I guess. The hold only happens once the pool is full of live pages, and main doesn't deal with that that either, it just gets there later because a prompt reserves its pages up front. What changed is the threshold, not the failure.

Two things worth naming in the docs sentence and the load warning: max_concurrent_requests in the YAML bounds admission directly today (pool over expected pages per request is a fine setting), and cpu_offload_pages gives the worker a victim. With those two named, #307 can take its time. Merging.

[stephen-dwq]

fuzzing of enable_prefix_cache, ingest_request, resolve_cached_prefix, apply_cached_prefix, admit, plan, commit, extend_prefix_chain, reset_request, remove_request comes clean

lgtm