feat(mem_cache): semantic KV cache reuse via a pluggable fuzzy-match radix backend - #31057
feat(mem_cache): semantic KV cache reuse via a pluggable fuzzy-match radix backend#31057zbennett10 wants to merge 31 commits into
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
github-actions
Bot
added
documentation
…radix backend Adds opt-in semantic KV cache reuse following the |exact|fuzzy|miss| prompt decomposition: when exact prefix matching covers only a prefix, a pluggable FuzzyMatchProvider may nominate donor KV from a previously finished request, which is realized into recipient-owned slots with RoPE position correction before the forward pass. - FuzzyRadixCache(RadixCache) registered as radix-cache backend "fuzzy_match"; all fuzzy behavior (provider match, donor validation and lock-ref pinning, realization-slot pre-allocation with exact-only fallbacks, donor registration) lives in the backend. - Core RadixCache/BasePrefixCache changes are two seams: MatchResult.fuzzy_matched_len and a no-op _on_finished_insert hook. - FuzzyKVRealizer collaborator performs pre-forward donor-KV realization; model_runner.py gains only orchestration (a gated maybe_init and one delegate call beside the deferred-mamba hook). - ReverseRotaryEmbedding / get_reverse_rope / reverse_rotary_emb added to layers/rotary_embedding (native path). - Default SemanticEmbedding provider implemented by the semblend package (Apache-2.0), installed via the sglang[fuzzy-semantic] extra; lazy import with a hard version gate. Semantic lookups are gated on a minimum miss-suffix length to bound no-hit overhead. - Prefill batch log shows #fuzzy-token when fuzzy reuse fired. - End-to-end call-chain doc in mem_cache/fuzzy_match/README.md; 26 CPU unit tests with semblend faked via sys.modules. Signed-off-by: Zach Bennett <zach@worldflowai.com> Co-authored-by: Chenxin Wu <ibifrostki@gmail.com>
2a91330 to
d193777
Compare
…-backend # Conflicts: # python/sglang/srt/model_executor/model_runner.py
| return torch.stack((o1, o2), dim=-1).flatten(-2) | ||
|
|
||
|
|
||
| def reverse_rotary_emb( |
There was a problem hiding this comment.
Could you add a numerical test for reverse_rotary_emb — i.e. applying RoPE and then reverse RoPE to an input x returns x' ≈ x — for both is_neox_style=True and False?
Requested in review: applying RoPE then reverse RoPE returns the input (orthonormal rotation, both Neox and GPT-J layouts), plus the relocation identity the KV realization path relies on: reverse at donor position p then re-apply at target position q equals rotating at q directly. A non-degenerate guard asserts the forward rotation actually changes the input so the roundtrip cannot pass for no-op helpers. Signed-off-by: Zach Bennett <zach@worldflowai.com>
…-backend Conflicts: python/sglang/srt/model_executor/model_runner.py Both sides added a pre-forward hook after the deferred-mamba step: kept fuzzy donor-KV realization (pool write, grouped with the mamba hook) followed by upstream's new DWDP first-layer prefetch. Semantic adaptations to upstream changes in this range: python/sglang/srt/mem_cache/fuzzy_match/fuzzy_radix_cache.py cache_finished_req gained the required keyword-only kv_len_to_handle parameter (sgl-project#29428) and forwards it to super(), matching the other backend overrides. test/registered/unit/mem_cache/test_fuzzy_radix_cache.py Pass kv_len_to_handle at the three finish call sites and drop the removed pop_committed_kv_cache stub. python/sglang/srt/layers/rotary_embedding/base.py ReverseRotaryEmbedding keeps its cos/sin cache dtype condition identical to RotaryEmbedding after the XPU/FP32 cache change. Signed-off-by: Zach Bennett <zach@worldflowai.com>
|
|
||
| mask_len = len(layer_recompute_mask) if layer_recompute_mask else 0 | ||
| for layer_id in range(pool.layer_num): | ||
| if layer_id < mask_len and layer_recompute_mask[layer_id]: |
There was a problem hiding this comment.
The mask comes from the provider, so it's not visible here which layers get zeroed — worth stating explicitly which layers / what fraction on Qwen2.5-7B. More importantly, the reported ROUGE 0.45–0.55 conflates zeroing with reuse, so we can't attribute the drift. An ablation would help here — e.g., rerunning the same three overlap types with zeroing disabled (copy all) to see how much of the drift comes from zeroing vs. reuse itself. (Note a zeroed K still draws softmax weight but contributes V=0 — worth comparing against reusing the donor KV.)
There was a problem hiding this comment.
Ran the ablation you suggested. Same probe that produced the PR table, same 8 clusters, Qwen2.5-7B AWQ on an A10G, one arm with the provider mask as shipped and one arm with zeroing disabled so every layer is copied from the donor. The same matches fired in both arms so the comparison is paired.
Warm vs cold on the pairs where a fuzzy match fired:
| arm | ROUGE-L | token-F1 |
|---|---|---|
| mask on (zeroing) | 0.384 | 0.442 |
| mask off (all copy) | 0.861 | 0.871 |
Per overlap type, mean over all pairs:
| type | fire rate | mask on | mask off |
|---|---|---|---|
| partial_80 | 0.50 | 0.632 | 0.908 |
| partial_60 | 0.25 | 0.778 | 0.929 |
| paraphrase | 0.38 | 0.825 | 1.000 |
| diverse (negative) | 0.00 | 1.000 | 1.000 |
So you called it. The zeroing is the dominant source of the drift on these workloads, not the reuse itself. A zeroed key still takes softmax mass at a zero logit while contributing nothing, which dilutes attention over the reused span exactly as you suspected. We're changing the provider default to copy all layers in the next semblend release and the engine side keeps the mask as an optional provider hook. On visibility, the mask is adaptive per request rather than a fixed list. On this preset it flags the edge layers, typically the first two and the last five of 28, capped around a quarter of the layers.
One caveat for reading the numbers against the June table in the PR description. Realized reuse mass per hit is smaller on the current merged stack than in those runs, so this ablation is the quality attribution and the June table remains the TTFT reference. I also reworked the engine docstring so it describes only the contract and the actual zeroing semantics.
| @@ -0,0 +1,64 @@ | |||
| --- | |||
There was a problem hiding this comment.
Currently, when fuzzy match fires, the quality drop is substantial — ROUGE-L 0.45–0.55 vs ~0.88 for the diverse/rerun control. Given that, it'd be good to clarify the intended application scenarios: what workloads is this aimed at, where the TTFT win justifies this much output drift? (Related: the README's "bound the drift" reads a bit strong against these numbers — might be worth qualifying.)
There was a problem hiding this comment.
Fair challenge. I qualified the wording in both places. The README now says the gates and the mask are intended to reduce the drift rather than bound it and the docs page states the measured range inline next to the control number so nobody reads the quality claim as stronger than the data.
On intended scenarios, the target is latency sensitive resend heavy traffic where the shared content does not sit at a common prefix. Think a shared document re-asked under different task instructions, session content behind a shifted header, or templated enterprise flows. Exact prefix caching recovers nothing there and time to first token is what the operator cares about. Workloads that need byte stable outputs for identical inputs should keep the default backend and the docs page now says that explicitly.
Also worth noting from the ablation in the thread above, most of the measured drift came from the per layer zeroing rather than from reuse. With all layers copied the hit-only aggregate is 0.86 ROUGE against cold and per overlap type it lands between 0.91 and 1.0. We're defaulting the mask off in the next provider release, which closes most of that gap.
| Used to apply RoPE in the recipient's reference frame. | ||
| layer_recompute_mask: Optional list of bools; when ``mask[i]`` is | ||
| True, layer ``i`` is zeroed instead of copied (the prefill | ||
| pass that bookends the block will produce fresh K/V for those |
There was a problem hiding this comment.
The comment says a "prefill pass that bookends the block will produce fresh K/V for those layers," but there's no such pass in this PR — the fuzzy span is in the matched prefix and the extend only computes the miss tail, so zeroed layers stay zeroed. Either point to the code that refills them or update the comment to match the actual behavior.
There was a problem hiding this comment.
You're right, there is no refill pass in this PR and the comment was wrong. I rewrote it to describe what actually happens. The zeros survive the forward pass since the extend only computes the missed tokens, and at request finish the realized slots are inserted into the radix tree, so they stay visible to later exact prefix matches and to later fuzzy matches that pick this request as a donor. I also added a unit test pinning the mask semantics, masked layers zeroed, unmasked layers position corrected, and short masks leaving trailing layers copied.
| _REVERSE_ROPE_DICT: Dict[Tuple, ReverseRotaryEmbedding] = {} | ||
|
|
||
|
|
||
| def get_reverse_rope( |
There was a problem hiding this comment.
get_reverse_rope doesn't appear to have any call site — the realization path uses copy_kv_with_rope_correction, which pulls reverse_rotary_emb from utils and the model's own cos_sin_cache.
There was a problem hiding this comment.
Removed. The ReverseRotaryEmbedding class behind it had no other caller so it went too, which puts layers/rotary_embedding back to the upstream tree except for the reverse helper in utils.py that the realization path actually uses. The numerical round trip test from the earlier thread lives in test/registered/unit/layers/test_rotary_embedding_utils.py and covers both neox styles.
| import torch | ||
|
|
||
|
|
||
| def as_long_tensor(obj, device) -> torch.Tensor: |
There was a problem hiding this comment.
as_long_tensor is also defined but never used.
There was a problem hiding this comment.
Removed as well, thanks for the catch.
Remove get_reverse_rope and the ReverseRotaryEmbedding class, which had no call site (the realization path uses reverse_rotary_emb from utils directly), restoring layers/rotary_embedding to the upstream tree. Remove the unused as_long_tensor helper. Correct the layer mask documentation, which described a refill pass that does not exist: zeroed layers stay zeroed through the forward pass and the realized slots are inserted into the radix tree, and at a zeroed layer reused positions score zero logits that take softmax mass while contributing zero V. Add a unit test pinning the mask semantics. Qualify the drift wording in the module README and docs page and state the measured quality range explicitly. Signed-off-by: Zach Bennett <zach@worldflowai.com>
Upstream now threads cache_salt from the request into RadixKey during finished-request insertion, so the stub carries the field. Signed-off-by: Zach Bennett <zach@worldflowai.com>
Correctness. Donor registration now skips inserts that deduplicated against existing tree content, because the captured kv indices in the deduplicated range are freed right after the callback and the provider would retain dangling slots. The cache salt joins the provider namespace on both the registration and match paths so donors never serve across salts. RadixKey.limit caps the tokens the provider sees. Empty segment lists normalize to None so every consumer has one no-segments sentinel. Startup guards. The backend now rejects MLA-style pools and non-RoPE models at model load instead of silently serving donor KV without position correction, requires page_size 1 and tp_size 1 and pp_size 1, validates the fuzzy flag ranges at argument parsing, and forces the overlap scheduler off: realization mutates request KV state on the forward stream and is not ordered against overlapped batch preparation. The realizer's silent degrade path is now a loud error. API surface. Removed dead provider API nothing consumed: the set_min_match_length method, the cached_token_ids result field, the provider-internal match_entry handle, the segment NodeRef addressing fields that no consumer implemented, and the rejection_reason quality field. Renamed layer_recompute_mask to layer_zero_mask to say what it does. The provider's fallback tokenizer now comes from the served model path instead of an undocumented environment variable. Observability and docs. The fuzzy match success log line now fires after validation so it marks an accepted match as documented, per request provider logs drop to debug, the metrics comment states that fuzzy tokens are a subset of the hit counter, and the README and docs page claims are aligned with measured behavior, including the client visibility of reused tokens and the startup rejections. The semblend extra takes an upper version bound. Tests. End to end realization math on real tensors for both the contiguous and scattered segment paths, mask zeroing semantics, and cache salt namespacing. Signed-off-by: Zach Bennett <zach@worldflowai.com>
|
Pushed the review fixes along with a fresh merge from main. Summary of what changed since the review. The dead code is gone so the rotary embedding package is back to upstream apart from the utils helper the realization path uses. The layer mask documentation now matches the real behavior and there is a new unit test pinning the mask semantics. The drift wording is qualified with the measured numbers stated inline in the docs. The ablation on the zeroing question is in the thread above, short version is that the zeroing caused most of the drift and the next provider release defaults it off. Beyond the review items I did a hardening pass on the backend. Donor registration now skips inserts that deduplicated against existing tree content so the provider does not retain freed slots. The cache salt joins the provider namespace on both registration and match so donors never serve across salts. Unsupported stacks are rejected at startup with clear errors instead of degrading silently, that covers MLA style pools, models without rotary embeddings, page size above 1, and multi GPU. The overlap scheduler is forced off with a warning. Flag values are validated at argument parsing. The provider interface also slimmed down, the fields nothing was consuming are gone and the mask field is renamed to say what it does. New tests cover the realization math end to end on real tensors for both the contiguous and scattered segment paths plus cache salt namespacing. |
Signed-off-by: Zach Bennett <zach@worldflowai.com>
Upstream moved per-request KV lifecycle state onto a ReqKvInfo record, so the fuzzy reuse fields move with it: they travel through detach_kv and clear on retraction, where a stale realization plan would otherwise point at slots the request no longer owns. Server-argument resolution moved to per-domain hooks, so the fuzzy backend's envelope checks move to the KV cache hook and register in the resolution pipeline. The cache-backend name now comes from the published memory namespace. Signed-off-by: Zach Bennett <zach@worldflowai.com>
The provider's per layer zeroing measurably degrades output on reused spans, and 0.3.17 is the first release that defaults to copying all of them. Raise both floors together so the runtime guard cannot admit an older build that the extra no longer installs. Signed-off-by: Zach Bennett <zach@worldflowai.com>
Reading model_path off the supplied server_args record returns the raw CLI input, so the provider's fallback tokenizer would miss a path that resolution or a later weight update rewrote. The backend factory already holds the resolved model config, so it passes the effective path in. Signed-off-by: Zach Bennett <zach@worldflowai.com>
Realization slots came from the free pool only, so once a few long donors and their recipients were resident every later provider match fell back to exact-only in silence: on an A10G with a 61K-token pool, 3 of 24 wrapper-shift items served at 8K tokens and 1 of 24 at 16K, exactly the number that fit before the pool filled. The donor is now pinned first so making room can never free it, then unlocked entries are evicted the way prefill admission does; a genuine capacity failure still degrades cleanly and releases the pin. Signed-off-by: Zach Bennett <zach@worldflowai.com>
A provider result whose span starts at the exact-matched length was dropped, and the contiguous realizer skipped it, on the assumption that token-and-position-aligned content belongs to the exact tree. A verified paraphrase served whole is different tokens at the same positions: it needs its own slots (delta 0 copy) or the recipient reads donor-owned slots and later inserts them into its own branch. The paraphrase_verified tier now admits the span and copies it; exact-tier aligned spans keep the old behavior. Signed-off-by: Zach Bennett <zach@worldflowai.com>
The adapter in this release serves verified paraphrases as a position-aligned contiguous span after the exact prefix, which the backend now realizes with a zero delta; earlier releases emitted a shape the backend drops. Signed-off-by: Zach Bennett <zach@worldflowai.com>
|
Three follow-up commits from long-context testing (8K to 24K token prompts on an A10G):
Unit tests cover all three (fuzzy radix, realizer, provider suites). |
A request that consumed donor KV holds approximate content for its fuzzy span and for everything computed after it. cache_finished_req inserted its whole sequence into the radix tree, so the next request sharing its token prefix received the donor's KV through a trusted exact match (observed: answers about one service naming another service's report). Served requests are no longer inserted or registered as donors; their slots beyond the exact prefix are freed as before. Signed-off-by: Zach Bennett <zach@worldflowai.com>
Signed-off-by: Zach Bennett <zach@worldflowai.com>
|
Two more follow-ups from long-context testing:
Both covered by unit tests in the fuzzy radix suite. |
…napshot A provider registers a donor with a snapshot of the request's slot mapping, which includes prefix slots shared with earlier tree nodes. If such a node is evicted and its slots reused before a recipient matches, the served span starts inside another request's content (observed: answers about one service opening with another service's report). Served slots are now resolved from the pinned donor's tree path at the same positions, for contiguous results and NodeRef segments; a span the path does not cover is dropped. Signed-off-by: Zach Bennett <zach@worldflowai.com>
0.3.21 aligns the leading donor candidates instead of only the top fast estimate, which keeps wrapper-shifted prompts on the position-corrected copy path among near-duplicate donors, and its fact gate compares identifiers as well as numbers. Signed-off-by: Zach Bennett <zach@worldflowai.com>
|
Head is now 2bcb005: the semblend floor moves to 0.3.21. What changed in 0.3.21 and why it matters here: among many near-duplicate donors, the fast reuse estimate could rank a decoy above the donor that actually shares the prompt's tokens, and only the top estimate was ever aligned, so wrapper-shifted prompts fell through to the paraphrase tier. The store now aligns the leading candidates and keeps the best alignment, and the fact gate also compares identifiers (service ids, incident codes), not just numbers. Re-measured on this head with 0.3.21 (A10G, Qwen2.5-7B fp16, 24 prompts per length, same document under a different instruction wrapper, each prompt paired with its own cold run):
All served on the exact tier (aligned copy, reuse 0.995 to 0.999). On LongBench multi-document QA (2K to 17K tokens) all 24 prompts were served whole; gold-answer pass count 12 served vs 13 cold. |
4 participants
Motivation
Exact prefix matching cannot reuse KV across prompts that share meaning but
not leading tokens: paraphrased questions over the same document, RAG
prompts with reordered context, multi-user workloads hitting the same
knowledge base. This PR adds opt-in semantic KV cache reuse following
the
|exact|fuzzy|miss|prompt decomposition from the fuzzy prefixmatching design discussion (Draft_Prefix_Matching.md, @ibifrost): when the
radix tree covers only a prefix, a pluggable
FuzzyMatchProvidermaynominate donor KV from a previously finished request, which is copied into
recipient-owned slots with RoPE position correction before the forward
pass.
Design goals, per maintainer guidance (@hzh0425): keep the default path
untouched, register the feature as a radix-cache backend, keep
model_runnerchanges minimal.Modifications
FuzzyRadixCache(RadixCache)registered as radix-cache backendfuzzy_match(flexkv pattern). All fuzzy behavior lives in the backend:provider match on exact-miss, donor validation + pinning (lock refs via a
TreeNode-id registry), realization-slot pre-allocation with clean
exact-only fallback, donor registration on request finish.
RadixCache/BasePrefixCachechanges = two seams:MatchResult.fuzzy_matched_len(new trailing field) and a no-op_on_finished_inserthook between the finished-request insert andduplicate-slot freeing.
FuzzyKVRealizercollaborator does pre-forward donor-KV realization(V copy; K reverse-RoPE at donor positions + apply at target positions;
per-layer zero-out mask).
model_runner.pygains only orchestration: agated
maybe_init_fuzzy_kv_realizer()and one delegate call in_forward_rawbeside the deferred-mamba hook (outside CUDA-graphcapture/replay, covers all extend branches).
ReverseRotaryEmbedding/get_reverse_rope/reverse_rotary_embadditions to
layers/rotary_embedding(native path; no CUDA kernelexists for the reverse direction).
SemanticEmbeddingprovider (default): implemented by thesemblendpackage(PyPI, Apache-2.0), installed via
the new extra
pip install "sglang[fuzzy-semantic]". SemBlend is anopen-source
library for semantic KV-cache donor discovery and reuse planning: it
embeds prompts (MiniLM ONNX), maintains the donor index, aligns donor
and recipient token chunks (N:M), scores reuse quality against
configurable gates, and emits the reuse plan — donor span, positions,
and the per-layer mask — that this PR's realization path consumes.
Keeping that pipeline in a package keeps SGLang's tree free of the
embedding/ANN stack while the reuse-planning algorithms iterate. Lazy
import + hard version gate at startup with an actionable error. The
FuzzyMatchProviderinterface is public so custom providers can plugin.
#fuzzy-token: Nwhen fuzzyreuse fired (previously realized reuse showed as
#cached-token: 0).--fuzzy-*flags (provider, threshold, min-reuse-ratio,min-match-length, model-arch). Provider-internal tuning stays in config
defaults.
in
python/sglang/srt/mem_cache/fuzzy_match/README.md(requested inreview of the experimental branch).
test_fuzzy_match_providers.py,test_fuzzy_radix_cache.py), semblend faked via sys.modules — no new CIdependency. Cover pool-accounting failure modes: stale-donor and
kv-length-mismatch fallbacks (leak-checked), donor pin/release lifecycle,
donor-registration ordering vs duplicate freeing, node-registry
maintenance, registry selection + unsupported-config rejection, and
concurrent recipients pinning one donor.
Not supported yet: MLA pools, EAGLE,
multi-region reuse, hierarchical-cache interaction.
Follow-on work
This PR deliberately ships the minimal
|exact|fuzzy|miss|pattern — onecontiguous fuzzy span re-anchored at the exact-prefix boundary — because it
reuses the existing inference pipeline with no structural changes. Planned
follow-ups, in rough order:
|exact|miss|fuzzy|miss|fuzzy|...),the realistic shape of RAG and multi-document workloads: the scheduler
builds a segmented prefill plan of fresh spans plus donor-KV spans, and
the attention backend consumes it via FlashInfer's
VariableBlockSparseAttentionWrapper(per the design discussion, abetter fit than
custom_mask). Early prototypes on 64K multi-donorprompts show mean 7.65x TTFT speedup versus cold prefill with good
quality — worth its own RFC once this foundation lands.
layer_recompute_maskis a zero-out signal); recompute is the cleanerlong-term answer for quality.
to shrink the remaining enabled-but-idle overhead on dense
short-request workloads.
host tiers.
Accuracy Tests
All runs on THIS branch (A10G, Qwen2.5-7B-Instruct-AWQ, fresh server per
configuration, 2026-07-13).
GSM8K few-shot (n=1000), fuzzy on vs off: baseline accuracy 0.862
(invalid 0.001) vs fuzzy-on 0.865 (invalid 0.001) — no degradation
(delta within batching noise). The provider's gates produced zero
fuzzy fires across all 1000 questions (short, semantically dissimilar
prompts never meet the reuse-ratio gate; the few-shot prefix hits the
exact radix path), and zero pool leaks. This is the negative control:
enabling the backend does not perturb workloads without semantic
overlap.
Warm-vs-cold A/B at 8K/16K (cnn_dailymail + wikihow clusters, at
default flags). ROUGE-L / token-F1 compare the warm (reused) answer
against the cold answer for the same prompt; the measurement scale:
identical-prompt reruns without any reuse score ~0.88/0.90
(batching-noise floor), exact-prefix radix reuse scores 0.94/0.96.
A sweep at looser settings (threshold 0.50, min-match 1) produced the
same hit rates and the same drift — on genuinely similar pairs the
threshold does not change what a hit does, it changes which pairs hit
(on GSM8K, defaults fire on 1 in 1000 requests; the diverse control
never fires). Read: a fuzzy hit trades TTFT for output drift beyond
the rerun-noise floor — reuse is quality-gated but not lossless;
answers remain on-topic and grounded (samples in the raw results), and
per-layer recompute (follow-on) is the path to closing the gap. Zero
pool-consistency failures across all sweeps and both GSM8K runs.
Experimental-branch reference @16k (n=1): PPL(warm)/PPL(cold) 0.9587.
Statement of semantics: reuse is quality-gated but not lossless; the
layer mask is a zero-out signal, not per-layer recompute.
Speed Tests and Profiling
104/119 donor tokens with RoPE correction; only 15 tokens prefilled.
Experimental-branch reference at 16K: cold TTFT 4548ms → warm 230ms
(19.8x); donor registration runs off the hot path (async embed).
[OPTIONAL: re-measure 16K TTFT pair on this branch]
1164 tok/s output throughput (~7%) and 127.3 → 136.8s wall on a
worst-case-dense workload (a thousand short requests, every one
indexed as a donor). The synchronous per-miss lookup cost is bounded
by
fuzzy_min_suffix_tokens(default 256): suffixes too short toamortize a semantic lookup never pay for one. The residual is the
asynchronous donor-registration embedding, which amortizes on
long-context workloads (fewer, larger requests); donor-registration
budgeting is listed as follow-on. Without the suffix gate the same
workload measured ~10-11%. The feature is opt-in per deployment.
Reproduction
Checklist / credits
FuzzyMatchProviderframework: @ibifrost (Draft_Prefix_Matching.md);backend-registry direction: @hzh0425.
Co-authored-by: Chenxin Wu ibifrostki@gmail.com
Big thanks to @ibifrost and @hzh0425 for all of the hard work/advice/reviewing over the past several months.
CI States
Latest PR Test (Base): ❌ Run #33812107578
Latest PR Test (Extra): ❌ Run #33812107456
Latest PR Test (AMD ROCm 7.2): ❌ Run #33812107729