Conversation
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: trueComment |
github-actions
Bot
added
documentation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| return value | ||
| if type(value) is int and 0 <= value < 2**64: | ||
| return value.to_bytes(8, "big") | ||
| raise ValueError("Invalid native block hash") |
There was a problem hiding this comment.
🟡 Nit: _hash rejects int block hashes outside [0, 2**64), but the legacy converter in the same bridge deliberately accepts the full int domain — to_int64 documents the engine hash as int | bytes and masks with & _U64_MASK, which handles negative values (vLLM's builtin prefix-caching hash algo is hash(), i.e. a signed 64-bit int, so roughly half of all hashes are negative there).
The consequence is not a per-event skip: convert_batch turns the ValueError into an INVALID group event, and GroupCache::apply then returns Err("invalid group event"), so apply_group_batch invalidates the worker's entire group view. If the installed engine ever emits int hashes in the signed range, every batch containing one is discarded and group affinity for that worker is permanently unavailable behind a single warn! — indistinguishable from "engine has no cache" at the routing layer.
Since the identity only needs to be deterministic and opaque, the safe form is to accept any int and encode it the way the legacy path does (e.g. (value & _U64_MASK).to_bytes(8, "big")), keeping the two converters' notions of a valid hash aligned. The -1 case at tests/test_vllm_kv_group_events.py:179 encodes the current stricter choice, so this needs a decision either way.
| } | ||
| GroupEvent::Invalid => return Err("invalid group event"), | ||
| GroupEvent::Remove { group_id, keys } => { | ||
| let group = self.groups.entry(*group_id).or_default(); |
There was a problem hiding this comment.
🟡 Nit: Remove uses entry(*group_id).or_default(), so a REMOVE for a group that has not yet been STOREd creates a Group with metadata: None. reusable_tokens then hits group.metadata.as_ref()? and returns None for the whole worker — every group's evidence is suppressed, not just the unknown one, until that group's first STORE arrives.
This is reachable on the normal path: after any invalidate() (sequence gap, stream reconnect, panic recovery) the cache is empty, and the next batches from a busy engine routinely carry BlockRemoved for evicted blocks before the next BlockStored for that group. The worker silently loses all cache affinity in the meantime.
Removing keys for an unknown group is a no-op anyway (reported is empty for it), so the entry does not need to be created — if let Some(group) = self.groups.get_mut(group_id) { for key in keys { group.reported.remove(key); } } keeps the metadata-less group out of the map.
| let limit = tokens.len().saturating_sub(1); | ||
| let mut prefix = 0; | ||
| let mut matching = Vec::new(); | ||
| for (i, &token) in tokens.iter().enumerate() { | ||
| prefix = extend_prefix(prefix, token); | ||
| if let Some(keys) = self.by_prefix.get(&Position { end: i + 1, prefix }) { | ||
| matching.extend(keys.iter().map(|key| (key, i + 1))); | ||
| } | ||
| } | ||
| let mut views = Vec::new(); |
There was a problem hiding this comment.
🟡 Nit: reusable_tokens walks every request token and computes one xxh3_128 plus one by_prefix lookup per token — but it is called once per candidate worker from the per-worker loop in overlap_candidates_with_groups (cache_aware.rs:1555). The prefix chain is worker-independent, so a 32K-token prompt over a 64-worker fleet redoes ~2M hash+lookup pairs per routing decision, in the path the neighbouring comment calls out as having "dominated routing CPU at scale".
Hoisting the chain would keep this linear in the prompt: compute Vec<Position> (or Vec<u128>) once per request in overlap_candidates_with_groups and pass it to a reusable_tokens(&[Position])-style entry point, leaving only the small matching/interval work per worker. The by_prefix lookups still have to be per-cache, but the hashing — the dominant cost — does not.
| let has_group_worker = self.kv_monitor.read().as_ref().is_some_and(|m| { | ||
| healthy_indices | ||
| .iter() | ||
| .any(|&idx| m.is_group_worker(workers[idx].url())) | ||
| }); | ||
| if self.has_event_indexer(model_id) || has_group_worker { |
There was a problem hiding this comment.
🟡 Nit: This scan runs on every request, for every healthy worker, even in fleets where no worker ever reports group events — each iteration is a DashMap shard lookup keyed by the worker URL String (so a full string hash + compare per worker), on top of the second identical scan added at line 1443 and the third per-worker group_caches.get() inside overlap_candidates_with_groups. That's three URL-keyed map probes per worker per request added to the hot path for a feature that is inactive in a legacy deployment.
Since group_caches is empty unless some worker has actually sent a group batch, an emptiness check short-circuits the common case:
| let has_group_worker = self.kv_monitor.read().as_ref().is_some_and(|m| { | |
| healthy_indices | |
| .iter() | |
| .any(|&idx| m.is_group_worker(workers[idx].url())) | |
| }); | |
| if self.has_event_indexer(model_id) || has_group_worker { | |
| let has_group_worker = self.kv_monitor.read().as_ref().is_some_and(|m| { | |
| m.has_any_group_worker() | |
| && healthy_indices | |
| .iter() | |
| .any(|&idx| m.is_group_worker(workers[idx].url())) | |
| }); | |
| if self.has_event_indexer(model_id) || has_group_worker { |
with fn has_any_group_worker(&self) -> bool { !self.group_caches.is_empty() } on the monitor.
| let monitor = guard.as_ref()?; | ||
| if info.cache_namespace.is_some() | ||
| && healthy_indices | ||
| .iter() | ||
| .any(|&idx| monitor.is_group_worker(workers[idx].url())) | ||
| { | ||
| return self.select_expected_wait(workers, healthy_indices, info); |
There was a problem hiding this comment.
🟡 Nit: The blast radius of this bail-out is the whole fleet, not the group workers. If a single healthy worker reports group events, every namespaced request loses cache affinity for all workers — including legacy workers whose PositionalIndexer evidence is unaffected by group reporting — and falls back to pure expected-wait. During a rolling upgrade (one upgraded worker among many legacy ones), namespaced traffic silently stops using KV-event routing fleet-wide.
It also reads as inconsistent with the decision documented 10 lines above at the dispatch site: "Event-driven mode re-hashes engine-reported blocks from their token ids on both sides, so a namespace marker on the request side alone would break every same-namespace match. It stays unpartitioned here". The group path re-hashes from token ids in exactly the same way, so it's not obvious why a namespace marker disqualifies it but not the legacy positional path.
If the intent is only "a namespaced request must not earn group affinity", scoping it per worker (skip the group score for group workers, keep the indexer score for the rest) preserves legacy behaviour instead of disabling the whole affinity stage.
| if batch.group_events_enabled || !batch.group_events.is_empty() { | ||
| Self::apply_group_batch(&group_caches, &worker_url, batch); | ||
| return; | ||
| } | ||
| if group_caches.contains_key(&worker_url) { | ||
| group_caches.insert(worker_url.clone(), GroupCache::default()); | ||
| warn!( | ||
| worker_url, | ||
| "Group stream lost its marker; using routing fallback" | ||
| ); | ||
| return; | ||
| } |
There was a problem hiding this comment.
🟡 Nit: This branch is sticky and noisy. The entry is re-inserted (not removed), so is_group_worker stays true forever: the worker keeps taking the group path with permanently empty evidence, gets no group score, and — because group_reusable_tokens returns Some(None) — is never scored by the legacy indexer either, even though the batches now arriving are legacy batches that learn_block_size/the positional indexer could use. The worker is excluded from cache affinity for the remaining life of the subscription.
It also warn!s once per batch, so a worker in this state emits a warning at engine event rate for as long as it runs.
Two milder options: remove the group_caches entry so the worker cleanly reverts to the legacy path (the comment's concern about a stale legacy tree is handled by the fact that the legacy tree is being fed by these same batches), or keep the entry but latch the warning so it fires only on the transition. Note the fall-through to learn_block_size is skipped in this branch, so the block size is never learned either.
Signed-off-by: ai-jz <ai-jz@users.noreply.github.com>
Signed-off-by: ai-jz <ai-jz@users.noreply.github.com>
Reuse the compressed Radix core for historical prefix identity and ordered endpoint matching. Prepare request content once, bound retained history, and retain group membership and common-position rules in GroupCache. Trim redundant tests and document the measured query, update and memory tradeoffs with a same-host Rust comparison chart. Signed-off-by: ai-jz <ai-jz@users.noreply.github.com>
496d1d6 to
d1b85ed
Compare
github-actions
Bot
added
dependencies
|
Updated to
The capacity-reset review finding is addressed. Other design discussion remains open; shared-index integration and combined validation still follow #2500/#2437/#2438. Current-head hosted checks are starting; no merge-readiness claim yet. |
| let matching: Vec<_> = endpoints | ||
| .iter() | ||
| .flat_map(|endpoint| { | ||
| self.by_prefix[endpoint] |
There was a problem hiding this comment.
🟡 Nit: self.by_prefix[endpoint] is a panicking index on the request path (inside a DashMap read guard, under the router's worker-selection loop). It is safe today only because of a non-local invariant: learn_context registers the endpoint in index at line 215, but the matching by_prefix insert is four lines and two return Err/one continue later — the gap is only benign because every learn error unwinds through apply_batch's invalidate(). Any future error path that returns without invalidating (or any partial apply) leaves a context registered in the tree with no by_prefix row, and this line panics rather than degrading to fallback.
Cheap to make structurally safe:
| self.by_prefix[endpoint] | |
| self.by_prefix | |
| .get(endpoint) | |
| .into_iter() | |
| .flatten() |
| if kind == "sliding_window" { | ||
| intervals.sort_unstable(); | ||
| } | ||
| candidates.extend( | ||
| intervals | ||
| .iter() | ||
| .map(|&(_, end)| end) | ||
| .filter(|&end| end <= limit), | ||
| ); | ||
| if matches!(kind.as_str(), "full_attention" | "mla_attention") { | ||
| // Full attention needs continuous coverage from zero. Compute | ||
| // its bound once, rather than again for every candidate. | ||
| let mut covered = 0; | ||
| for &(start, end) in &intervals { | ||
| // Endpoints arrive in depth order. A later span may | ||
| // bridge a gap; earlier skipped spans cannot extend it. | ||
| if start <= covered { | ||
| covered = end; | ||
| } | ||
| } |
There was a problem hiding this comment.
🟡 Nit: making sort_unstable() conditional on sliding_window turns "endpoints arrive in ascending depth" into a load-bearing, unenforced invariant for two consumers, and only one of them says so.
- Here (
full_attention/mla_attention):covered = endis only equivalent tocovered.max(end)because ends are non-decreasing. Out of order,coveredcan shrink. - Below at line 319 (
mamba):rangesis built straight from the same unsortedintervalsand then consumed by the monotonecursorincandidates.retain(line 342), which silently drops matches unlessrangesis sorted by end. That branch carries no comment about the ordering it depends on.
The order does hold today (matching_contexts → endpoints in depth order → flat_map over by_prefix preserves it, non-strictly), but it is three layers away from either use site. Two low-cost hardenings, neither of which gives up the win — intervals is small, so an unconditional sort_unstable() costs little, or keep the fast path and assert it:
debug_assert!(intervals.is_sorted_by_key(|&(_, end)| end));plus covered = covered.max(end); here, which is free.
| let mut endpoints = Vec::new(); | ||
| self.index.matching_contexts( | ||
| &request.contents, | ||
| &mut OverlapScratch::default(), | ||
| &mut endpoints, | ||
| ); |
There was a problem hiding this comment.
🟡 Nit: GroupRequest hoists the u32 → u64 conversion out of the per-worker loop, but the scratch does not come with it. OverlapScratch::default() is built fresh on every call, and reusable_tokens is called once per candidate worker per request (cache_aware.rs:1557), so a 100-worker fleet allocates and grows 100 scratches — plus the endpoints Vec — per request. That is exactly the per-query heap traffic radix_tree's caller-owned-scratch design exists to avoid ("no per-query heap once warm — the alloc gate holds this", chain.rs:588).
Since matching_contexts only reads scratch.segments and the buffer is not held across calls, a thread_local! reused scratch (or threading one through GroupRequest, which is already the per-request scratchpad this push introduced) keeps it warm across the whole worker loop.
| run: | | ||
| cd bindings/python | ||
| python3 -m pip install pytest pytest-cov pytest-xdist | ||
| python3 -m pip install pytest pytest-asyncio pytest-cov pytest-xdist |
There was a problem hiding this comment.
🟡 Nit: pytest-asyncio is installed by the bindings/python step but the tests that need it live in grpc_servicer/tests (the @pytest.mark.asyncio cases this PR adds in test_vllm_kv_group_events.py, plus the pre-existing ones in test_vllm_kv_events_stream.py / test_tokenspeed_kv_events_stream.py). It works only because both steps share one job's interpreter.
That coupling is invisible and load-bearing: if the grpc_servicer step is ever split into its own job, reordered, or the bindings/python step's deps are trimmed, the async tests do not fail loudly — under strict mode with no plugin they are skipped with a warning, so the group-event subscription and corrupt-frame tests would go green while never running. Installing it in the step that consumes it (or adding it to grpc_servicer's dev extras) makes the dependency survive that refactor.
| } | ||
| } | ||
| assert_eq!((total, native_subset), (159, 63)); | ||
| } |
There was a problem hiding this comment.
🟡 Nit: this push deletes nine_native_hash_vectors_are_opaque_identities_not_recomputed_hashes, which was the only test that drove the recorded vLLM hash vectors through the cache — and the doc's headline contract claim ("Native hashes remain opaque byte identities, including unsigned 64-bit exports") no longer has a dedicated regression. It also strands data: hash_vectors (9 entries, with hash/next_hash/token_ids) is still in tests/fixtures/vllm-group-cache.json and is now referenced by nothing in the repo.
The removal is understandable — the swap from the local xxh3 prefix chain to learn_context means keys are no longer hashed at all — but that arguably makes the opacity assertion cheaper to keep, not less relevant: it is now the check that a raw engine hash is never reinterpreted as content. Either re-point the vectors at the new path or drop hash_vectors from the fixture so the corpus does not carry dead scenarios.
| impl RadixTree { | ||
| /// Append from a previously learned parent, without claiming residency. | ||
| /// Only the final endpoint is registered; intermediate content is shared. | ||
| pub fn learn_context( |
There was a problem hiding this comment.
🟡 Nit: learn_context is a mutator of the shared chain structure — it calls the same place_content that store uses, so it creates roots, extends chains in place, and forks children — yet it sits entirely outside the harness that this crate's correctness story rests on. tests/common/mod.rs's Op is Store | Remove | Clear, and fuzz_differential.rs's COp adds only Truncate | Retire | Recreate | StaleProbe | Query. No generated run ever learns a context, so no run ever exercises a store landing in a tree whose shape was partly built by contexts, or a remove whose GC walk has to respect the new third pin condition (!cd.contexts.is_empty(), chain.rs:1041).
That is the same class the harness has already caught twice — both hazards documented in place_block (the in-batch move freeing the cursor chain, the per-holder key canonicity divergence) were chaos-fuzz findings on exactly this GC/placement interaction. tests/contexts.rs covers two hand-written interleavings and calls audit(), which is the right instinct, but two fixed sequences is thin next to what the generator reaches.
Adding a COp::Learn { parent_slot, contents } looks cheap: the reference model needs no change (contexts claim no membership, so overlap and holder sets must be unaffected — that invariance is itself the assertion worth making), and audit() already validates context lineage coherence and the orphan/GC rule.
Present query and state-memory scaling for the current Radix matcher without experimental implementation labels. Signed-off-by: ai-jz <ai-jz@users.noreply.github.com>
|
|
||
| The bridge automatically selects this path when the installed vLLM event schema supplies group metadata. There is no additional environment variable, manual group list or request restriction. Existing KV-event publisher and gateway cache-aware routing configuration still apply. Older vLLM event schemas keep the legacy converter. Updated protocol bindings, bridge and gateway are needed for group-aware scores; older gateways ignore the additive group fields and receive no legacy cache evidence from these batches. | ||
|
|
||
|  |
There was a problem hiding this comment.
🟡 Nit: this push replaces images/kv-cache-group-routing-scaling.png (287 KB), but unlike its sibling on this line, that figure is referenced by nothing in the repo — grep -rn kv-cache-group-routing-scaling . returns no hits. It renders only in the PR body, via a raw.githubusercontent.com URL pinned to bd67585.
So the scaling evidence lands in the tree as an orphaned binary: a reader of vllm-group-events.md never sees it, and the "Complexity and scaling" material the PR body builds on it (linear-in-W fan-out, the 1,048,576 × 1,000 stress point) has no in-repo home at all — this doc's "Producer and transport limits" section covers correctness limits but says nothing about query cost. Once the PR is merged the body is the only thing carrying it, and body text is not versioned with the code it describes.
Either embed it under a short scaling subsection here, alongside the O(W·T) caveat, or drop the file and keep the figure purely as PR-discussion material rather than committing an unreferenced asset.
Signed-off-by: ai-jz <ai-jz@users.noreply.github.com>
| Self { | ||
| index: RadixTree::new(Config { | ||
| // Token positions use u32; this is not a cache capacity policy. | ||
| max_chain_len: u32::MAX, |
There was a problem hiding this comment.
🔴 Important: this push removes every capacity bound on a GroupCache — MAX_ENTRIES (100k), MAX_CONTEXT_TOKENS (1.6M), the per-event tokens.len() guard in the Store arm, and now max_chain_len — and replaces them with reclamation that is purely event-driven. Nothing else in the gateway backstops it: KvEventMonitor::prune_all (kv_event_monitor.rs:124) only walks indexers, so group_caches is never pruned, and the only remaining reset is a sequence gap / disconnect / malformed input.
That makes reclamation entirely contingent on the producer emitting a matching REMOVE (or CLEAR) for every key it STOREs, which this PR's own description documents as not holding:
- Phantom keys are never reclaimed. Under
kv_cache_report_mode="full",emit_cached_block_eventsadvertises block hashes for blocks that are not in the pool (the PR body's second producer issue, vLLM #52371 — "blocks 0 and 2 are absent … nevertheless reports all six logical hashes"). The report has dense-looking lengths, so it passes thekeys.len() * block_size == tokens.len()check at line 162 and islearned. A block that was never cached can never be evicted, so noBlockRemovedwill ever name its hash →needed()stays true →retain_contextis never released → that chain and all its ancestors are pinned for the life of the subscription. Every such report adds monotonically. - Lost removals are undetectable. The PR states "silent restarts or missing final events are not fully detectable". Previously the 100k/1.6M ceiling was the safety valve that reset a diverged view; now a view that drifts simply grows.
- Dead tails are never compacted.
new_childforks rather than splits, so a chain keeps its fullcontentsas long as any one context on it is pinned —retained_contents()is explicitly "not a total memory bound" per the README change in this push.
The growth is unbounded and per worker, in a DashMap<String, GroupCache> on the gateway, so the failure mode is gateway OOM rather than degraded routing.
The previous round's objection was to the cliff (reset-on-overflow), not to having a bound at all — the GC added here fixes the cliff, and keeping a much looser ceiling on top of it costs little. Either re-add a coarse check in reclaim (e.g. index.retained_contents() + reported entries over a generous limit → invalidate(), now a genuine last resort rather than the primary mechanism), or extend prune_all to bound group caches the way it bounds indexers. Separately, max_chain_len: u32::MAX disables the learn_context length guard at context.rs:52 entirely, so a single wire-supplied tokens vector is now the only limit on one chain's length — worth keeping a sane cap there even if the aggregate bound moves elsewhere.
Note the PR body's "History bound" paragraph still describes the deleted 1,600,000 / 100,000 limits; vllm-group-events.md was updated but the description was not.
| /// Release an endpoint. Whole chains are reclaimed once no retained | ||
| /// endpoint, holder membership or child needs them. Drain retired contexts | ||
| /// after mutations to remove corresponding caller-owned identity records. | ||
| pub fn release_context(&mut self, context: PrefixContext) -> bool { |
There was a problem hiding this comment.
🟡 Nit: release_context turns contexts into a GC trigger, which is a materially bigger surface than the read-mostly learn_context of the previous push — and it is still entirely outside the chaos/differential harness.
Concretely, this push makes cd.context_pins load-bearing in three places that the fuzzer exercises hard for holders but never with contexts: the free condition (chain.rs:1046), the orphan/GC-leak assertion (chain.rs:1388), and the parent-walk in maybe_gc_chain_pinned that now retires contexts on ancestors it frees (chain.rs:1065). fuzz_differential.rs's COp (line 421) is still Store | Remove | Clear | Truncate | Retire | Recreate | StaleProbe | Query — no generated run ever interleaves a context release with a holder remove or retire that frees the same chain from the other direction, or with a store that reuses the freed slot from free_chains.
Slot reuse is the specific hazard: tests/contexts.rs now covers it once, by hand, for learn_context(None, &[7, 8]) where the lineage differs. The case it does not cover is a freed slot re-allocated with an identical path, where (chain, position, lineage) compares equal to a retired handle — set_context_retained accepts it, and GroupCache::reclaim's drain_retired_contexts loop would then evict the live by_prefix/positions row for the new context. It's unreachable today only because GroupCache never learns between release_context and the drain in the same reclaim; that's a caller-side ordering invariant this crate doesn't state or enforce, and a generator that emits Learn/Release/Store in arbitrary order is exactly what would pin it down.
The reference model needs no new state — contexts claim no membership, so overlap and holder sets must be unchanged by a Learn/Release pair, and that invariance plus audit() (which this push already extended with the pin-count check at chain.rs:1195) is most of the assertion.
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
The shared index now answers per holder with the covered runs of the request and the publisher's opaque lane metadata, and a worker with several independently evicted caches (full attention, sliding window, recurrent state) is published as one holder per lane. The gateway is where those answers become one number per worker: the largest prefix position every lane accepts at that position, which is what vLLM's hybrid coordinator computes natively. policies/reuse.rs holds the rules: full and MLA attention need coverage from 0, a sliding window needs cdiv(W - 1, block) contiguous engine blocks (at least one) before the candidate, a recurrent lane needs a saved state exactly there (any block end inside a run), and candidates walk down from the request bound aligned to the least common multiple of the lanes' engine blocks. A holder without lane metadata keeps its depth from 0, so single-cache workers score exactly as before. A lane the gateway cannot read disqualifies its worker rather than guessing. Checked against the native-lookup fixture recorded in #2497 (vLLM find_longest_cache_hit expectations, 160 cases): all 125 whole-block cases agree; the 35 partial-hash-hit cases are not representable in the event feed, whose partial blocks the converter drops, and are skipped with that reason. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Targeted cache-selection test and AgentX replayFYI @slin1237. The targeted Qwen test measures the benefit in the cache-selection scenario this PR addresses, even with original vLLM. AgentX workload replay measures the overall impact: Qwen results were mixed; Gemma improved with vLLM PR44488 and PR55092, while the same SMG candidate without those fixes regressed.
Tested build: both PR2497 columns include the same local support for the proposed vLLM event format, beyond the published head. This is evidence for the proposal, not qualification of the published head unchanged. Gemma's observed benefit relies on the repaired event semantics; the fixes were tested together, so their individual contributions are not isolated. Test setup and interpretationOriginal: SMG Four workers per AgentX run: Qwen/Qwen3.6-35B-A3B, two H100s per worker; google/gemma-4-12B-it, one H100 per worker. Both used BF16 and a 262,144-token context limit. Each version replayed the same corpus/settings from fresh engine caches; 3,600 seconds of dispatch plus about 900 seconds of grace, one run per version. Parallel sessions are agent trajectories, not worker count. Cache hit rate and mean TTFT use common successful turns: 9,926 / 6,525 / 1,544 for Qwen-64 / Qwen-256 / Gemma-64. Cache hits are vLLM's cached prompt tokens at first prefill, before any later recomputation; throughput covers the full run. Qwen-256 had four timeouts per version. Gemma had 12 / 12 / 11 timeouts, 2 / 1 / 1 cancellations, and warmup parsing failures; its mean latency gain was concentrated in longer waits. These limits preclude a general speedup claim. The targeted Qwen test used two workers and eight target requests per version, each with 57,128 input and 64 output tokens. It tests choosing usable cached state instead of a longer prefix that cannot resume the request; it does not establish how often that case occurs in real traffic. No original-SMG + patched-vLLM comparison or shared-index integration was tested. |
1 participant
Description
This is a draft for design feedback on group-aware cache routing. The implementation and tests make the proposal concrete; feedback on the cache semantics, event contract and scope is welcome before merge readiness.
Problem
The legacy KV-event index models a contiguous cached prefix. Hybrid models need different cache data at the same prefix position: full-attention keys/values, a sliding-window history, or a saved recurrent state. A KV cache group is vLLM's set of model layers sharing one cache block table; several distinct groups can have the same cache kind.
A single
min(group_maxima)is insufficient. A Mamba state saved at a later position cannot restore every earlier position, and a sliding-window cache may have gaps. Each candidate position must satisfy all reported groups at that same position.Solution: select the largest common position
The matching principle in this PR is consistent with vLLM's native hybrid prefix-cache lookup: select the largest prefix position supported by every participating KV cache group. In
HybridKVCacheCoordinator.find_longest_cache_hit, each cache manager accepts the candidate length or returns a shorter hit, and the coordinator reconciles the results until it reaches a common hit. The native implementation uses an iterative fixed-point search; SMG computes full-attention coverage bounds and intersects ordered sliding-window ranges and Mamba checkpoints against the same common-position criterion.The per-kind reuse conditions below correspond to vLLM's native
find_longest_cache_hitimplementations. Links are pinned to the engine revision used in validation:full_attention,mla_attentionFullAttentionManager: scan the cached prefix from the beginning.sliding_windowWincludes the current token, so up toW - 1previous tokens are needed.SlidingWindowManager: search backward for a sufficient contiguous history.mambaMambaManager.MambaManager: search backward for an eligible saved state.The SMG matcher applies these conditions to resolved received reports and returns the largest endpoint supported by every KV cache group reported so far. If none qualifies, existing routing fallback applies. The same rule works when a model uses only one kind or a different combination; it does not require every model to contain all three kinds.
The selection principle is shared; matching live-engine token counts additionally requires the same participating groups, cache state and boundary eligibility. vLLM queries its actual cache and applies scheduler alignment, partial-hit and speculative-decoding rules. SMG has the received-event view and the three conditions above. Those information/eligibility boundaries, including the observed Gemma sparse-report limitation, remain explicit below.
In the illustrative 32,768-token request, the saved state at 29,568 lacks the full-attention prefix. The state at 25,344 has that prefix, but its required SWA history crosses a gap. 16,896 is the largest position accepted by all three groups, inside a retained SWA span of 11,648 tokens. The earlier saved states cannot improve that result. The combined example is hypothetical; the model cards show actual tested configurations, including Gemma's 1,024-token window.
Design principle: reuse semantics are independent of retention policy
For the supported cache kinds, reuse is defined by the state required to resume computation. SMG selects the largest position supported by every observed KV cache group. Retention and eviction change which positions qualify, while the matching rule stays the same. This principle follows from the computation requirements and is consistent with vLLM's native lookup.
Full attention needs its earlier prefix; SWA needs its local history; a recurrent layer needs the corresponding saved state.
Eviction withdraws cached blocks or states. It changes which positions satisfy each condition, but does not require a new matching rule for each eviction pattern. Retention density and earlier requests can likewise change SWA spans and Mamba points without changing the intersection rule. This small separation between reported state and reuse conditions is the source of the design's simplicity and expected maintainability. It does not, by itself, prove that every implementation or producer report is correct.
We also checked retention directly. vLLM's
prefix_cache_retention_intervalretains replay/shared-prefix boundaries at0, adds periodic checkpoints for positive intervals, and permits dense retention atNone. Previously stored blocks and subsequent eviction determine what remains available. Thus SWA can retain longer spans or several fragments; Mamba checkpoints follow engine-selected positions. The extracted SWA/Mamba retention methods were identical between tested revision56d001faand inspected revisionc3ccc0e9, and their unchanged methods reproduced the figure's checkpoint selection. This is concrete compatibility evidence across those revisions, not a guarantee about every past or future version.Received events are the interface
The bridge preserves native group identifiers, opaque hashes, parents and token spans. Groups are discovered when their events arrive; there is no preconfigured model roster. Dense, anchored reports establish positions. Sparse reports keep their membership and token span, but missing offsets are not guessed; another report may resolve a shared hash, and unresolved reports remain pending while live evidence depends on them.
Repeated STORE reports are idempotent evidence. REMOVE withdraws the named group's key, CLEAR empties reported memberships, and known sequence gaps, disconnects or malformed input reset the observation session. Unsupported or unresolved evidence yields no group-affinity score and uses the existing fallback; it does not reject inference requests.
The names
full_attention,mla_attention,sliding_windowandmambaare unchanged vLLM event-kind values. MLA uses the same prefix-coverage condition as full attention. The bridge mapsgroup_idxto the candidate'sKvGroupEvent.group_id,kv_cache_spec_kindtokind, andkv_cache_spec_sliding_windowtosliding_window. These are field mappings for the same concepts. Other vLLM cache kinds are outside this PR's scope.The installed vLLM event schema selects this converter automatically through the existing KV-events path. There is no additional opt-in flag or configured group list. Schemas without the required group metadata retain the legacy converter, including #2495's alignment guard. Group-capable schemas may omit the later-added
localityandownershipfields; absence now retains the older local-GPU semantics, while explicit remote locality and non-None ownership remain filtered. This repairs a compatibility regression in this PR; it requires no vLLM upgrade or producer change.Relationship to the shared prefix-cache index work
The group matching rules complement the shared-index work. The merged #2436 supplies the Radix core reused here; #2500 adds recency-based eviction, #2437 supplies the shared index service, and #2438 connects gateways to it. Those changes address how cache knowledge is maintained and shared across gateways. This PR addresses how reports from different KV cache groups jointly support one reuse position.
The implementation has two integration boundaries. The small Radix context API retains prefix identity independently of live membership and now releases unused history through the existing whole-chain collector. Its interaction with #2500's recency-based reclamation still needs joint validation. The group matcher currently lives in the gateway's
KvEventMonitor; enabling #2438's remote-index mode skips those local subscriptions, and #2437's inspected bridge consumes legacyeventsrather thangroup_events.The shared query currently carries block-level content hashes and returns matched blocks, while this matcher uses token-level paths and group-specific boundaries. Integration therefore needs to align the event/query contract and preserve group membership, historical parents, pending reconstruction and observation resets through replication and state recovery. Reusing the same Radix crate does not provide that integration automatically.
The plan is to rebase onto the relevant shared-index changes after they land and wire the group support into the shared ingestion/query path. The matching rules and current implementation can be reviewed now, including the two boundaries above. Shared-index integration and combined performance evaluation remain follow-up work. #2439's independent simulation harness may support that evaluation; its existing results do not establish group-aware behavior. The result remains a hint derived from received events, subject to the producer and freshness limits below.
Alternatives considered: worker-side native prefix-cache lookup
A reasonable alternative is to extend the existing SMG gRPC integration with a query that executes vLLM's native prefix-cache lookup inside each candidate worker's EngineCore and returns its currently reusable token count for the request. This could avoid sparse-offset and full-report-membership defects in KV events and reduce duplicated matching semantics in SMG.
The existing gRPC library can carry the request, but does not itself own EngineCore's cache. In the pinned engine, AsyncLLM uses a multiprocess EngineCore client; its utility RPC mechanism is a possible transport into that process. A concrete query entry point is still needed.
Two implementation conditions deserve separate discussion:
KVCacheManager.get_computed_blockscan emitBlockStoredevents in full-report mode, so calling it unchanged is not automatically side-effect-free. A query path must audit/isolate such effects; that is a vLLM interface/implementation dependency.We have no latency or throughput measurements for this alternative and have not rejected it experimentally. With a small candidate-worker set, it may be the better trade-off. The draft invites feedback on that choice; implementing the query or a reservation mechanism is outside this patch.
Changes
One group matcher, a thin Python/protobuf bridge, and monitor/policy integration provide the complete routing behavior. The matcher now reuses
smg-radix-treefor exact historical prefix identity and traversal. Its small context API appends from a stable parent and returns only completely matched, explicitly registered endpoints.GroupCacheowns native-key memberships, sparse reconstruction, stream lifecycle and group reuse conditions.Request tokens are prepared once across worker queries. Each worker retains its own tree; this does not yet provide shared-index integration or cross-worker traversal. Cleanup runs after each event batch: live descendants keep their ancestry, unused pending reports are discarded, and unneeded contexts are released through Radix's existing whole-chain collector. Native identities are removed when their backing chain is collected.
The patch has 2,731 changed text lines across 29 files, including two image files (2,591 additions and 140 deletions). The PNGs are counted as files, not text lines. This exceeds the repository's small-PR guideline, so feedback on reviewing this end-to-end change together is explicitly requested.
kv-indexis bumped to 1.5.0 with a dependency onsmg-radix-tree0.1.1; the release workflow publishes Radix first.Test Plan
Confidence comes from three separate layers:
A minimal model-free regression demonstrates why independent maxima cannot be minimized:
cargo test -p kv-index sparse_joint_resume_is_24_even_though_independent_maxima_are_32_and_40The groups' independent maxima are 32 and 40; their latest common supported position is 24. Additional tests cover later group discovery, sparse positions learned across groups, pending-parent resolution and reclamation, shared native identities, same-batch eviction/append, repeated STORE/removal, CLEAR, sequence gaps and malformed final frames. A Python-to-protobuf-to-Rust fixture changes worker selection after removal of a Mamba checkpoint.
Commands and recorded CPU validation
September 11 lifecycle update: the frozen candidate passed the complete serial Rust workspace suite on a CPU pod: 5,539 passed, 0 failed, 43 ignored (including 4 passing doctests). Nightly formatting, workspace/all-target Clippy using the documented alternative without OpenCV, and applicable file hooks passed. All 1,417 source/symlink paths matched before and after validation. The 23 GroupCache and 3 context tests include the new lifecycle cases and two independent 1,048,576-token prefixes on one worker; these focused results are included in the workspace total. Python/protocol code is unchanged by this update; its earlier validation is recorded below.
A local optimized Rust diagnostic exercised one worker with four live prompts, 65,536 or 1,048,576 tokens each, with independent prefixes or a shared 1,024-token prefix. Across 256 STORE rounds (four initial fills and 252 replacements), retained history plateaued at 4,194,304 and 5,238,784 token units in the million-token cases; the shared-prefix case includes an existing dead tail. All four cases emptied their logical history on CLEAR, and 1,024 removed unresolved reports left zero pending tokens. Native event replay confirmed the documented lost-parent miss and preserved hits when full-attention evidence or a same-batch parent remained available. These are finite CPU lifecycle checks, not a total memory bound or shared-index/GPU validation.
For the previous published snapshot, the complete serialized workspace suite passed 5,535 tests, 0 failed, 43 ignored across 129 result blocks. Nightly formatting, workspace/all-target Clippy using CONTRIBUTING.md's documented alternative without OpenCV, source-built Python bindings/native import, Ruff and applicable file hooks passed. The available Python suite passed 178, with two skips for unavailable NumPy-dependent TokenSpeed loads and vLLM;
test_tokenspeed_dp_rank_pin.pywas explicitly excluded because its NumPy/Torch/TokenSpeed runtime dependencies were unavailable.The original default-parallel workspace run stopped on one unchanged HTTP metrics test: the process-global interner grew by 102 against its slack limit of 100. That same test binary passed the test in isolation, and the complete suite subsequently passed with one test thread. The original failure is retained as a validation caveat; no assertion was weakened or test removed to get this result.
The previous snapshot was validated on main
3d36721ewith all 1,416 source paths matching before and after the checks. A final packaging-only change bumpedkv-indexto 1.5.0 in three Cargo files; all Rust/Python bytes remained identical. The final locked/offline dependency metadata and package suite passed separately (254 unit tests plus 1 doc test), with all 1,416 final source paths matching again. Those 255 package results are not added to the full-suite total. That snapshot's GroupCache source SHA-256 isa016f3456e6bcae62cd171feb97aa6be4ec58c59ea2725bd0fd5cc3545bb7cdb.Protocol bindings were regenerated through their owning build. The new core API, group matcher, Python/protobuf boundary and monitor/policy tests ran in the validation above; all 159 fixture snapshots remain. A separate mutation check removed the sequence-width guard in memory: the old doubly-invalid input still passed, while the corrected valid-payload input failed as intended. Source hashes were unchanged by the mutation experiment.
The ten earlier STORE/REMOVE compatibility cases explicitly distinguish absent attributes from attributes set to None or LOCAL. Before repair, six failed and four passed; all ten now pass. A separate native-schema audit exercised real msgspec encode/decode with unchanged event-class definitions from five pinned vLLM revisions: the legacy-converter boundary, earliest array-form group events without either optional field, locality-only map events, and modern events with ownership. It also checked CLEAR, hash representations, default/tuple extra keys, empty stores and explicit remote/owned filtering. No additional missing accessed field was found in those representative schemas. The exact version/image of the reported failing deployment remains unconfirmed; this is bounded schema compatibility evidence, not a deployed integration rerun.
The model runs below were recorded for the earlier candidate; they are historical correctness/routing evidence, not GPU validation of the current Radix implementation.
Pinned engine: vLLM
56d001faf0f53c72fcedbbdd77e5418f68fe7494.full_attention+ three GDNmambagroups; all block size 1,056event_hitpassed.event_hitpassed.Both models passed 23 received-state replay queries each. All tested HTTP pairs returned 200 and identical output IDs. Runs used BF16, text-only requests, 2,048-token prefill chunks, a 16K context limit and sequential single-worker checks; the Gemma prefill setting was not reduced to hide the limitation. These are correctness/routing checks, not a performance benchmark. The long Gemma cases remain failed affinity checks despite successful inference.
Rust CPU tests varied cached prompt length from 1K to 1M tokens on one worker, and worker count from 1 to 1,000 at 64K tokens per worker, using one cached prompt per worker with full attention and one Mamba checkpoint. Query-time growth was consistent with the expected prefix traversal and sorting costs. Checking every worker still becomes expensive at large scale. These measurements and the complexity analysis apply to the current independent per-worker matcher; they do not predict performance after shared-index integration or end-to-end serving latency.
Limits, rollout and follow-up
There are two distinct correctness questions: whether SMG correctly applies received events, and whether those events fully and promptly describe the engine. This proposal addresses the first. The score is resolved received cache evidence, not a reservation or the engine's exact current reusable-token count.
The investigation identified two different producer issues; they should not be conflated with router matching correctness:
block_hashes, whiletoken_idsstill spans the logical range. Ordinal pairing gives incorrect positions; unresolved offsets also explain the long Gemma misses.skipped_parent_block_hash,skipped_token_idsandskipped_extra_keys; it has moved beyond the initialblock_offsetsproposal. Consume the agreed upstream format and rerun captured-trace/GPU integration after it lands. This PR does not guess missing positions or adopt the unmerged extension.emit_cached_block_eventsnevertheless reports all six logical hashes. The event has dense-looking lengths, so #2495's alignment check cannot detect this; the consumer can overestimate cache availability.The second finding is established by executing the extracted native helper against a constructed sparse cache, not by the Qwen/Gemma HTTP runs. Neither SMG change claims to make inaccurate full reports exact. The intended follow-up is to repair and validate the producer contract upstream, then strengthen integration coverage here.
The upstream minimal regression uses ordinary Mamba
align, block size 4 and an 8-token producer prefix, with no partial hash matching. The token-8 state exists and the token-4 state does not, yet a consumer'skv_cache_report_mode="full"lookup publishes the token-4 hash. The test runs throughKVCacheManageron CPU and marks the desired membership assertion as strictxfail.Completeness and freshness also matter: unreported groups or stale removals can overestimate current engine state. The live rank-zero subscriber does not recover publisher replay, and silent restarts or missing final events are not fully detectable. Native scheduler alignment, partial-hit rules and speculative decoding eligibility are not reconstructed from these events. CPU/remote/offload reports, nonzero ranks and namespaced requests do not establish supported plain-token group affinity.
History lifecycle: there is no additional aggregate capacity limit or capacity-triggered reset. Unused history follows Radix's existing whole-chain reclamation; dead tails within live chains may remain, so this is not a total memory bound. If all supporting history has been collected, a later child waits for a usable parent report. This can lose affinity—for example, a Mamba checkpoint whose parent was removed in an earlier event batch—but never turns an unknown parent into a shorter rooted prefix.
Scope and fallback: only local-GPU, rank-zero reports establish supported group affinity; engine-specific reuse eligibility is not reconstructed. If any eligible worker uses group mode, namespaced requests use load routing for the whole candidate set. A subscription that loses its group marker stays unavailable until valid group batches or re-registration restore it.
Deploy the updated bridge and gateway together. Older gateways ignore the additive fields and receive no group-affinity evidence from those batches. Before broader rollout, canary cold/warm requests and stream-disconnect fallback on the intended model; roll back bridge/gateway versions if behavior regresses.
Feedback requested