fix(router): index only the attention KV-cache groups in kv-aware routing - #98
Open
jiejingzhangamd wants to merge 4 commits into
Open
fix(router): index only the attention KV-cache groups in kv-aware routing#98jiejingzhangamd wants to merge 4 commits into
jiejingzhangamd wants to merge 4 commits into
Conversation
kv-aware routing reported cache_hits=0 on every decision against Kimi-K3 on
two mixed workers, and pinned all traffic to one of them. Not a degraded hit
rate -- half the fleet idle.
vLLM emits one BlockStored PER KV-CACHE GROUP. On a hybrid model only the
attention group pairs one hash with one block; Kimi-K3 has 24 MLA layers and
69 KDA layers, and the KDA groups run prefix caching in "align" mode, which
nulls all but one block per step. Those events go out with token_ids spanning
the whole chunk and a single surviving hash, and nothing on the wire says
which block that hash covers. Captured off a live worker (block_size 768,
--max-num-batched-tokens 4096, so chunked prefill splits at 3840):
kv_cache_spec_kind group_idx token_ids block_hashes
mamba 0 3840 1
mamba 1 3840 1
mamba 2 3840 1
mla_attention 3 3840 5
Indexing all four is worse than indexing none of them. vLLM mixes no group id
into the block hash, so at equal block sizes a Mamba hash COLLIDES with an
attention hash and overwrites its entry in map[engine_hash] -> router_hash;
the next chunk then resolves its parent to the wrong node and every block
after it is hashed off a poisoned chain. Which of the four wins depended on
intra-batch emission order, which vLLM does not promise.
So filter by kv_cache_spec_kind, the way Dynamo's is_main_attention() does --
which is why Dynamo survives hybrid models and we did not. The two fields the
filter needs, group_idx and kv_cache_spec_kind, were on the wire all along;
our msgspec struct simply did not declare them, so they were discarded before
anything could look at them. Same change in the Rust decoder.
mla_attention is NOT optional in the whitelist. Kimi-K3's attention layers are
MLA, so a filter written as == "full_attention" discards 100% of that model's
usable events -- the same empty view, reached from the other side. The filter
fails OPEN when the field is absent (SGLang, older vLLM builds), because those
engines were indexed before this change and a closed default would switch
kv-aware routing off for them silently.
On a length disagreement the event is NOT dropped whole. Its view entries are
correct -- chained from a resolved parent over a contiguous token span -- and
dropping them measurably costs hits: on Qwen3.5-0.8B (same event shapes), 18
of 32 requests hit a full prefix with the whole-event drop versus 22 of 32
without the filter at all. What is not correct is the hash-to-block pairing,
so the view is indexed and the map is not. A later event then misses its
parent and is dropped, which under-reports hits instead of mis-reporting them.
Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…'s rejects The per-field tests prove the group filter rejects what it should. They do not prove a request can HIT afterwards, and that is the property that matters -- production still reported cache_hits=0 with the filter in place, and nothing in the suite would have noticed. This replays what a live worker actually emits, captured off the wire: four events per prefill chunk at block_size 768, three Mamba groups reporting the whole 3840-token span against a single hash and one MLA group reporting it against five, with the Mamba groups repeating the attention group's hashes because vLLM mixes no group id into them. Two chunks, not one, because a single chunk never exercises the parent map: chunk two names chunk one's last block as its parent, and resolving that lookup against a colliding Mamba hash is the failure mode the filter exists to prevent. Then it asks the question routing asks -- hash the prompt, count its blocks in the view -- and requires all ten. Also asserts a divergent tail hits the shared prefix and stops, since a filter that over-matched would route a new prompt onto another prefix's KV. Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
…orders Three drafts of this test passed against the unfixed code before it caught anything, and each failure mode is worth keeping in mind when reading it. Draft 1 constructed BlockStored objects in Python. That cannot catch the defect at all: group_idx and kv_cache_spec_kind were on the wire the whole time and were dropped because the struct did not declare them, so a test that builds the struct itself starts from the fixed state. This one publishes msgpack dicts over a real socket and lets the decoder do its job. Draft 2 emitted the four per-chunk events in the order the capture happened to show, Mamba first and attention last. In that order the attention group's five correct hashes overwrite whatever the Mamba groups wrote, and the view comes out right even with no filter at all. vLLM promises no order; the test now runs both and the fix is what makes the outcome independent of it. Draft 3 asserted the view's SIZE and which worker won. Both survive a poisoned chain: ten entries still appear, five of them hashed from a clobbered parent, and a worker with five real blocks still beats an empty one. It now counts how many of the prompt's blocks actually match, which is the number routing acts on. Against the pre-fix implementation the mamba-last case reports 5/10 blocks visible on the worker that just served the prompt. Fixed: 10/10, both orders. Worth stating plainly: on the shapes captured from this cluster, the Mamba events are always sparse, so the length check alone rejects them and the group filter is not what makes this test pass. The filter earns its place on the orders and group kinds the length check cannot see -- a non-attention group whose lengths happen to agree, which sliding-window groups produce under VLLM_PREFIX_CACHE_RETENTION_INTERVAL. Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
The suite could answer "does a hit happen". It could not answer "does the RIGHT worker get it", because no test ever gave the policy a choice. That gap is exactly the shape of the production failure this branch started from: 32 of 32 requests to one worker while the second node sat idle, and nothing red. These cases have one correct answer each, arrived at by hand rather than by degree: a prompt one worker holds entirely (20/20 to the holder, and the same with the candidate list reversed, so order cannot be what decided it); two disjoint prefixes that must sort themselves 10/10 BY CONTENT -- round-robin also produces 10/10 overall while sending half of each prefix to the wrong worker, so the split is asserted per prefix; a strict subset losing to a longer match, down to a single block of difference; a cold fleet that must spread instead of piling up; and the load-versus-locality tradeoff at the default weight, where one cached block and one in-flight request are both worth 1, plus its release on request completion. Checked against two degenerate policies to confirm the cases are not vacuous: replacing the pick with targets[0] fails 8 of 10, and dropping the cache term so only load counts fails 5 of 10. Signed-off-by: Zhang, Jiejing <jiejing.zhang@amd.com>
jiejingzhangamd
force-pushed
the
fix/kv-event-group-filter
branch
from
August 7, 2026 02:21
c019b78 to
234e47f
Compare
Collaborator
Author
|
Force-pushed: A stray Now green locally: |
Collaborator
Author
|
The red e2e jobs are a bad CI host, not this branch.
GitHub annotates every one of them with "The self-hosted runner lost communication with the server", and the step logs 404 ( The one exception is attempt 1 Nothing failed an assertion. Every tier that actually issued a request passed, all of them through the router this PR modifies: @xiaobochen-amd — the worker-side logs are under |
Collaborator
Author
|
ci failed because no resource. @xiaobochen-amd |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The failure
kv-aware routing against Kimi-K3 on two mixed workers logged
cache_hits=0on every decision and pinned all traffic to one worker. The second node's GPUs
sat idle. Not a degraded hit rate — half the fleet.
Cause
vLLM emits one
BlockStoredper KV-cache group, and only the attentiongroup pairs one hash with one block. Kimi-K3 is hybrid (24 MLA + 69 KDA); the
KDA groups run prefix caching in
alignmode, which nulls all but one blockper step, so those events carry
token_idsspanning the whole chunk and asingle surviving hash — with nothing on the wire saying which block it covers.
Captured off a live worker (
block_size 768,--max-num-batched-tokens 4096,so chunked prefill splits at 3840):
kv_cache_spec_kindgroup_idxtoken_idsblock_hashesIndexing all four is worse than indexing none. vLLM mixes no group id into
the block hash, so at equal block sizes a Mamba hash collides with an
attention hash and overwrites its entry in
map[engine_hash] -> router_hash.The next chunk then resolves its parent to the wrong node and everything after
it hashes off a poisoned chain. Which group won depended on intra-batch
emission order, which vLLM does not promise.
Upstream: vllm#44451 is
open and untriaged, vllm#44488
(block offsets) has no reviews. NVIDIA Dynamo survives hybrid models because it
filters with
is_main_attention()— which is what this PR does.The fix
Filter by
kv_cache_spec_kind. The two fields the filter needs (group_idx,kv_cache_spec_kind) were on the wire all along — our msgspec struct just didnot declare them, so they were discarded before anything could look. Same
change in the Rust decoder.
Two decisions worth reviewing specifically:
mla_attentionis mandatory in the whitelist. Kimi-K3's attention layersare MLA, so a filter written as
== "full_attention"discards 100% of thatmodel's usable events — the same empty view, reached from the other side.
engines were indexed before this change; a closed default would switch
kv-aware routing off for them silently.
A length disagreement no longer drops the event whole. An earlier revision
of this branch did, and it measurably costs hits: on Qwen3.5-0.8B (same event
shapes) 18/32 requests hit a full prefix with the whole-event drop, versus
22/32 with no filter at all. The view entries are correct — chained from a
resolved parent over a contiguous token span. What is not correct is the
hash-to-block pairing, so the view is indexed and the map is not. A later
event then misses its parent and is dropped, which under-reports hits instead
of mis-reporting them.
Tests
tests/unit/router239 passed; Rustcargo test80 passed,cargo fmt --checkclean.test_kv_event_group_filter.pytest_kv_event_kimi_k3_shapes.pytest_kv_event_e2e_hybrid.pytest_kv_aware_routing_extremes.pyTwo of these exist because earlier drafts were vacuous. The e2e test builds
events as plain dicts because a test that constructs our own structs starts
from the fixed state and cannot catch a field the schema never declared; and it
runs both emission orders because with the attention event last, its correct
hashes overwrite the Mamba groups' and the bug disappears by luck.
The extremes file exists because nothing in the suite ever gave the policy a
choice — which is precisely why "32/32 to one worker" was green. Verified
non-vacuous against two degenerate policies: replacing the pick with
targets[0]fails 8/10, dropping the cache term so only load counts fails 5/10.Not claimed
A local A/B on Qwen3.5-0.8B gives 19/32 full hits with this fix versus
22/32 on stock v0.2.3. I have not explained that 3-hit gap and am not going
to invent a reason for it. Note the local rig is a dense-path comparison —
the production zero-hit case is not reproduced by it, and its most likely cause
was a concurrency flaw in the affinity test I was measuring with (rounds fired
concurrently, so repeats raced KV-event propagation). The hybrid-group bug here
is real and independently demonstrated by the replay tests, but it is not
established that it alone explains the production symptom.