From 0ed827345f4f1f978d8d0259e63126cba20fb3bc Mon Sep 17 00:00:00 2001 From: krakhit <61501745+krakhit@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:08:32 +0000 Subject: [PATCH 1/2] exact:content fuzzy match provider --- python/sglang/srt/environ.py | 6 + .../srt/mem_cache/fuzzy_match/chunker.py | 180 +++++++++ .../srt/mem_cache/fuzzy_match/config.py | 4 +- .../fuzzy_match/exact_hash_provider.py | 198 ++++++++++ .../fuzzy_match/fuzzy_match_provider.py | 8 +- .../fuzzy_match/fuzzy_radix_cache.py | 20 + .../srt/mem_cache/fuzzy_match/realizer.py | 44 ++- python/sglang/srt/server_args.py | 5 +- .../fuzzy_match/test_exact_hash_e2e_safety.py | 350 ++++++++++++++++++ .../test_exact_hash_shifted_offset_kl.py | 167 +++++++++ .../mem_cache/fuzzy_match/test_chunker.py | 119 ++++++ .../fuzzy_match/test_exact_hash_provider.py | 166 +++++++++ 12 files changed, 1256 insertions(+), 11 deletions(-) create mode 100644 python/sglang/srt/mem_cache/fuzzy_match/chunker.py create mode 100644 python/sglang/srt/mem_cache/fuzzy_match/exact_hash_provider.py create mode 100644 test/registered/fuzzy_match/test_exact_hash_e2e_safety.py create mode 100644 test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py create mode 100644 test/registered/unit/mem_cache/fuzzy_match/test_chunker.py create mode 100644 test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b4ff61057eb6..7cac95a2167e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -364,6 +364,12 @@ class Envs: SGLANG_TEST_MAMBA_LAZY_ALLOC_FAIL = EnvBool(False) # KL tests: skip the cache-hit count assertion (e.g. when alloc failure reduces hits) SGLANG_TEST_SKIP_CACHE_HIT_ASSERT = EnvBool(False) + # Fuzzy-match E2E tests: force every chunk fingerprint to the same + # constant, so two genuinely different chunks collide by construction — + # exercises the mandatory token-ID equality-check fallback at the real + # system level (through the live scheduler), which can't otherwise be + # forced without an astronomically unlikely real hash collision. + SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION = EnvBool(False) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY = EnvInt(0) SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE = EnvBool(True) # Physical KV-page checks: committed<=allocated + no page alias. diff --git a/python/sglang/srt/mem_cache/fuzzy_match/chunker.py b/python/sglang/srt/mem_cache/fuzzy_match/chunker.py new file mode 100644 index 000000000000..aaecc5dc8912 --- /dev/null +++ b/python/sglang/srt/mem_cache/fuzzy_match/chunker.py @@ -0,0 +1,180 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Content-defined chunking (CDC) over token-ID sequences. + +Boundaries depend only on the last ``WINDOW_SIZE`` tokens, not on absolute +position or anything upstream — the property RadixAttention's chained hash +lacks by design. Used by ``ExactHashProvider`` to find donor spans +independent of where they sit in the current prompt. + +The chunking parameters below (window size, boundary mask, chunk-size +clamp, sink carve-out) follow the methodology in Ma, Eitzinger, and +Köstler, "Irminsul: MLA-Native Position-Independent Caching for Agentic +LLM Serving" (arXiv:2605.05696), Section 4 and Appendix B. That paper's +own reuse mechanism (RoPE delta-rotation over MLA's decomposed KV) is +unrelated to this module — only its CDC/chunking layer is used here. +""" + +from __future__ import annotations + +import random +from typing import List, NamedTuple, Sequence + +_MASK64 = (1 << 64) - 1 + +# Gear-hash rolling window, in tokens. A boundary decision only ever depends +# on the last WINDOW_SIZE tokens (Gear hash's left-shift makes older tokens' +# contributions fall off the top of a 64-bit accumulator automatically — +# no explicit "remove the oldest byte" step needed, unlike Rabin-Karp). +WINDOW_SIZE = 64 + +# Tuned so a boundary fires with probability 2^-BOUNDARY_BITS per token, +# i.e. an expected chunk size of 2**BOUNDARY_BITS tokens. k=7 is the value +# the cited paper's mask-exponent ablation (Appendix B) settled on — wider +# masks push expected chunk length above MAX_CHUNK_TOKENS, which clamps +# chunks to fixed-size and degenerates back to the fixed-block hashing this +# design exists to avoid. +BOUNDARY_BITS = 7 # 2**7 = 128 +_BOUNDARY_MASK = (1 << BOUNDARY_BITS) - 1 + +MIN_CHUNK_TOKENS = 32 +MAX_CHUNK_TOKENS = 512 + +# Content before this position is never chunked/cached — attention-sink +# territory: early positions absorb a disproportionate, content-independent +# share of attention regardless of what's actually there, so a content hash +# match in that zone isn't a trustworthy signal. 32 tokens matches the +# cited paper's own sink-extent measurement (Section 7.1). +SINK_TOKENS = 32 + +# Fixed seed: the Gear table must be identical across the write side +# (cache_on_request_finished) and the read side (match_on_prefix_miss), +# and across process restarts, or chunk boundaries silently stop lining up. +_GEAR_SEED = 0x516D5A4F2F3A7C5C +_GEAR_TABLE_SIZE = 1 << 16 # token_id % this size indexes the table + + +def _build_gear_table() -> List[int]: + rng = random.Random(_GEAR_SEED) + return [rng.getrandbits(64) for _ in range(_GEAR_TABLE_SIZE)] + + +_GEAR_TABLE = _build_gear_table() + +try: + import xxhash + + def _real_fingerprint(token_ids: Sequence[int]) -> int: + h = xxhash.xxh64(seed=0) + h.update(_tokens_to_bytes(token_ids)) + return h.intdigest() + +except ImportError: + import hashlib + + def _real_fingerprint(token_ids: Sequence[int]) -> int: + # xxhash isn't installed in this environment; blake2b is stdlib, + # fast enough for a fingerprint (not a security boundary), and + # gives the same "cheap, non-cryptographic-role" fingerprint + # xxHash64 would — the token-ID equality check downstream is what + # actually guards correctness, not this hash's collision odds. + digest = hashlib.blake2b( + _tokens_to_bytes(token_ids), digest_size=8 + ).digest() + return int.from_bytes(digest, "little") + + +def _fingerprint(token_ids: Sequence[int]) -> int: + # Test-only: force every chunk to the same fingerprint so a real E2E + # test can exercise the mandatory token-ID equality-check fallback + # (never trust the hash alone) without needing an astronomically + # unlikely real collision. Off by default in production. + from sglang.srt.environ import envs + + if envs.SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION.get(): + return 0 + return _real_fingerprint(token_ids) + + +def _tokens_to_bytes(token_ids: Sequence[int]) -> bytes: + return b"".join(t.to_bytes(4, "little", signed=False) for t in token_ids) + + +class Chunk(NamedTuple): + """One content-defined chunk: token IDs plus its offset within the + sequence that was chunked (not an absolute sequence position — callers + add their own base offset).""" + + start: int # offset into the chunked sequence, inclusive + end: int # offset into the chunked sequence, exclusive + token_ids: List[int] + fingerprint: int + + +def chunk_tokens(token_ids: Sequence[int]) -> List[Chunk]: + """Split ``token_ids`` into content-defined chunks. + + Boundaries are declared where a Gear-hash rolling hash over the last + ``WINDOW_SIZE`` tokens has its low ``BOUNDARY_BITS`` bits all zero, + clamped to ``[MIN_CHUNK_TOKENS, MAX_CHUNK_TOKENS]``. The same content at + a different position in a different call produces the same boundaries + and the same per-chunk fingerprint, as long as ``WINDOW_SIZE`` tokens of + identical context precede each boundary. + """ + n = len(token_ids) + chunks: List[Chunk] = [] + chunk_start = 0 + h = 0 + for i, tok in enumerate(token_ids): + h = ((h << 1) + _GEAR_TABLE[tok % _GEAR_TABLE_SIZE]) & _MASK64 + chunk_len = i - chunk_start + 1 + if chunk_len < MIN_CHUNK_TOKENS: + continue + at_boundary = (i - chunk_start + 1 >= WINDOW_SIZE) and ( + h & _BOUNDARY_MASK == 0 + ) + if at_boundary or chunk_len >= MAX_CHUNK_TOKENS: + span = token_ids[chunk_start : i + 1] + chunks.append( + Chunk( + start=chunk_start, + end=i + 1, + token_ids=list(span), + fingerprint=_fingerprint(span), + ) + ) + chunk_start = i + 1 + h = 0 + + # Trailing remainder: the rolling hash may never hit a boundary before + # the input runs out, regardless of how much is left — this is not + # bounded to be small. If what's left is at least MIN_CHUNK_TOKENS, + # emit it as a final chunk (standard CDC practice: the last chunk of + # any given input is "whatever's left" once you hit the end); only drop + # it when it's genuinely too small to amortize a lookup. This only + # affects the *final* chunk of whatever span is being chunked — content + # that isn't at the tail end still gets purely content-defined + # boundaries, unaffected by this. + if len(token_ids) - chunk_start >= MIN_CHUNK_TOKENS: + span = token_ids[chunk_start:] + chunks.append( + Chunk( + start=chunk_start, + end=len(token_ids), + token_ids=list(span), + fingerprint=_fingerprint(span), + ) + ) + + return chunks diff --git a/python/sglang/srt/mem_cache/fuzzy_match/config.py b/python/sglang/srt/mem_cache/fuzzy_match/config.py index cc17f769234d..ecb47d2eb352 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/config.py +++ b/python/sglang/srt/mem_cache/fuzzy_match/config.py @@ -75,9 +75,9 @@ def __post_init__(self): f"got {self.fuzzy_semantic_threshold}" ) - if self.fuzzy_match_provider not in ("SemanticEmbedding",): + if self.fuzzy_match_provider not in ("SemanticEmbedding", "ExactHash"): raise ValueError( - f"fuzzy_match_provider must be 'SemanticEmbedding', " + f"fuzzy_match_provider must be 'SemanticEmbedding' or 'ExactHash', " f"got {self.fuzzy_match_provider}" ) diff --git a/python/sglang/srt/mem_cache/fuzzy_match/exact_hash_provider.py b/python/sglang/srt/mem_cache/fuzzy_match/exact_hash_provider.py new file mode 100644 index 000000000000..817671a0764e --- /dev/null +++ b/python/sglang/srt/mem_cache/fuzzy_match/exact_hash_provider.py @@ -0,0 +1,198 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exact-content fuzzy-match provider. + +Finds donor KV for content that is byte-identical to the current prompt's +unmatched tail but sits at a *different* offset than where it was originally +computed — as opposed to ``SemanticEmbeddingProvider``, which matches +merely-similar content and is explicitly not lossless. This provider's own +correctness bar is exact reuse: the token-ID equality check on every hit is +mandatory, not a nice-to-have — a fingerprint collision must never be trusted +without confirming the underlying token IDs actually match. + +Current limitation: single-chunk, non-segmented matches only +(``FuzzyMatchResult.segments=None``, realized via +``FuzzyKVRealizer._realize_contiguous``). Multi-chunk / N:M segment matches +are a natural follow-up, not required for the mechanism to be correct. +""" + +from __future__ import annotations + +import logging +from typing import Dict, List, Optional, Tuple + +import msgspec +import torch + +from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX +from sglang.srt.mem_cache.fuzzy_match.chunker import SINK_TOKENS, chunk_tokens +from sglang.srt.mem_cache.fuzzy_match.config import FuzzyMatchConfig +from sglang.srt.mem_cache.fuzzy_match.fuzzy_match_provider import ( + FuzzyMatchProvider, + FuzzyMatchResult, +) + +logger = logging.getLogger(__name__) + +# (extra_key, fingerprint) -> candidates; a list because two different +# chunks can share a fingerprint (rare, but the point of the equality check +# below is to never trust the hash alone regardless of how rare). +_StoreKey = Tuple[Optional[str], int] + + +class _ChunkEntry(msgspec.Struct): + """One registered donor chunk.""" + + token_ids: List[int] + kv_indices: torch.Tensor + start_pos: int # p_src: absolute position this chunk was computed at + donor_last_node_id: Optional[int] = None + + +class ExactHashProvider(FuzzyMatchProvider): + """Content-defined-chunking + exact-hash donor matching.""" + + def __init__(self, config: FuzzyMatchConfig): + super().__init__(config) + self._store: Dict[_StoreKey, List[_ChunkEntry]] = {} + # request_id -> keys registered by that request's own + # cache_on_request_finished call, so the immediately-following + # on_donor_inserted callback can attach the real TreeNode id. + self._pending_by_request: Dict[str, List[_StoreKey]] = {} + logger.info("ExactHashProvider initialized") + + # ------------------------------------------------------------------ + # FuzzyMatchProvider contract + # ------------------------------------------------------------------ + + def cache_on_request_finished( + self, + request, + token_ids: List[int], + kv_cache: torch.Tensor, + cache_start_pos: int, + cache_end_pos: int, + radix_tree=None, + ) -> bool: + request_id = _request_id(request) + if request_id is None or request_id.startswith(HEALTH_CHECK_RID_PREFIX): + return False + if cache_end_pos <= cache_start_pos: + return False + + # Never register content whose original occurrence started inside + # the attention-sink zone — sink behavior is position-driven, not + # reliably content-only, so a hash match there isn't trustworthy. + chunk_region_start = max(cache_start_pos, SINK_TOKENS) + if chunk_region_start >= cache_end_pos: + return False + + extra_key = getattr(request, "extra_key", None) + region = token_ids[chunk_region_start:cache_end_pos] + chunks = chunk_tokens(region) + + registered_keys: List[_StoreKey] = [] + for c in chunks: + abs_start = chunk_region_start + c.start + abs_end = chunk_region_start + c.end + entry = _ChunkEntry( + token_ids=c.token_ids, + kv_indices=kv_cache[abs_start:abs_end].detach().clone(), + start_pos=abs_start, + ) + key = (extra_key, c.fingerprint) + self._store.setdefault(key, []).append(entry) + registered_keys.append(key) + + if registered_keys: + self._pending_by_request[request_id] = registered_keys + logger.info( + "[EXACT_HASH] cache_on_request_finished: rid=%s tokens=%d chunks=%d", + request_id, + cache_end_pos - cache_start_pos, + len(chunks), + ) + return bool(registered_keys) + + def match_on_prefix_miss( + self, + prompt_token_ids: List[int], + already_matched_len: int, + request=None, + extra_key=None, + ) -> Optional[FuzzyMatchResult]: + tail = prompt_token_ids[already_matched_len:] + chunks = chunk_tokens(tail) + if not chunks: + return None + + # Only the first chunk of the unmatched tail. It always starts at + # offset 0 of `tail`, so its target position in the current prompt + # is exactly `already_matched_len` — no extra offset math needed + # for the single-chunk case. + c = chunks[0] + candidates = self._store.get((extra_key, c.fingerprint)) + if not candidates: + return None + + entry = _first_equal_match(candidates, c.token_ids) + if entry is None: + logger.debug( + "[EXACT_HASH] fingerprint hit but token-ID mismatch — " + "hash collision, never trusting the hash alone" + ) + return None + + return FuzzyMatchResult( + cached_token_count=len(entry.token_ids), + cached_token_ids=entry.token_ids, + prompt_token_count=len(prompt_token_ids), + kv_cache_indices=entry.kv_indices, + position_offset=already_matched_len - entry.start_pos, + cached_start_pos=entry.start_pos, + segments=None, + donor_last_node_id=entry.donor_last_node_id, + ) + + def on_donor_inserted(self, request, donor_last_node_id: int) -> None: + request_id = _request_id(request) + if request_id is None: + return + keys = self._pending_by_request.pop(request_id, None) + if not keys: + return + for key in keys: + for entry in self._store.get(key, []): + if entry.donor_last_node_id is None: + entry.donor_last_node_id = donor_last_node_id + + def on_cache_reset(self) -> None: + self._store.clear() + self._pending_by_request.clear() + logger.info("[EXACT_HASH] cleared store on cache reset") + + +def _request_id(request) -> Optional[str]: + rid = getattr(request, "rid", None) or getattr(request, "request_id", None) + return str(rid) if rid is not None else None + + +def _first_equal_match( + candidates: List[_ChunkEntry], token_ids: List[int] +) -> Optional[_ChunkEntry]: + """Mandatory token-ID equality check — never trust the hash alone.""" + for entry in candidates: + if entry.token_ids == token_ids: + return entry + return None diff --git a/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_match_provider.py b/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_match_provider.py index 3c0d7867a9f5..b8af636b3a25 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_match_provider.py +++ b/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_match_provider.py @@ -132,8 +132,14 @@ def create_fuzzy_match_provider( ) return SemanticEmbeddingProvider(config) + elif provider_name == "ExactHash": + from sglang.srt.mem_cache.fuzzy_match.exact_hash_provider import ( + ExactHashProvider, + ) + + return ExactHashProvider(config) else: raise ValueError( f"Unknown fuzzy match provider: {provider_name}. " - f"Supported providers: 'SemanticEmbedding'" + f"Supported providers: 'SemanticEmbedding', 'ExactHash'" ) diff --git a/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_radix_cache.py b/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_radix_cache.py index f85163077bd9..cabc65b5426a 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_radix_cache.py +++ b/python/sglang/srt/mem_cache/fuzzy_match/fuzzy_radix_cache.py @@ -374,6 +374,26 @@ def fuzzy_match_backend_factory(ctx: TreeCacheBuildContext) -> BasePrefixCache: "--radix-cache-backend fuzzy_match does not support EAGLE " "speculative decoding yet" ) + if ctx.is_hybrid_ssm: + # FuzzyRadixCache(RadixCache) has no awareness of a model's + # separate Mamba/SSM state pool the way UnifiedRadixCache's MAMBA + # component does (registry.py's is_hybrid_ssm branch routes there + # instead, for every other backend). Registering fuzzy-matched + # content still inserts into the tree normally, but the + # invariant checker's mamba-pool accounting (which assumes + # whatever cache impl is active correctly tracks mamba slot + # protection) goes out of sync, surfacing as a confusing + # "pool memory leak detected!" crash with no fuzzy-match frames + # in the stack, found empirically running this backend against a + # real hybrid full-attention + GatedDeltaNet model. Fail loudly at + # startup instead. + raise ValueError( + "--radix-cache-backend fuzzy_match does not support hybrid " + "SSM/Mamba models yet (e.g. Qwen3.5's GatedDeltaNet layers) — " + "FuzzyRadixCache has no mamba-pool accounting, unlike " + "UnifiedRadixCache's MAMBA component that every other backend " + "routes through for these models" + ) config = FuzzyMatchConfig.from_server_args(ctx.server_args) provider = create_fuzzy_match_provider(config) cache = FuzzyRadixCache(params=ctx.params) diff --git a/python/sglang/srt/mem_cache/fuzzy_match/realizer.py b/python/sglang/srt/mem_cache/fuzzy_match/realizer.py index 0077148b9d9d..eb6ef1e3ceb4 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/realizer.py +++ b/python/sglang/srt/mem_cache/fuzzy_match/realizer.py @@ -33,7 +33,7 @@ from sglang.srt.mem_cache.fuzzy_match.rope_correction import ( copy_kv_with_rope_correction, ) -from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req @@ -47,7 +47,19 @@ class FuzzyKVRealizer: def __init__(self, req_to_token_pool, token_to_kv_pool_allocator, model): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - self.pool = token_to_kv_pool_allocator.get_kvcache() + pool = token_to_kv_pool_allocator.get_kvcache() + if isinstance(pool, HybridLinearKVPool) and not pool.use_mla: + # Hybrid full-attention + linear-attention models (e.g. Qwen3.5's + # GatedDeltaNet layers) keep the full-attention layers' KV in a + # nested pool (`full_kv_pool`), already indexed 0..N-1 scoped to + # just those layers — exactly what copy_kv_with_rope_correction's + # `range(pool.layer_num)` loop expects, with no extra layer-id + # remapping needed. The linear-attention layers have no + # per-token K to correct at all (irreversible recurrent state, + # not addressable K) and live entirely outside this nested pool, + # so there's nothing else to unwrap for them. + pool = pool.full_kv_pool + self.pool = pool # MLA-style pools have no separate K/V buffers; realization is # MHA-only for now. self.pool_supported = isinstance(self.pool, MHATokenToKVPool) @@ -59,20 +71,38 @@ def __init__(self, req_to_token_pool, token_to_kv_pool_allocator, model): ) if self.rotary_emb is None: logger.warning( - "[FUZZY] model exposes no layer-0 rotary_emb; fuzzy " + "[FUZZY] model exposes no rotary_emb on any layer; fuzzy " "realization disabled" ) @staticmethod def _resolve_rotary_emb(model): - # Model class layouts differ per architecture; probe the common - # llama-style path (model.model.layers[0].self_attn.rotary_emb). + # Model class layouts differ per architecture: some hold rotary_emb + # under a self_attn submodule (e.g. layers[i].self_attn.rotary_emb), + # others hold it directly on the decoder layer with no self_attn + # indirection at all. Hybrid architectures also mix layer types + # across model.model.layers (e.g. Qwen3.5's GatedDeltaNet layers + # expose neither), so probing layer 0 specifically is unsafe. Scan + # for the first layer that actually has one — every full-attention + # layer in a given model shares the same + # rope_theta/rotary_dim/max_position_embeddings, so any one match is + # as good as any other. inner = getattr(model, "model", None) layers = getattr(inner, "layers", None) if not layers: return None - self_attn = getattr(layers[0], "self_attn", None) - return getattr(self_attn, "rotary_emb", None) + for layer in layers: + self_attn = getattr(layer, "self_attn", None) + rotary_emb = ( + getattr(self_attn, "rotary_emb", None) + if self_attn is not None + else None + ) + if rotary_emb is None: + rotary_emb = getattr(layer, "rotary_emb", None) + if rotary_emb is not None: + return rotary_emb + return None def realize(self, fuzzy_reqs: List[Req]) -> None: """Realize every pending fuzzy match, then clear per-request state. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6061c314efe7..5b94f19140e0 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -2889,7 +2889,10 @@ class ServerArgs: ( "FuzzyMatchProvider implementation used by the fuzzy_match " "radix-cache backend. 'SemanticEmbedding' finds donor KV by " - "semantic similarity (requires the 'semblend' package)." + "semantic similarity (requires the 'semblend' package). " + "'ExactHash' finds donor KV by exact content match at a " + "different offset (content-defined chunking, no external " + "dependency) — lossless, unlike 'SemanticEmbedding'." ), NS("memory"), ] = "SemanticEmbedding" diff --git a/test/registered/fuzzy_match/test_exact_hash_e2e_safety.py b/test/registered/fuzzy_match/test_exact_hash_e2e_safety.py new file mode 100644 index 000000000000..80eb69fd1cbc --- /dev/null +++ b/test/registered/fuzzy_match/test_exact_hash_e2e_safety.py @@ -0,0 +1,350 @@ +"""E2E safety tests for ``ExactHashProvider``: hash-collision fallback and +cross-tenant isolation. Companion to +``test_exact_hash_shifted_offset_kl.py`` (which tests the *positive* +correctness case) — these test the two ways a fuzzy-match provider could +silently corrupt output if it were built carelessly: + +1. **Hash collision must fall back to a token-ID equality check, never trust + the hash alone** (``exact_hash_provider.py``'s ``_first_equal_match``). + A real 64-bit hash collision is astronomically unlikely to occur by + chance in a controlled test, so this test forces one, via the + ``SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION`` env var (collapses every + chunk's fingerprint to a constant). With two *different* donor contents + colliding into the same store bucket, a correct implementation still + disambiguates by token-ID equality; a broken one (hash-only) would either + serve the wrong donor's KV (silent corruption, catchable via KL + divergence against a fresh recompute) or crash. A third, genuinely novel + query (also forced to the same fingerprint) must fall back to a clean + miss, proving the equality check actively *rejects* false candidates, not + just "happens to usually work." + +2. **Different tenants must never share fuzzy-matched content** + (``extra_key``, already namespaced into the store key as + ``(extra_key, fingerprint)``). Uses the real, already-wired-through + ``extra_key`` field on ``/generate`` requests — no new plumbing. + +Both tests verify behavior via server log content (``[FUZZY RADIX] fuzzy +match success`` firing or not firing across a request), not just HTTP status +codes, since a silently-wrong-but-200 response and a correct response are +indistinguishable from status code alone. +""" + +import os +import random +import unittest + +from sglang.srt.environ import envs +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_test_utils import ( + _extract_output_logprobs, + _flush_cache, + _generate, + _get_input_logprobs, + compare_kl_divergence, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Same checkpoint as test_exact_hash_shifted_offset_kl.py — see that file's +# module docstring for why (PR #31057's own tested checkpoint; the -1M +# variant uses Dual Chunk Attention, not plain RoPE). +MODEL = "Qwen/Qwen2.5-7B-Instruct-AWQ" + +SINK_TOKENS = 32 # must match chunker.SINK_TOKENS +NUM_SAMPLES = 8 # matches test_exact_hash_shifted_offset_kl.py's sample size +DONOR_TOKENS = 400 +QUERY_PREFIX_TOKENS = 200 +SYNTHETIC_TOKEN_LOW = 1000 +SYNTHETIC_TOKEN_HIGH = 140000 + +ACC_THRESHOLDS = {MODEL: {"kl_div": 0.02}} + +FUZZY_SUCCESS_MARKER = "[FUZZY RADIX] fuzzy match success" + +FUZZY_ARGS = [ + "--radix-cache-backend", + "fuzzy_match", + "--fuzzy-match-provider", + "ExactHash", +] + +# One registration per file (both test classes below run under it) — see +# ci_register.py's AST-based collection, which registers the whole file per +# call, not per class. +register_cuda_ci(est_time=480, stage="base-b", runner_config="1-gpu-large") + + +def _rand_tokens(rng, n): + return [rng.randint(SYNTHETIC_TOKEN_LOW, SYNTHETIC_TOKEN_HIGH) for _ in range(n)] + + +def _generate_with_extra_key(base_url, input_ids, max_new_tokens, extra_key): + # kl_test_utils._generate has no extra_key parameter — it's a real, + # already-wired-through top-level GenerateReqInput field + # (srt/managers/io_struct.py), just not exposed by that shared helper. + import requests + + json_data = { + "input_ids": input_ids, + "extra_key": extra_key, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": max_new_tokens, + "ignore_eos": True, + }, + } + response = requests.post(base_url + "/generate", json=json_data) + response.raise_for_status() + return response.json() + + +def _log_tail(log_paths, start_pos_by_path): + """Concatenate everything written to any of ``log_paths`` since the + corresponding recorded position — logger output location (stdout vs + stderr) isn't asserted on, so both are scanned.""" + chunks = [] + for path in log_paths: + with open(path) as f: + f.seek(start_pos_by_path[path]) + chunks.append(f.read()) + return "\n".join(chunks) + + +def _log_positions(log_paths): + return {path: os.path.getsize(path) for path in log_paths} + + +class TestHashCollisionFallback(CustomTestCase): + """Forces every chunk fingerprint to collide (fingerprint=0) and checks + that ExactHashProvider still (a) serves the *correct* donor's KV when + one exists, via KL divergence against a fresh recompute, and (b) falls + back to a clean miss for content that collides but doesn't equality- + match anything registered.""" + + STDOUT_PATH = "/tmp/fuzzy_hash_collision_stdout.txt" + STDERR_PATH = "/tmp/fuzzy_hash_collision_stderr.txt" + + @classmethod + def setUpClass(cls): + cls.model = MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = open(cls.STDOUT_PATH, "w") + cls.stderr = open(cls.STDERR_PATH, "w") + with envs.SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION.override(True): + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=FUZZY_ARGS, + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + if hasattr(cls, "stdout"): + cls.stdout.close() + if hasattr(cls, "stderr"): + cls.stderr.close() + + def test_collision_disambiguates_and_rejects_false_positives(self): + # NUM_SAMPLES donors matches test_exact_hash_shifted_offset_kl.py's + # sample size — an earlier version of this test used only 2 donors + # and got a marginal KL-divergence failure (0.028 vs the 0.02 + # threshold) purely from small-sample variance (log-based checks + # all passed; a *wrong*-donor substitution would produce KL + # divergence an order of magnitude higher, not a 40% overshoot). + # Averaging over NUM_SAMPLES colliding candidates, like the sibling + # test does for non-colliding ones, is the honest fix — and is + # incidentally a *harder* disambiguation test than 2 donors, since + # every query must pick the right one out of NUM_SAMPLES colliding + # candidates in the same store bucket, not just one alternative. + rng = random.Random(20260727) + + donors = [_rand_tokens(rng, DONOR_TOKENS) for _ in range(NUM_SAMPLES)] + novel = _rand_tokens(rng, DONOR_TOKENS) # never registered + prefixes = [_rand_tokens(rng, QUERY_PREFIX_TOKENS) for _ in range(NUM_SAMPLES)] + novel_prefix = _rand_tokens(rng, QUERY_PREFIX_TOKENS) + + _flush_cache(self.base_url) + sink_prefix = list(range(50000, 50000 + SINK_TOKENS)) + + # Register NUM_SAMPLES genuinely different donors. With the forced + # collision, every one of their chunks lands in the same + # (extra_key=None, fingerprint=0) store bucket. + _generate(self.base_url, [sink_prefix + d for d in donors], max_new_tokens=0) + + # Pre-cache each query prefix standalone so the combined request + # below is a full exact match on the prefix, leaving donor/novel + # content as the entire unmatched tail (see + # test_exact_hash_shifted_offset_kl.py's docstring for why this + # exact construction is required by the current "chunks[0] only" + # limitation). + _generate(self.base_url, prefixes + [novel_prefix], max_new_tokens=0) + + log_paths = [self.STDOUT_PATH, self.STDERR_PATH] + + # One batched request for all NUM_SAMPLES donor queries plus the + # novel-content query — like test_exact_hash_shifted_offset_kl.py, + # not a Python loop of individual requests. An earlier version sent + # NUM_SAMPLES+1 separate back-to-back /generate calls and hit a + # server crash ("cannot reshape tensor of 0 elements ... [0, -1, + # 128]") inside vanilla qwen2.py RotaryEmbedding — a pre-existing + # rapid-cycling edge case unrelated to fuzzy matching (that crash + # site has no fuzzy-match frames in its stack); batching into one + # request avoids the rapid-cycling trigger entirely rather than + # working around an unrelated, pre-existing bug. + # + # This also must run *before* the KL check below — + # _get_input_logprobs calls _flush_cache internally (needs a clean + # slate for a fresh full-recompute baseline), which would wipe + # ExactHashProvider's store (on_cache_reset) and silently defeat + # this query if it ran afterward instead. + queries = [p + d for p, d in zip(prefixes, donors)] + [novel_prefix + novel] + start = _log_positions(log_paths) + results = _generate( + self.base_url, queries, max_new_tokens=64, return_logprob=True + ) + self.assertEqual(len(results), NUM_SAMPLES + 1) + log_tail = _log_tail(log_paths, start) + + # Exactly NUM_SAMPLES successes: proves every one of the + # NUM_SAMPLES real donors was found despite colliding into the same + # bucket (not fewer — a missed disambiguation), and that the novel + # query — which also collides into that bucket but equality-matches + # none of them — did *not* get a spurious match (not more). + actual_successes = log_tail.count(FUZZY_SUCCESS_MARKER) + self.assertEqual( + actual_successes, + NUM_SAMPLES, + f"expected exactly {NUM_SAMPLES} fuzzy match successes (one per " + "real donor, despite all colliding into the same forced-" + "fingerprint bucket, and none for the novel never-registered " + f"content that also collides into it), got {actual_successes}", + ) + + # --- Correctness oracle, batched (like + # test_exact_hash_shifted_offset_kl.py): if the equality check + # picked the *wrong* donor for any query above, the realized KV + # would be numerically wrong and this KL comparison would fail even + # though the log-based check above already passed. Excludes the + # novel query (never registered — a real recompute, no fuzzy + # correction to validate). --- + donor_queries = queries[:NUM_SAMPLES] + donor_results = results[:NUM_SAMPLES] + new_input_ids = [ + q + r["output_ids"] for q, r in zip(donor_queries, donor_results) + ] + output_logprobs = [_extract_output_logprobs(r) for r in donor_results] + input_logprobs = _get_input_logprobs( + self.base_url, new_input_ids, output_logprobs + ) + compare_kl_divergence( + input_logprobs, + output_logprobs, + ACC_THRESHOLDS, + self.model, + "test_collision_disambiguates_and_rejects_false_positives", + ) + + +class TestCrossTenantIsolation(CustomTestCase): + """Content registered under one extra_key must never fuzzy-match a + lookup under a different extra_key, even when the token content is + byte-identical.""" + + STDOUT_PATH = "/tmp/fuzzy_cross_tenant_stdout.txt" + STDERR_PATH = "/tmp/fuzzy_cross_tenant_stderr.txt" + + @classmethod + def setUpClass(cls): + cls.model = MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.stdout = open(cls.STDOUT_PATH, "w") + cls.stderr = open(cls.STDERR_PATH, "w") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=FUZZY_ARGS, + return_stdout_stderr=(cls.stdout, cls.stderr), + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + if hasattr(cls, "stdout"): + cls.stdout.close() + if hasattr(cls, "stderr"): + cls.stderr.close() + + def test_cross_tenant_lookup_never_matches(self): + rng = random.Random(20260728) + donor = _rand_tokens(rng, DONOR_TOKENS) + query_prefix = _rand_tokens(rng, QUERY_PREFIX_TOKENS) + + _flush_cache(self.base_url) + sink_prefix = list(range(50000, 50000 + SINK_TOKENS)) + + # Register donor content under tenant_a only. + _generate_with_extra_key( + self.base_url, [sink_prefix + donor], max_new_tokens=0, extra_key="tenant_a" + ) + + # Pre-cache the *same* query prefix standalone under both tenants' + # namespaces, so both scenarios below get an identical, full exact + # match on the prefix — the only variable is extra_key on the + # combined request. (Exact RadixCache is itself namespaced by + # extra_key — srt/mem_cache/radix_cache.py's RadixKey — so without + # this, tenant_b's combined request would get exact_matched_len=0 + # and a different, unrelated chunk-alignment path instead of a + # clean isolation test.) + _generate_with_extra_key( + self.base_url, [query_prefix], max_new_tokens=0, extra_key="tenant_a" + ) + _generate_with_extra_key( + self.base_url, [query_prefix], max_new_tokens=0, extra_key="tenant_b" + ) + + log_paths = [self.STDOUT_PATH, self.STDERR_PATH] + query = query_prefix + donor + + # --- Negative case first: a different tenant must get a clean miss + # for byte-identical content. --- + start = _log_positions(log_paths) + _generate_with_extra_key( + self.base_url, [query], max_new_tokens=8, extra_key="tenant_b" + ) + log_cross = _log_tail(log_paths, start) + self.assertNotIn( + FUZZY_SUCCESS_MARKER, + log_cross, + "tenant_b must never fuzzy-match content registered only " + "under tenant_a's extra_key", + ) + + # --- Positive control: the *same* tenant must still get the hit — + # proves the negative result above is isolation working, not the + # mechanism silently failing to fire at all. --- + start = _log_positions(log_paths) + _generate_with_extra_key( + self.base_url, [query], max_new_tokens=8, extra_key="tenant_a" + ) + log_same = _log_tail(log_paths, start) + self.assertIn( + FUZZY_SUCCESS_MARKER, + log_same, + "tenant_a must still get a fuzzy match on its own registered " + "content (positive control for the isolation test above)", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3) diff --git a/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py b/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py new file mode 100644 index 000000000000..60602054b531 --- /dev/null +++ b/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py @@ -0,0 +1,167 @@ +"""E2E correctness test: does ExactHashProvider's shifted-offset reuse +produce numerically consistent output vs. a fresh full recompute? + +This is deliberately *not* a reuse of +``test_input_output_logprobs_match_prefill_cache_hit_helper`` — that helper +tests exact-prefix content reused at the *same* offset (today's ordinary +RadixCache behavior). This test constructs the scenario ``ExactHashProvider`` +exists for: content registered as a donor at one absolute position, found +and reused when it reappears at a *different* position in a later, +unrelated prompt. + +Scenario construction is deliberate, not arbitrary, because of a real +current limitation: ``ExactHashProvider`` only checks the *first* chunk of +the unmatched tail. For that chunk to align with a registered chunk, the +content immediately after the exact-matched prefix must be chunked +identically on both the registration and query side — which requires the +registration prefix to be exactly ``SINK_TOKENS`` long (so +``chunk_region_start`` lands exactly at the donor content's start) and the +query prefix to be *fully* exact-matched beforehand (so the unmatched tail +is exactly the donor content, chunked fresh from its own start). Real, +un-arranged traffic won't reliably line up this way yet — multi-chunk / +N:M segment matching is a natural follow-up (see the provider's own +docstring). + +Donor/query content is **synthetic random token IDs**, not real dataset text +(e.g. LongBench, used by the sibling KL tests). An earlier version used real +LongBench-v2 samples and got spurious `cached_tokens` hits with no +`[FUZZY RADIX]`/`[EXACT_HASH]` log line ever firing — LongBench-v2 entries +frequently share source documents across different questions, so two +samples drawn from "different" indices in the pool aren't guaranteed to be +non-overlapping content, which is exactly what this test needs to control +for precisely. Random token IDs make non-overlap essentially certain and +are fine for a pure numerical self-consistency check — the KL comparison is +between two computations of the same thing, not a judgment about text +quality. +""" + +import random +import unittest + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kl_test_utils import ( + _extract_output_logprobs, + _flush_cache, + _generate, + _get_input_logprobs, + compare_kl_divergence, +) +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +# Qwen2.5-7B-Instruct-AWQ: the exact checkpoint PR #31057 was tested +# against. NOTE: Qwen2.5-7B-Instruct-1M (already cached locally) was tried +# first as a same-family substitute to avoid a download, on the assumption +# that it's architecturally identical modulo rope_theta — that assumption +# was wrong. The -1M variant uses Dual Chunk Attention for its extended +# context (a genuinely different attention mechanism, requiring a specific +# dual_chunk_flash_attn backend), not plain RoPE-theta scaling. Confirmed +# by actually trying to launch it, not by re-reading the config more +# carefully beforehand — worth being honest about in case this recurs. +MODEL = "Qwen/Qwen2.5-7B-Instruct-AWQ" + +SINK_TOKENS = 32 # must match chunker.SINK_TOKENS +NUM_SAMPLES = 8 +DONOR_TOKENS = 400 +QUERY_PREFIX_TOKENS = 200 +# Safe ordinary-token range for Qwen2.5's ~151.6k vocab — avoids the +# special-token IDs clustered near the top of the vocab and near 0. +SYNTHETIC_TOKEN_LOW = 1000 +SYNTHETIC_TOKEN_HIGH = 140000 + +ACC_THRESHOLDS = {MODEL: {"kl_div": 0.02}} + +register_cuda_ci(est_time=240, stage="base-b", runner_config="1-gpu-large") + + +class TestExactHashShiftedOffset(CustomTestCase): + @classmethod + def setUpClass(cls): + cls.model = MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--radix-cache-backend", + "fuzzy_match", + "--fuzzy-match-provider", + "ExactHash", + ], + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_shifted_offset_reuse_matches_full_recompute(self): + # Synthetic, guaranteed-non-overlapping content — see module + # docstring for why real dataset text doesn't work for this. + rng = random.Random(20260726) + + def rand_tokens(n): + return [ + rng.randint(SYNTHETIC_TOKEN_LOW, SYNTHETIC_TOKEN_HIGH) + for _ in range(n) + ] + + donor_contents = [rand_tokens(DONOR_TOKENS) for _ in range(NUM_SAMPLES)] + query_prefixes = [ + rand_tokens(QUERY_PREFIX_TOKENS) for _ in range(NUM_SAMPLES) + ] + + _flush_cache(self.base_url) + + # Fixed SINK_TOKENS-length prefix: makes chunk_region_start land + # exactly at donor content's start during registration (see module + # docstring). Distinct constant range from the donor/query pools, + # still within the safe ordinary-token window (not near 0, where + # low IDs risk colliding with special/reserved tokens). + sink_prefix = list(range(50000, 50000 + SINK_TOKENS)) + registration_prompts = [sink_prefix + d for d in donor_contents] + _generate(self.base_url, registration_prompts, max_new_tokens=0) + + # Pre-cache each query prefix standalone, so the combined request + # below gets a full *exact* match on it — leaving donor content as + # the entire unmatched tail, chunked fresh from its own start. + _generate(self.base_url, query_prefixes, max_new_tokens=0) + + query_prompts = [ + query_prefixes[i] + donor_contents[i] for i in range(NUM_SAMPLES) + ] + results = _generate( + self.base_url, + query_prompts, + max_new_tokens=64, + return_logprob=True, + ) + self.assertEqual(len(results), NUM_SAMPLES) + + new_input_ids = [] + output_logprobs = [] + for i, result in enumerate(results): + new_input_ids.append(query_prompts[i] + result["output_ids"]) + output_logprobs.append(_extract_output_logprobs(result)) + + input_logprobs = _get_input_logprobs( + self.base_url, new_input_ids, output_logprobs + ) + + compare_kl_divergence( + input_logprobs, + output_logprobs, + ACC_THRESHOLDS, + self.model, + "test_shifted_offset_reuse_matches_full_recompute", + ) + + +if __name__ == "__main__": + unittest.main(verbosity=3) diff --git a/test/registered/unit/mem_cache/fuzzy_match/test_chunker.py b/test/registered/unit/mem_cache/fuzzy_match/test_chunker.py new file mode 100644 index 000000000000..2d78fb72f3da --- /dev/null +++ b/test/registered/unit/mem_cache/fuzzy_match/test_chunker.py @@ -0,0 +1,119 @@ +"""Unit tests for srt/mem_cache/fuzzy_match/chunker.py — no server, no model loading.""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import random +import unittest + +from sglang.srt.environ import envs +from sglang.srt.mem_cache.fuzzy_match.chunker import ( + MAX_CHUNK_TOKENS, + MIN_CHUNK_TOKENS, + chunk_tokens, +) +from sglang.test.test_utils import CustomTestCase + + +class TestChunkTokens(CustomTestCase): + def test_content_defined_not_position_defined(self): + """Identical content at a different absolute offset must still + produce at least one matching chunk (same fingerprint, same token + IDs). This is the entire property CDC exists for — the property + a fixed-boundary / position-defined chunker provably lacks. + Regressing to fixed-boundary chunking would silently defeat the + whole mechanism without failing any obviously-related test, since + fixed boundaries still "work" in isolation. + """ + rng = random.Random(42) + shared_content = [rng.randint(1000, 50000) for _ in range(1500)] + prefix_a = [rng.randint(1000, 50000) for _ in range(50)] + prefix_b = [rng.randint(1000, 50000) for _ in range(137)] + + chunks_a = chunk_tokens(prefix_a + shared_content) + chunks_b = chunk_tokens(prefix_b + shared_content) + + by_fp_a = {c.fingerprint: c.token_ids for c in chunks_a} + by_fp_b = {c.fingerprint: c.token_ids for c in chunks_b} + shared_fingerprints = set(by_fp_a) & set(by_fp_b) + + self.assertGreater( + len(shared_fingerprints), + 0, + "expected at least one chunk boundary to align across two " + "sequences sharing 1500 tokens of content behind different-" + "length prefixes", + ) + for fp in shared_fingerprints: + self.assertEqual( + by_fp_a[fp], + by_fp_b[fp], + "same fingerprint but different token IDs — the fingerprint " + "itself must never be trusted without this equality holding " + "for the CDC layer to be sound", + ) + + def test_chunk_sizes_bounded(self): + """Every emitted chunk must respect the [MIN, MAX] clamp — an + off-by-one in the boundary/clamp bookkeeping could otherwise emit + pathologically tiny (defeats amortizing the hash lookup) or huge + (defeats fine-grained reuse) chunks silently. + """ + rng = random.Random(7) + tokens = [rng.randint(0, 100000) for _ in range(5000)] + chunks = chunk_tokens(tokens) + self.assertGreater(len(chunks), 0) + for c in chunks: + size = c.end - c.start + self.assertGreaterEqual(size, MIN_CHUNK_TOKENS) + self.assertLessEqual(size, MAX_CHUNK_TOKENS) + self.assertEqual(size, len(c.token_ids)) + + def test_chunks_tile_without_gaps_or_overlap(self): + """Chunks must exactly tile the input in order (chunk[i].end == + chunk[i+1].start), with only a final under-MIN_CHUNK_TOKENS + remainder legitimately dropped. A bug here (double-counting or + skipping tokens at a boundary) would silently register or match + content that doesn't correspond to what's actually at that offset. + """ + rng = random.Random(99) + tokens = [rng.randint(0, 100000) for _ in range(3000)] + chunks = chunk_tokens(tokens) + + self.assertEqual(chunks[0].start, 0) + for prev, nxt in zip(chunks, chunks[1:]): + self.assertEqual(prev.end, nxt.start) + remainder = len(tokens) - chunks[-1].end + self.assertLess(remainder, MIN_CHUNK_TOKENS) + + def test_force_hash_collision_env_var_collapses_fingerprints(self): + """The E2E hash-collision-fallback test + (test_exact_hash_e2e_safety.py) relies entirely on + SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION making two genuinely + different chunks collide by construction. If that wiring silently + broke (renamed env var, hook removed), the E2E test would degrade + into exercising nothing while still passing — this pins the hook's + actual effect at the unit level, independent of any server. + """ + rng = random.Random(1234) + tokens_a = [rng.randint(0, 100000) for _ in range(200)] + tokens_b = [rng.randint(0, 100000) for _ in range(200)] + self.assertNotEqual(tokens_a, tokens_b) + + with envs.SGLANG_TEST_FUZZY_FORCE_HASH_COLLISION.override(True): + fp_a = chunk_tokens(tokens_a)[0].fingerprint + fp_b = chunk_tokens(tokens_b)[0].fingerprint + self.assertEqual(fp_a, fp_b, "override(True) must force a collision") + + fp_a_real = chunk_tokens(tokens_a)[0].fingerprint + fp_b_real = chunk_tokens(tokens_b)[0].fingerprint + self.assertNotEqual( + fp_a_real, + fp_b_real, + "override must not leak past its `with` block", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py b/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py new file mode 100644 index 000000000000..e721b3c6acc7 --- /dev/null +++ b/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py @@ -0,0 +1,166 @@ +"""Unit tests for srt/mem_cache/fuzzy_match/exact_hash_provider.py. + +No server, no model loading — pure provider logic against a fake Req. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +import types +import unittest +from unittest.mock import patch + +import torch + +from sglang.srt.mem_cache.fuzzy_match.chunker import SINK_TOKENS, Chunk +from sglang.srt.mem_cache.fuzzy_match.config import FuzzyMatchConfig +from sglang.srt.mem_cache.fuzzy_match.exact_hash_provider import ExactHashProvider +from sglang.test.test_utils import CustomTestCase + + +def _fake_req(rid: str, extra_key=None): + return types.SimpleNamespace(rid=rid, extra_key=extra_key) + + +def _provider() -> ExactHashProvider: + return ExactHashProvider( + FuzzyMatchConfig(enable_fuzzy_match=True, fuzzy_match_provider="ExactHash") + ) + + +class TestExactHashProviderMatching(CustomTestCase): + def test_exact_match_at_shifted_offset(self): + """The entire point of the mechanism: content registered at one + absolute position must be found and returned (with the correct + p_src/position_offset) when the same content later appears at a + different offset. + """ + provider = _provider() + content = list(range(2000, 2000 + 200)) # well past SINK_TOKENS + donor_tokens = [0] * SINK_TOKENS + content + kv = torch.arange(len(donor_tokens)) + + ok = provider.cache_on_request_finished( + request=_fake_req("donor-1"), + token_ids=donor_tokens, + kv_cache=kv, + cache_start_pos=0, + cache_end_pos=len(donor_tokens), + ) + self.assertTrue(ok) + provider.on_donor_inserted(_fake_req("donor-1"), donor_last_node_id=42) + + # Same content, arriving at a different offset in a new prompt. + already_matched_len = 500 + prompt = list(range(9999, 9999 + already_matched_len)) + content + result = provider.match_on_prefix_miss( + prompt_token_ids=prompt, + already_matched_len=already_matched_len, + ) + + self.assertIsNotNone(result) + self.assertEqual(result.cached_start_pos, SINK_TOKENS) + self.assertEqual(result.position_offset, already_matched_len - SINK_TOKENS) + self.assertEqual(result.donor_last_node_id, 42) + self.assertEqual(result.cached_token_ids, content[: result.cached_token_count]) + + def test_sink_region_never_registered(self): + """Content whose original occurrence started inside the + attention-sink zone must never be registered as a donor: early + positions absorb a disproportionate, content-independent share of + attention regardless of what's actually there, so a content hash + match in that zone isn't a trustworthy signal. A regression here + would silently start serving sink-adjacent content as if it were a + reliable, content-only signal. + """ + provider = _provider() + # Entirely within the sink zone. + donor_tokens = list(range(SINK_TOKENS)) + kv = torch.arange(len(donor_tokens)) + + ok = provider.cache_on_request_finished( + request=_fake_req("donor-sink"), + token_ids=donor_tokens, + kv_cache=kv, + cache_start_pos=0, + cache_end_pos=len(donor_tokens), + ) + self.assertFalse(ok) + self.assertEqual(len(provider._store), 0) + + def test_cross_tenant_isolation(self): + """Identical content registered under one tenant's extra_key must + never match a lookup under a different tenant's extra_key — this + provider must enforce multi-tenant isolation unconditionally. A + regression here is a real cross-tenant KV leak. + """ + provider = _provider() + content = list(range(3000, 3000 + 200)) + donor_tokens = [0] * SINK_TOKENS + content + kv = torch.arange(len(donor_tokens)) + + provider.cache_on_request_finished( + request=_fake_req("donor-2", extra_key="tenant_a"), + token_ids=donor_tokens, + kv_cache=kv, + cache_start_pos=0, + cache_end_pos=len(donor_tokens), + ) + + already_matched_len = 10 + prompt = list(range(10)) + content + result = provider.match_on_prefix_miss( + prompt_token_ids=prompt, + already_matched_len=already_matched_len, + extra_key="tenant_b", + ) + self.assertIsNone(result) + + # Sanity: the same lookup under the correct tenant does match. + result_same_tenant = provider.match_on_prefix_miss( + prompt_token_ids=prompt, + already_matched_len=already_matched_len, + extra_key="tenant_a", + ) + self.assertIsNotNone(result_same_tenant) + + def test_hash_collision_falls_back_to_miss(self): + """Never trust the hash alone: two different chunks forced to share + a fingerprint must not produce a match — the mandatory token-ID + equality check is what actually guards correctness here, not the + fingerprint's collision odds. + """ + provider = _provider() + donor_content = [111] * 200 + query_content = [222] * 200 # different content, forced same fingerprint + + def fake_chunks(tokens): + return [Chunk(start=0, end=len(tokens), token_ids=list(tokens), fingerprint=1)] + + with patch( + "sglang.srt.mem_cache.fuzzy_match.exact_hash_provider.chunk_tokens", + side_effect=fake_chunks, + ): + donor_tokens = [0] * SINK_TOKENS + donor_content + provider.cache_on_request_finished( + request=_fake_req("donor-3"), + token_ids=donor_tokens, + kv_cache=torch.arange(len(donor_tokens)), + cache_start_pos=0, + cache_end_pos=len(donor_tokens), + ) + + result = provider.match_on_prefix_miss( + prompt_token_ids=[0] * 10 + query_content, + already_matched_len=10, + ) + self.assertIsNone( + result, + "fingerprint collided by construction but content differs — " + "must fall back to miss, not serve the wrong donor", + ) + + +if __name__ == "__main__": + unittest.main() From 6338fed5ca0a9bbb38f35483c0e9153d8f5deb71 Mon Sep 17 00:00:00 2001 From: krakhit <61501745+krakhit@users.noreply.github.com> Date: Wed, 29 Jul 2026 10:33:24 +0000 Subject: [PATCH 2/2] format --- .../srt/mem_cache/fuzzy_match/README.md | 47 ++++++++++++++----- .../srt/mem_cache/fuzzy_match/chunker.py | 8 +--- .../test_exact_hash_shifted_offset_kl.py | 7 +-- .../fuzzy_match/test_exact_hash_provider.py | 4 +- 4 files changed, 41 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/mem_cache/fuzzy_match/README.md b/python/sglang/srt/mem_cache/fuzzy_match/README.md index a5fa65088376..d53170a9c503 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/README.md +++ b/python/sglang/srt/mem_cache/fuzzy_match/README.md @@ -1,12 +1,21 @@ # Fuzzy KV Cache Reuse (`--radix-cache-backend fuzzy_match`) -Semantic KV cache reuse for prompts that share meaning but not tokens. -When exact prefix matching leaves part of a prompt uncovered, a pluggable -`FuzzyMatchProvider` may nominate donor KV from a previously finished -request; the donor KV is position-corrected (RoPE) into recipient-owned -slots before the forward pass. Reuse follows the `|exact|fuzzy|miss|` -prompt decomposition: one contiguous fuzzy span anchored at the exact -prefix boundary. +KV cache reuse for prompts whose content reappears at a different offset +than where it was originally computed. When exact prefix matching leaves +part of a prompt uncovered, a pluggable `FuzzyMatchProvider` may nominate +donor KV from a previously finished request; the donor KV is +position-corrected (RoPE) into recipient-owned slots before the forward +pass. Reuse follows the `|exact|fuzzy|miss|` prompt decomposition: one +contiguous fuzzy span anchored at the exact prefix boundary. + +Two providers are available: +- `SemanticEmbedding` (default): matches merely-similar content by cosine + similarity. Requires the `semblend` package. Not lossless — reuse changes + model outputs by construction (see Scope and guarantees below). +- `ExactHash`: matches content that is byte-identical to the current + prompt's unmatched tail but sits at a different offset (content-defined + chunking, no external dependency). Lossless: every hash hit is confirmed + by a token-ID equality check before being served. ## Enabling @@ -20,6 +29,15 @@ python -m sglang.launch_server \ --fuzzy-model-arch qwen2.5-7b ``` +Or, with no external dependency: + +```bash +python -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-7B-Instruct-AWQ \ + --radix-cache-backend fuzzy_match \ + --fuzzy-match-provider ExactHash +``` + Selecting the backend enables the feature; the default provider is `SemanticEmbedding`. Reuse markers in the server log: `fuzzy match success` (match accepted), `[FUZZY] Realized N fuzzy tokens` (KV copied + @@ -98,7 +116,7 @@ detectably stale — never dangling. | Flag | Default | Why it exists | |---|---|---| | `--radix-cache-backend fuzzy_match` | off | The enable switch; registers nothing and costs nothing when unset. | -| `--fuzzy-match-provider` | `SemanticEmbedding` | Provider selection; the interface admits out-of-tree providers. | +| `--fuzzy-match-provider` | `SemanticEmbedding` | `SemanticEmbedding` or `ExactHash`; the interface admits out-of-tree providers too. | | `--fuzzy-semantic-threshold` | `0.60` | Precision knob: cosine floor for accepting a donor. | | `--fuzzy-min-reuse-ratio` | `0.50` | Hit gate: donors covering less of the prompt are rejected. | | `--fuzzy-min-match-length` | `16` | Skips fuzzy lookup behind weak partial exact anchors. | @@ -115,10 +133,13 @@ flags. the missed suffix. The default backend path is byte-identical when the backend is not selected (two seams: one `MatchResult` field, one no-op hook in `cache_finished_req`). -- Reuse changes model outputs by construction: donor K/V attended to the - donor's context. The provider's quality gates plus the per-layer - zero-out mask bound the drift; accuracy methodology and results are in - the PR description. -- Not yet supported: MLA-style KV pools, EAGLE speculative decoding, +- `SemanticEmbedding` reuse changes model outputs by construction: donor + K/V attended to the donor's context. The provider's quality gates plus + the per-layer zero-out mask bound the drift. `ExactHash` reuse is + lossless by construction — matched content is byte-identical to the + current prompt, confirmed by a mandatory token-ID equality check on + every hit. Accuracy methodology and results for both are in the PR + descriptions that introduced them. +- Not yet supported (both providers): MLA-style KV pools, EAGLE speculative decoding, multi-region (`|exact|miss|fuzzy|miss|...`) reuse, hierarchical (host) cache interaction. Each is rejected explicitly rather than silently. diff --git a/python/sglang/srt/mem_cache/fuzzy_match/chunker.py b/python/sglang/srt/mem_cache/fuzzy_match/chunker.py index aaecc5dc8912..aa861d3776b1 100644 --- a/python/sglang/srt/mem_cache/fuzzy_match/chunker.py +++ b/python/sglang/srt/mem_cache/fuzzy_match/chunker.py @@ -89,9 +89,7 @@ def _real_fingerprint(token_ids: Sequence[int]) -> int: # gives the same "cheap, non-cryptographic-role" fingerprint # xxHash64 would — the token-ID equality check downstream is what # actually guards correctness, not this hash's collision odds. - digest = hashlib.blake2b( - _tokens_to_bytes(token_ids), digest_size=8 - ).digest() + digest = hashlib.blake2b(_tokens_to_bytes(token_ids), digest_size=8).digest() return int.from_bytes(digest, "little") @@ -141,9 +139,7 @@ def chunk_tokens(token_ids: Sequence[int]) -> List[Chunk]: chunk_len = i - chunk_start + 1 if chunk_len < MIN_CHUNK_TOKENS: continue - at_boundary = (i - chunk_start + 1 >= WINDOW_SIZE) and ( - h & _BOUNDARY_MASK == 0 - ) + at_boundary = (i - chunk_start + 1 >= WINDOW_SIZE) and (h & _BOUNDARY_MASK == 0) if at_boundary or chunk_len >= MAX_CHUNK_TOKENS: span = token_ids[chunk_start : i + 1] chunks.append( diff --git a/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py b/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py index 60602054b531..643a2fc868d1 100644 --- a/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py +++ b/test/registered/fuzzy_match/test_exact_hash_shifted_offset_kl.py @@ -108,14 +108,11 @@ def test_shifted_offset_reuse_matches_full_recompute(self): def rand_tokens(n): return [ - rng.randint(SYNTHETIC_TOKEN_LOW, SYNTHETIC_TOKEN_HIGH) - for _ in range(n) + rng.randint(SYNTHETIC_TOKEN_LOW, SYNTHETIC_TOKEN_HIGH) for _ in range(n) ] donor_contents = [rand_tokens(DONOR_TOKENS) for _ in range(NUM_SAMPLES)] - query_prefixes = [ - rand_tokens(QUERY_PREFIX_TOKENS) for _ in range(NUM_SAMPLES) - ] + query_prefixes = [rand_tokens(QUERY_PREFIX_TOKENS) for _ in range(NUM_SAMPLES)] _flush_cache(self.base_url) diff --git a/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py b/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py index e721b3c6acc7..40dbfa30c44a 100644 --- a/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py +++ b/test/registered/unit/mem_cache/fuzzy_match/test_exact_hash_provider.py @@ -136,7 +136,9 @@ def test_hash_collision_falls_back_to_miss(self): query_content = [222] * 200 # different content, forced same fingerprint def fake_chunks(tokens): - return [Chunk(start=0, end=len(tokens), token_ids=list(tokens), fingerprint=1)] + return [ + Chunk(start=0, end=len(tokens), token_ids=list(tokens), fingerprint=1) + ] with patch( "sglang.srt.mem_cache.fuzzy_match.exact_hash_provider.chunk_tokens",