Skip to content

kv-router: opt-in refcount + SWA-aware hits for hybrid-attention models (Step-5) - #52

Draft
temirulan wants to merge 2 commits into
prod-kv-indexersfrom
kv-router-step5-hybrid
Draft

temirulan wants to merge 2 commits into
prod-kv-indexersfrom
kv-router-step5-hybrid

Conversation

@temirulan

@temirulan temirulan commented Oct 15, 2026 •

Copy link
Copy Markdown

Opt-in --hybrid-swa-hits mode for the standalone KV indexer. It makes /query predict the prefix hit a hybrid-attention engine actually serves. It is off by default: with the flag unset, nothing new is constructed, and the listener and query paths are unchanged for MiMo, Qwen and other models.

Why

DEEP-309 (deepinfra/backend#6255) tested Step-5-Preview: 4x SWA-512 groups plus 1 full-attention (DSA) group, block 64, MTP run as EAGLE. The full-attention-only, set-semantics tree mispredicts its hits:

probe set stock indexer, within 64 tok
multi-turn continuations 0/20 (+192: EAGLE block + previous completion)
shared 4K–31K scaffolds 16/20 (first reuse of a new junction over-predicted by 3.4K–30.7K)
long prompts 10/10 (+64)

After eviction it under-predicts by 9–12% of served cached tokens. The engine re-stores duplicate copies, and evicting the older copy drops a block that is still cached.

The cause is structural:

  • SWA groups keep blocks only at sparse checkpoints.
  • Duplicate stores clash with set semantics.
  • EAGLE drops the last block.

A reference model fed from the same event stream matched the engine's cached_tokens on 56/56 parity probes and 60/60 eviction probes.

What

standalone_indexer/hybrid.rs holds a per-worker oracle fed from the same ZMQ stream:

  • Refcounts: every group keeps a per-copy count. A Removed for a block that still has another copy is held back from the radix tree.
  • SWA groups: the main path filters these out today. The oracle tracks them by engine block hash, which vLLM shares across groups for the same 64-token position.
  • /query and /query_by_hash: both report min(tree, oracle) per worker. The oracle returns max{k ≤ min(chain, (n−1)/bs) : every SWA group holds blocks [k−need, k)} − mtp, with need = ceil((window−1)/bs) + mtp, which is 9 for Step-5.
    • /query_by_hash is deepapi's routing path. deepapi hashes the full blocks of the prompt with xxh3, seed 1337, which is this indexer's XXH3_SEED. The oracle therefore treats every hashed block as reusable: n = blocks·bs + 1.

Flags:

  • --hybrid-swa-hits
  • --mtp-drop-blocks (default 1)
  • --swa-need-blocks (an override; it is derived per group by default)

--hybrid-swa-hits requires --block-size.

Tests

  • cargo test -p dynamo-kv-router --features standalone-indexer --lib passes 606/606: 7 new tests in hybrid.rs and 7 new by-hash tests in server.rs (apply_hybrid).
  • The new tests reproduce the DEEP-309 shapes: the SWA junction without a checkpoint, a continuation capped at the previous prompt's checkpoint, a duplicate copy surviving its first removal, apply_main holding back the tree removal, full coverage, and the no-SWA cap.

Not done

Offline parity replay. DEEP-309 did not capture raw event streams, only per-probe outcomes. The live gate is to re-run deepbot's tools/probes.py and tools/evict2.py (on di-slc6-61, under /data/deepbot/step-5-preview-deep-309/) against this image in shadow.

Engine-side /tokenize and reasoning_content. DEEP-309 also found that /tokenize in chat form drops reasoning_content. That needs a separate engine-side change and is out of scope here.

Image: localhost:30500/dynamo-indexer:kvtest-9bdca6760d

Engine: KV events must be on (--kv-events-config). VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES already defaults to 1 in the Step-5 image. The indexer only decodes int hashes, so never set it to 0.

🤖 Generated with Claude Code

…ls (Step-5)

New --hybrid-swa-hits mode for the standalone indexer (default off; other
models unchanged). For engines that interleave sliding-window groups with full
attention and run MTP as EAGLE (Step-5-Preview: 4x SWA-512 + 1 full-attention
group, block 64), the FA-only, set-semantics tree mispredicts the engine's
prefix hit (DEEP-309: continuations 0/20 within 64 tokens, first reuse of a
new scaffold junction over-predicted by 3.4K-30.7K, 9-12% under-prediction
after eviction).

The mode adds a per-worker oracle fed from the same ZMQ stream:
- per-copy refcounts for every group; a Removed for a block with another
  cached copy is held back from the radix tree
- sliding-window groups (filtered by the main path) are tracked by engine
  block hash, which vLLM shares across groups
- /query reports min(tree, oracle) per worker, where the oracle predicts
  max{k <= min(chain, (n-1)/bs) : every SWA group holds [k-need, k)} - mtp,
  need = ceil((window-1)/bs) + mtp (9 for Step-5)
Flags: --hybrid-swa-hits, --mtp-drop-blocks (default 1), --swa-need-blocks.

Port of the reference model that matched engine cached_tokens on 56/56 parity
and 60/60 eviction probes in DEEP-309.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@temirulan
temirulan deployed to external_collaborator October 15, 2026 21:22 — with GitHub Actions Active
@temirulan
temirulan changed the base branch from main to prod-kv-indexers October 15, 2026 21:23
@github-actions github-actions Bot added documentation Improvements or additions to documentation router container labels Oct 15, 2026
deepapi's KV routing calls /query_by_hash with xxh3 local block hashes it
computes itself (seed 1337, the indexer's XXH3_SEED), not /query. With
--hybrid-swa-hits the by-hash path now reports min(tree, oracle) per worker
too. The probe hashes only the full blocks of a longer prompt, so the oracle
treats every hashed block as reusable (n = blocks * bs + 1); with
n = blocks * bs it would under-predict one block for the 63/64 of prompts
that are not an exact block multiple.

The override is factored into apply_hybrid() and covered by 7 by-hash tests
of the DEEP-309 shapes (resend, new scaffold junction, continuation,
duplicate-copy eviction, never above the tree, unknown worker).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@temirulan
temirulan had a problem deploying to external_collaborator October 15, 2026 21:55 — with GitHub Actions Failure

This branch had an error being deployed

1 failed deployment
external_collaborator — 9bdca676 Deployed Oct 15, 2026 by temirulan via ok-to-test #60
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

container documentation Improvements or additions to documentation router

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant