From caf29d31cae03177e6dbdfb45f5c1ad7e3d5fe3f Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Mon, 19 Oct 2026 22:57:38 -0700 Subject: [PATCH 1/5] kvcr-foundation(endpoints,non_mechanical_provable): validate rank endpoints Signed-off-by: Ziqi Fan --- .../srt/mem_cache/storage/kvcr/__init__.py | 2 + .../srt/mem_cache/storage/kvcr/endpoints.py | 42 ++++++++++ .../unit/mem_cache/test_kvcr_endpoints.py | 82 +++++++++++++++++++ 3 files changed, 126 insertions(+) create mode 100644 python/sglang/srt/mem_cache/storage/kvcr/__init__.py create mode 100644 python/sglang/srt/mem_cache/storage/kvcr/endpoints.py create mode 100644 test/registered/unit/mem_cache/test_kvcr_endpoints.py diff --git a/python/sglang/srt/mem_cache/storage/kvcr/__init__.py b/python/sglang/srt/mem_cache/storage/kvcr/__init__.py new file mode 100644 index 000000000000..a864f5018474 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional KVCR integration primitives; importing this package does not load KVCR.""" diff --git a/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py b/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py new file mode 100644 index 000000000000..bd430d5f73d5 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Rank-local control endpoints shared by KVCR storage and direct delivery. + +Only callers with an explicitly homogeneous source topology may apply a rank +offset. A router hint is advisory; endpoint validation does not establish a +source lease or prove that a transfer completed. +""" + +from urllib.parse import urlsplit + + +def rank_offset(config) -> int: + """TP is the innermost rank within one advertised DP replica.""" + return config.tp_rank + config.tp_size * ( + config.attn_cp_rank + config.attn_cp_size * config.pp_rank + ) + + +def rank_endpoint(endpoint: str, offset: int) -> str | None: + """Offset an advertised TCP endpoint, or reject an ambiguous address.""" + if not isinstance(endpoint, str) or type(offset) is not int or offset < 0: + return None + try: + source = urlsplit(endpoint) + port = source.port + if ( + source.scheme != "tcp" + or not source.hostname + or port is None + or not 0 < port <= port + offset <= 65535 + or source.path + or source.query + or source.fragment + or source.username is not None + ): + return None + except ValueError: + return None + host = source.hostname + if ":" in host: + host = f"[{host}]" + return f"tcp://{host}:{port + offset}" diff --git a/test/registered/unit/mem_cache/test_kvcr_endpoints.py b/test/registered/unit/mem_cache/test_kvcr_endpoints.py new file mode 100644 index 000000000000..bec7eac5351f --- /dev/null +++ b/test/registered/unit/mem_cache/test_kvcr_endpoints.py @@ -0,0 +1,82 @@ +"""One source replica's rank ports must not cross endpoint or topology bounds.""" + +import subprocess +import sys +import unittest +from types import SimpleNamespace + +from sglang.srt.mem_cache.storage.kvcr import endpoints +from sglang.srt.mem_cache.storage.kvcr.endpoints import rank_endpoint, rank_offset +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestKVCRRankEndpoints(CustomTestCase): + def test_shard_offsets_are_unique_within_one_replica(self): + config = SimpleNamespace(tp_size=4, attn_cp_size=2) + offsets = [] + for pp in range(2): + config.pp_rank = pp + for cp in range(2): + config.attn_cp_rank = cp + for tp in range(4): + config.tp_rank = tp + offsets.append(rank_offset(config)) + self.assertEqual(offsets, list(range(16))) + self.assertEqual( + rank_endpoint("tcp://source.example:32100", offsets[-1]), + "tcp://source.example:32115", + ) + + def test_ipv6_host_and_zero_offset_remain_unambiguous(self): + for source, offset, expected in ( + ("tcp://127.0.0.1:32100", 0, "tcp://127.0.0.1:32100"), + ("tcp://source.example:32100", 2, "tcp://source.example:32102"), + ("tcp://[2001:db8::1]:32100", 2, "tcp://[2001:db8::1]:32102"), + ("tcp://source.example:65535", 0, "tcp://source.example:65535"), + ): + with self.subTest(source=source, offset=offset): + self.assertEqual(rank_endpoint(source, offset), expected) + + def test_invalid_addresses_cannot_select_another_source(self): + for source in ( + "http://source.example:32100", + "tcp://source.example", + "tcp://source.example:bad", + "tcp://source.example:0", + "tcp://source.example:65536", + "tcp://source.example:32100/path", + "tcp://source.example:32100?query", + "tcp://source.example:32100#fragment", + "tcp://user:password@source.example:32100", + "tcp://[2001:db8::1:32100", + "tcp://:32100", + None, + ): + with self.subTest(source=source): + self.assertIsNone(rank_endpoint(source, 1)) + self.assertIsNone(rank_endpoint("tcp://source.example:65535", 1)) + for offset in (-1, True, 0.5, "1"): + with self.subTest(offset=offset): + self.assertIsNone(rank_endpoint("tcp://source.example:32100", offset)) + + def test_shared_helpers_do_not_import_the_optional_runtime(self): + result = subprocess.run( + [ + sys.executable, + "-c", + "import runpy,sys; sys.modules['kvcr']=None; " + "runpy.run_path(sys.argv[1])", + endpoints.__file__, + ], + capture_output=True, + text=True, + timeout=10, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + +if __name__ == "__main__": + unittest.main() From e17419bd4a3e8d6783202359fe6eab1ad41d2183 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Mon, 19 Oct 2026 22:59:34 -0700 Subject: [PATCH 2/5] kvcr-foundation(inverse-port,non_mechanical_provable): retain reverse offsets Signed-off-by: Ziqi Fan --- python/sglang/srt/mem_cache/storage/kvcr/endpoints.py | 5 +++-- test/registered/unit/mem_cache/test_kvcr_endpoints.py | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py b/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py index bd430d5f73d5..6bd504f7d616 100644 --- a/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py +++ b/python/sglang/srt/mem_cache/storage/kvcr/endpoints.py @@ -18,7 +18,7 @@ def rank_offset(config) -> int: def rank_endpoint(endpoint: str, offset: int) -> str | None: """Offset an advertised TCP endpoint, or reject an ambiguous address.""" - if not isinstance(endpoint, str) or type(offset) is not int or offset < 0: + if not isinstance(endpoint, str) or type(offset) is not int: return None try: source = urlsplit(endpoint) @@ -27,7 +27,8 @@ def rank_endpoint(endpoint: str, offset: int) -> str | None: source.scheme != "tcp" or not source.hostname or port is None - or not 0 < port <= port + offset <= 65535 + or not 1 <= port <= 65535 + or not 1 <= port + offset <= 65535 or source.path or source.query or source.fragment diff --git a/test/registered/unit/mem_cache/test_kvcr_endpoints.py b/test/registered/unit/mem_cache/test_kvcr_endpoints.py index bec7eac5351f..a73cc6335263 100644 --- a/test/registered/unit/mem_cache/test_kvcr_endpoints.py +++ b/test/registered/unit/mem_cache/test_kvcr_endpoints.py @@ -34,6 +34,7 @@ def test_ipv6_host_and_zero_offset_remain_unambiguous(self): for source, offset, expected in ( ("tcp://127.0.0.1:32100", 0, "tcp://127.0.0.1:32100"), ("tcp://source.example:32100", 2, "tcp://source.example:32102"), + ("tcp://source.example:32102", -2, "tcp://source.example:32100"), ("tcp://[2001:db8::1]:32100", 2, "tcp://[2001:db8::1]:32102"), ("tcp://source.example:65535", 0, "tcp://source.example:65535"), ): @@ -58,7 +59,8 @@ def test_invalid_addresses_cannot_select_another_source(self): with self.subTest(source=source): self.assertIsNone(rank_endpoint(source, 1)) self.assertIsNone(rank_endpoint("tcp://source.example:65535", 1)) - for offset in (-1, True, 0.5, "1"): + self.assertIsNone(rank_endpoint("tcp://source.example:1", -1)) + for offset in (True, 0.5, "1"): with self.subTest(offset=offset): self.assertIsNone(rank_endpoint("tcp://source.example:32100", offset)) From 17ab9da0dd197b66abf9e698b533073f251374c1 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Tue, 20 Oct 2026 09:06:24 -0700 Subject: [PATCH 3/5] feat(cache): add attempt-scoped linker hooks and owned inventory events Port the generic parts of Moein and Lin linker work to the current request-private loading lifecycle. Keep GPU and storage inventory ownership separate, preserve event namespaces after eviction, and confirm sparse prefix boundaries without early tree publication. This slice adds no KVCR runtime or backend. Co-authored-by: Moein Khazraee Co-authored-by: Lin Hu --- python/sglang/srt/disaggregation/kv_events.py | 4 +- python/sglang/srt/mem_cache/events.py | 3 +- .../srt/mem_cache/rust_tree_core/adapter.py | 12 + .../unified_cache/unified_cache_linker.py | 138 +++++- .../unified_cache/unified_tree_core.py | 75 +++- .../unified_tree_core_interface.py | 12 + .../srt/mem_cache/unified_radix_cache.py | 31 +- rust/sglang-radix-tree/src/python_bindings.rs | 47 ++- .../src/tests/unified_tree_core.rs | 100 +++++ .../src/unified_tree_core.rs | 83 ++++ .../unit/disaggregation/test_kv_events.py | 19 +- .../mem_cache/test_unified_cache_linker.py | 393 ++++++++++++++++-- 12 files changed, 861 insertions(+), 56 deletions(-) diff --git a/python/sglang/srt/disaggregation/kv_events.py b/python/sglang/srt/disaggregation/kv_events.py index 84d3855034c7..00f5f387471b 100644 --- a/python/sglang/srt/disaggregation/kv_events.py +++ b/python/sglang/srt/disaggregation/kv_events.py @@ -293,15 +293,17 @@ class BlockStored(KVCacheEvent): # Public adapter name for router lookup. Published hashes are opaque IDs # scoped to the internal cache namespace, including the adapter load. lora_name: Optional[str] = None + ownership: Optional[str] = None class BlockRemoved(KVCacheEvent): block_hashes: list[int] medium: Optional[str] = None + ownership: Optional[str] = None class AllBlocksCleared(KVCacheEvent): - pass + ownership: Optional[str] = None # The dynamo format publishes namespaced hashes as block hashes. If you change diff --git a/python/sglang/srt/mem_cache/events.py b/python/sglang/srt/mem_cache/events.py index 861d8589d24e..d08b3269aab8 100644 --- a/python/sglang/srt/mem_cache/events.py +++ b/python/sglang/srt/mem_cache/events.py @@ -148,13 +148,14 @@ def enqueue(self, event) -> None: tail = self._queue[-1] if isinstance(tail, BlockRemoved) and isinstance(event, BlockRemoved): - if tail.medium == event.medium: + if tail.medium == event.medium and tail.ownership == event.ownership: tail.block_hashes.extend(event.block_hashes) return elif isinstance(tail, BlockStored) and isinstance(event, BlockStored): if ( tail.medium == event.medium + and tail.ownership == event.ownership and tail.lora_id == event.lora_id and tail.block_size == event.block_size and tail.cache_salt == event.cache_salt diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index 2f44ea44866d..73f8229ff316 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -1187,6 +1187,18 @@ def finish_external_linker_offload( ) -> None: self._binding.finish_external_linker_offload(list(node_ids), ack_id, success) + def record_external_cache_store( + self, node_ids: Sequence[NodeId], medium: StorageMedium + ) -> None: + self._binding.record_external_cache_store(list(node_ids), medium.value) + + def remove_external_cache_storage( + self, hashes: Sequence[str], medium: StorageMedium, *, publish: bool = True + ) -> None: + self._binding.remove_external_cache_storage( + list(hashes), medium.value if publish else None + ) + @property def write_back_duplicate_reclaim_digest(self) -> int: return self._binding.write_back_duplicate_reclaim_digest() diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py b/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py index 64f341e17857..ffaf31e3423f 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py @@ -30,14 +30,13 @@ import torch +from sglang.srt.disaggregation.kv_events import StorageMedium from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, DecLockRefParams, MatchResult, ) -from sglang.srt.mem_cache.hicache_storage import ( - PoolName, - PoolTransfer, -) +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.components import ( ComponentType, @@ -48,6 +47,7 @@ from sglang.srt.mem_cache.utils import get_storage_hash_str if TYPE_CHECKING: + from sglang.srt.managers.kv_hints import KvHintsEnvelope from sglang.srt.managers.schedule_batch import Req from sglang.srt.mem_cache.unified_cache.unified_tree_core_interface import NodeId from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache @@ -133,6 +133,21 @@ class UnifiedCacheLinker(ABC): layer_done_counter: object pools: dict[PoolName, object] + local_storage_medium: StorageMedium | None = None + local_storage_ownership: str | None = None + requires_lookup_confirmation: bool = False + + def set_request_context( + self, handle: CacheRequestHandle, kv_hints: KvHintsEnvelope | None + ) -> None: + """Snapshot attempt-scoped hints; accepted loads retain their own scope.""" + + def release_request_context(self, handle: CacheRequestHandle) -> None: + """Release request scope without revoking accepted transfers.""" + + def pop_storage_removals(self) -> list[str]: + """Consume storage hashes whose local KVCR copies were removed.""" + return [] @abstractmethod def lookup(self, rid: str, transfers: list[PoolTransfer]) -> list[int]: @@ -146,8 +161,15 @@ def lookup(self, rid: str, transfers: list[PoolTransfer]) -> list[int]: let the tree pick a length that is invalid on another rank. Local to this rank; the tree intersects the sets across ranks. + Backends requiring confirmation may return optimistic candidates. """ + def confirm_lookup( + self, rid: str, transfers: list[PoolTransfer], hit_pages: int + ) -> bool: + """Confirm a shared boundary; zero discards unaccepted lookup claims.""" + return True + @abstractmethod def load(self, rid: str, transfers: list[PoolTransfer]) -> bool: """Queue a load into the given device indices. @@ -237,6 +259,8 @@ def __init__( ) # rid -> what match found, consumed by the next init_load_back. self.hit_markers: dict[str, ExternalCacheHitMarker] = {} + self._request_handles: dict[str, CacheRequestHandle] = {} + self._load_request_handles: dict[str, CacheRequestHandle] = {} # Requests loading in the batch being built, then its counter index once # the batch starts; finish_loads consumes both after its forward. self.inflight_load_rids: list[str] = [] @@ -259,6 +283,9 @@ def has_hit(self, rid: str) -> bool: # ---- match: probe the remote store and report host_hit_length ---- def match(self, key: RadixKey, req: Req, result: MatchResult) -> MatchResult: + previous = self.hit_markers.pop(req.rid, None) + if previous is not None and self.cache_linker.requires_lookup_confirmation: + self.cache_linker.confirm_lookup(req.rid, [], 0) cache = self.cache key, _ = key.maybe_to_bigram_view(cache.tree_core.is_eagle) page = cache.page_size @@ -280,11 +307,22 @@ def match(self, key: RadixKey, req: Req, result: MatchResult) -> MatchResult: lookup_transfers.append(transfer) by_pool = {transfer.name: transfer for transfer in lookup_transfers} + handle = req.cache_request_handle + previous = self._request_handles.get(req.rid) + if previous is not None and previous != handle: + self.release_request(previous, cancel_load=False) + self._request_handles[req.rid] = handle + self.cache_linker.set_request_context(handle, req.kv_hints) + + def confirm(pages): + return self.cache_linker.confirm_lookup(req.rid, lookup_transfers, pages) + # Tail-relative: page 0 of `tail_hashes` is the first uncached page. hit_pages = self._sync_restorable_prefix( self.cache_linker.lookup(req.rid, lookup_transfers), num_pages=len(tail_hashes), device_hit_pages=0, + confirm=confirm if self.cache_linker.requires_lookup_confirmation else None, ) if hit_pages == 0: return result @@ -313,7 +351,12 @@ def match(self, key: RadixKey, req: Req, result: MatchResult) -> MatchResult: ) def _sync_restorable_prefix( - self, restorable: list[int], *, num_pages: int, device_hit_pages: int + self, + restorable: list[int], + *, + num_pages: int, + device_hit_pages: int, + confirm: Callable[[int], bool] | None = None, ) -> int: """Intersect the per-rank sets of restorable prefix lengths and return the longest one, or 0 when the ranks share none beyond the device prefix. @@ -328,9 +371,15 @@ def _sync_restorable_prefix( mask[pages] = 1 self.cache._all_reduce_attn_groups(mask, torch.distributed.ReduceOp.MIN) common = mask.nonzero() - if common.numel() == 0: - return 0 - return int(common[-1].item()) + if confirm is None: + return int(common[-1].item()) if common.numel() else 0 + for pages in reversed(common.flatten().tolist()): + accepted = torch.tensor([int(confirm(pages))], dtype=torch.int) + self.cache._all_reduce_attn_groups(accepted, torch.distributed.ReduceOp.MIN) + if accepted.item(): + return pages + confirm(0) + return 0 def _tail_hashes( self, key: RadixKey, result: MatchResult, device_hit_len: int @@ -385,6 +434,8 @@ def load_back(self, req: Req) -> tuple[int, NodeId]: component_transfers, prefix_len, ) + if self.cache_linker.requires_lookup_confirmation: + self.cache_linker.confirm_lookup(req.rid, [], 0) return 0, req.last_node component_transfers.append((component, transfer)) @@ -406,6 +457,8 @@ def load_back(self, req: Req) -> tuple[int, NodeId]: self._update_load( ExternalLinkerLoadPhase.ABORT, req, component_transfers, prefix_len ) + if self.cache_linker.requires_lookup_confirmation: + self.cache_linker.confirm_lookup(req.rid, [], 0) return 0, req.last_node self._update_load( ExternalLinkerLoadPhase.PREPARE, @@ -436,6 +489,8 @@ def load_back(self, req: Req) -> tuple[int, NodeId]: req.swa_branching_seqlen = None self.private_load_indices[req.rid] = full_transfer.device_indices self.inflight_load_rids.append(req.rid) + if req.rid in self._request_handles: + self._load_request_handles[req.rid] = self._request_handles[req.rid] return len(full_transfer.device_indices), req.last_node def _update_load( @@ -523,8 +578,44 @@ def commit_completed_offloads(self, successes: Sequence[bool]) -> None: self.cache.tree_core.finish_external_linker_offload( pending.publish_node_ids, pending.lock_node_id, success ) + if success and self.cache_linker.local_storage_medium is not None: + self.cache.tree_core.record_external_cache_store( + pending.publish_node_ids, self.cache_linker.local_storage_medium + ) self.cache.dec_lock_ref(pending.lock_node_id, pending.lock_params) + def drain_storage_removals(self) -> None: + medium = self.cache_linker.local_storage_medium + if medium is None: + return + hashes = self.cache_linker.pop_storage_removals() + any_removed = torch.tensor([bool(hashes)], dtype=torch.int) + self.cache._all_reduce_attn_groups(any_removed, torch.distributed.ReduceOp.MAX) + if not any_removed.item(): + return + all_hashes = set(hashes) + groups = [ + group + for group in (self.cache.attn_cp_group, self.cache.attn_tp_group) + if group is not None and torch.distributed.get_world_size(group=group) > 1 + ] + if not groups and self.cache.tp_world_size > 1: + groups = [self.cache.tp_group] + for group in groups: + gathered = [None] * torch.distributed.get_world_size(group=group) + torch.distributed.all_gather_object( + gathered, sorted(all_hashes), group=group + ) + all_hashes.update(key for rank_hashes in gathered for key in rank_hashes) + # Offload decisions must agree, but each rank advertises its own inventory. + remote_hashes = all_hashes.difference(hashes) + if remote_hashes: + self.cache.tree_core.remove_external_cache_storage( + sorted(remote_hashes), medium, publish=False + ) + if hashes: + self.cache.tree_core.remove_external_cache_storage(hashes, medium) + def start_layer_wise_loading(self) -> int: self.inflight_load_index = self.cache_linker.start_layer_wise_loading() return self.inflight_load_index @@ -550,6 +641,7 @@ def finish_loads(self) -> list[str]: self.cache._all_reduce_attn_groups(landed, torch.distributed.ReduceOp.MIN) for rid in rids: self.private_load_indices.pop(rid, None) + self._load_request_handles.pop(rid, None) if landed.item(): return [] logger.error( @@ -566,6 +658,8 @@ def reset(self) -> None: self.private_load_indices.clear() self.inflight_load_index = -1 self._release_pending_offloads() + self._load_request_handles.clear() + self._release_request_contexts() def _release_pending_offloads(self) -> None: for pending in self.pending_offloads: @@ -575,16 +669,40 @@ def _release_pending_offloads(self) -> None: self.cache.dec_lock_ref(pending.lock_node_id, pending.lock_params) self.pending_offloads.clear() - def release_request(self, rid: str) -> None: + def release_request( + self, handle: CacheRequestHandle, *, cancel_load: bool = True + ) -> None: + rid = handle.rid + current = self._request_handles.get(rid) + if current is not None and current != handle: + return self.hit_markers.pop(rid, None) # A queued load targets slots the request owns and is about to free. # A started one is finished by finish_loads before anything is freed. - if rid in self.inflight_load_rids and self.cache_linker.cancel_queued_load(rid): + owns_load = self._load_request_handles.get(rid, handle) == handle + if ( + cancel_load + and owns_load + and rid in self.inflight_load_rids + and self.cache_linker.cancel_queued_load(rid) + ): self.inflight_load_rids.remove(rid) self.private_load_indices.pop(rid, None) + self._load_request_handles.pop(rid, None) + if current is not None: + self.cache_linker.release_request_context(handle) + self._request_handles.pop(rid) + + def _release_request_contexts(self) -> None: + for handle in self._request_handles.values(): + self.cache_linker.release_request_context(handle) + self._request_handles.clear() def close(self) -> None: self.cache_linker.close() + self.hit_markers.clear() self.inflight_load_rids.clear() self.private_load_indices.clear() self._release_pending_offloads() + self._load_request_handles.clear() + self._release_request_contexts() diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 0714b191ae5d..9f2d888a0b7a 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -26,7 +26,11 @@ import msgspec import torch -from sglang.srt.disaggregation.kv_events import StorageMedium +from sglang.srt.disaggregation.kv_events import ( + BlockRemoved, + StorageMedium, + namespaced_block_hash, +) from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, @@ -84,6 +88,7 @@ from sglang.srt.mem_cache.utils import ( compute_node_hash_values, get_eviction_strategy, + hash_str_to_int64, split_node_hash_value, ) @@ -466,6 +471,8 @@ def __init__( dynamo_format=params.dynamo_kv_event_format, lora_names=params.kv_event_lora_names, ) + # Local storage outlives device tree nodes and engine cache resets. + self._external_cache_event_hashes: dict[str, int] = {} self.reset() @@ -729,9 +736,9 @@ def inc_lock_ref( def _assert_receipt_anchor(node: UnifiedTreeNode, params: DecLockRefParams) -> None: """A receipt releases only the node its acquire returned; a mispaired or unanchored one would silently release (or steal) another holder's segment.""" - assert params.node_id == node.id, ( - f"lock receipt anchored on node {params.node_id} released on node {node.id}" - ) + assert ( + params.node_id == node.id + ), f"lock receipt anchored on node {params.node_id} released on node {node.id}" def _release_components( self, @@ -1688,12 +1695,12 @@ def _finish_internal_component_eviction( ) -> EvictDeviceNextNodeResult: """Resume an internal tombstone after the controller's backup attempt.""" component = self.components_by_type[component_type] - assert component.is_evict_device_ongoing, ( - f"{component_type.name} device eviction not started" - ) - assert self._pending_internal_evictions.get(component_type) == node_id, ( - f"no matching pending internal {component_type.name} eviction" - ) + assert ( + component.is_evict_device_ongoing + ), f"{component_type.name} device eviction not started" + assert ( + self._pending_internal_evictions.get(component_type) == node_id + ), f"no matching pending internal {component_type.name} eviction" del self._pending_internal_evictions[component_type] # Consuming the pending request advances the walk even if I/O changed # this victim's eligibility. No frees occur before backup completion. @@ -2753,6 +2760,48 @@ def finish_external_linker_offload( node.write_through_pending_id = None node.external_cache_stored |= success + def record_external_cache_store( + self, node_ids: Sequence[NodeId], medium: StorageMedium + ) -> None: + if not self.kv_events.enabled: + return + for node_id in node_ids: + node = self.node_by_id(node_id) + self.kv_events.record_store(node, medium=medium) + self._external_cache_event_hashes.update( + zip( + node.hash_value, + ( + namespaced_block_hash( + hash_str_to_int64(value), + namespace_seed=self.kv_events._namespace_seed(node), + ) + for value in self.kv_events._node_event_hash_values(node) + ), + ) + ) + + def remove_external_cache_storage( + self, hashes: Sequence[str], medium: StorageMedium, *, publish: bool = True + ) -> None: + removed = set(hashes) + for node in self._node_arena.values(): + if node.external_cache_stored and removed.intersection( + node.hash_value or () + ): + node.external_cache_stored = False + if not publish: + return + block_hashes = [ + self._external_cache_event_hashes.pop(key) + for key in dict.fromkeys(hashes) + if key in self._external_cache_event_hashes + ] + if block_hashes: + self.kv_events.enqueue( + BlockRemoved(block_hashes=block_hashes, medium=medium) + ) + def finish_write_through(self, node_ids: list[NodeId], ack_id: int) -> None: """Clear the write-through-pending mark (when it matches ack_id) and record the host store event for each acked node.""" @@ -2844,9 +2893,9 @@ def attach_swa_window( covered = assign_end if covered >= window_end: break - assert covered == window_end, ( - f"attach_swa_window covered {covered} of [{window_start}, {window_end})" - ) + assert ( + covered == window_end + ), f"attach_swa_window covered {covered} of [{window_start}, {window_end})" return actions def _attach_swa_segment( diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index a2163b1c92b1..744904273637 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -709,6 +709,18 @@ def finish_external_linker_offload( """ ... + @abstractmethod + def record_external_cache_store(self, node_ids: Sequence[NodeId], medium) -> None: + """Publish confirmed local storage, retaining its hashes beyond tree eviction.""" + ... + + @abstractmethod + def remove_external_cache_storage( + self, hashes: Sequence[str], medium, *, publish: bool = True + ) -> None: + """Invalidate stored marks and publish actual local storage removals.""" + ... + # Order-sensitive digest of write_back duplicate-reclaim victim ids, # cross-checked across TP ranks; cores that never reclaim keep 0. write_back_duplicate_reclaim_digest: int = 0 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index afcaa90b3173..e045e7db0026 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -21,6 +21,7 @@ import msgspec import torch +from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored from sglang.srt.distributed.communication_tags import P2PTag from sglang.srt.environ import envs from sglang.srt.managers.cache_controller import CacheOperation @@ -30,6 +31,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, CacheRequestHandle, + CacheRequestOutcome, DecLockRefParams, DecLockRefResult, EvictParams, @@ -1108,10 +1110,15 @@ def claim_kv_row(self, req: Req) -> bool: and req.finished() ): self.cache_controller.release_pp_prefetch(req.rid) - return self.session.try_cache_finished_req(req) + claimed = self.session.try_cache_finished_req(req) + if claimed and self.linker is not None and req.finished(): + self.linker.release_request(req.cache_request_handle, cancel_load=False) + return claimed @rank_consensus(same_params=["req.rid", "checkpointed"]) def on_release(self, req: Req, *, checkpointed: bool) -> None: + if self.linker is not None and req.finished(): + self.linker.release_request(req.cache_request_handle, cancel_load=False) if checkpointed: return for comp in self._components_tuple: @@ -2703,12 +2710,17 @@ def pop_prefetch_loaded_span( self.prefetch_loaded_storage_start_by_reqid.pop(request, None), ) + def finish(self, handle: CacheRequestHandle, outcome: CacheRequestOutcome) -> None: + if outcome == CacheRequestOutcome.SUCCESS and self.linker is not None: + self.linker.release_request(handle, cancel_load=False) + super().finish(handle, outcome) + @rank_consensus(same_params=True) def release_aborted_request(self, request: CacheRequestHandle) -> None: rid = request.rid self._storage_prefetch_wait.pop(request, None) if self.linker is not None: - self.linker.release_request(rid) + self.linker.release_request(request) self.prefetch_loaded_tokens_by_reqid.pop(request, None) self.prefetch_loaded_storage_start_by_reqid.pop(request, None) self.storage_prefetch_retries.cancel(rid) @@ -3603,6 +3615,7 @@ def check_hicache_events(self) -> None: self.linker.commit_completed_offloads( [bool(success) for success in successes.tolist()] ) + self.linker.drain_storage_removals() return # Reap the previous round's PP-sync sends before issuing new ones. @@ -3877,8 +3890,18 @@ def root_node(self): return self.tree_core.root_node def take_events(self): - # Drain the KV event queue from the TreeCore. - return self.tree_core.take_events() + events = self.tree_core.take_events() + if self.linker is not None: + backend = self.linker.cache_linker + if backend.local_storage_ownership is not None: + # Linker and HiCache are exclusive, so only the Linker owns this tier. + for event in events: + if ( + isinstance(event, (BlockStored, BlockRemoved)) + and event.medium == backend.local_storage_medium + ): + event.ownership = backend.local_storage_ownership + return events def resolve_node_handle(self, node_handle): """Look up the node object from its NodeId. diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 879dc7ee6c4a..25b6cbd45b5a 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -16,13 +16,13 @@ use crate::components::{ComponentSet, ComponentType, FULL, MAMBA, SWA}; use crate::node::ChildKeyType; use crate::node::{KeyNamespaceRef, NodeAccessError, NodeId, TreeCoreRuntimeError}; use crate::unified_lru_list::{TlruFloatConfig, TlruPromptEstimate}; -use crate::unified_tree_core::KvCacheEvent; use crate::unified_tree_core::{ BufferBackupSnapshot, BufferBackupState, CacheAction, CacheInitParams, CacheTransferPhase, DecLockRefParams, EvictLayer, EvictionStepResult, InsertParams, InsertResult, InsertStepResult, MatchPrefixParams, MatchResult, PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, Req, UnifiedTreeCore, }; +use crate::unified_tree_core::{KvCacheEvent, StorageMedium}; /// Translate only Rust panics; ordinary Python errors pass through unchanged. /// Keep this boundary outside the whole binding call so its MutexGuard unwinds @@ -45,6 +45,16 @@ fn catch_native_panic(operation: impl FnOnce() -> PyResult) -> PyResult } } +fn parse_storage_medium(medium: &str) -> PyResult { + match medium { + "GPU" => Ok(StorageMedium::Gpu), + "CPU_PINNED" => Ok(StorageMedium::Cpu), + _ => Err(PyValueError::new_err(format!( + "unknown storage medium: {medium}" + ))), + } +} + /// Parse a torch-style device string (e.g. "cpu", "cuda", "cuda:1"); a bare /// "cuda" means index 0, so callers must resolve the index themselves. fn parse_device(device: &str) -> PyResult { @@ -2328,6 +2338,28 @@ impl TreeCoreBinding { .map_err(tree_core_assertion_error) } + fn record_external_cache_store( + &self, + py: Python<'_>, + node_ids: Vec, + medium: &str, + ) -> PyResult<()> { + let medium = parse_storage_medium(medium)?; + py.allow_threads(|| self.core().record_external_cache_store(&node_ids, medium)) + .map_err(node_access_error) + } + + fn remove_external_cache_storage( + &self, + py: Python<'_>, + hashes: Vec, + medium: Option<&str>, + ) -> PyResult<()> { + let medium = medium.map(parse_storage_medium).transpose()?; + py.allow_threads(|| self.core().remove_external_cache_storage(&hashes, medium)); + Ok(()) + } + /// Order-sensitive digest of reclaimed coexisting host values. fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> i64 { py.detach(|| self.core().write_back_coexist_reclaim_digest) @@ -3591,6 +3623,19 @@ macro_rules! tree_core_binding { }) } + fn record_external_cache_store( + &self, py: Python<'_>, node_ids: Vec, medium: &str, + ) -> PyResult<()> { + catch_native_panic(|| self.inner.record_external_cache_store(py, node_ids, medium)) + } + + #[pyo3(signature = (hashes, medium=None))] + fn remove_external_cache_storage( + &self, py: Python<'_>, hashes: Vec, medium: Option<&str>, + ) -> PyResult<()> { + catch_native_panic(|| self.inner.remove_external_cache_storage(py, hashes, medium)) + } + /// Order-sensitive digest of reclaimed coexisting host values. #[pyo3(name = "write_back_duplicate_reclaim_digest")] fn write_back_coexist_reclaim_digest(&self, py: Python<'_>) -> PyResult { diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 600809dd37f8..70788aaadf8f 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -2850,6 +2850,62 @@ fn failed_external_offload_preserves_independently_confirmed_state() { assert!(leaf.external_cache_stored); } +#[test] +fn external_linker_inventory_survives_device_eviction() { + let mut tc = core(); + tc.enable_kv_cache_events = true; + tc.set_enable_external_cache_linker(true).unwrap(); + tc.insert(&insert_params(&vec![1, 2], &[10, 11])); + let leaf = tc + .match_prefix(&match_params(&vec![1, 2])) + .best_match_node_id; + let idx = tc.arena.resolve(leaf).unwrap(); + let hashes = tc.arena.node(idx).hash_value.clone().unwrap(); + tc.take_events(); + tc.mark_external_linker_offload_pending(leaf).unwrap(); + tc.finish_external_linker_offload(&[leaf], leaf, true) + .unwrap(); + tc.record_external_cache_store(&[leaf], StorageMedium::Cpu) + .unwrap(); + assert!(matches!( + tc.take_events().as_slice(), + [KvCacheEvent::BlockStored { + medium: StorageMedium::Cpu, + .. + }] + )); + + tc.remove_external_cache_storage(&hashes[..1], Some(StorageMedium::Cpu)); + assert!(!tc.arena.node(idx).external_cache_stored); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockRemoved { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[0])], + medium: StorageMedium::Cpu, + extra_key: None, + }] + ); + let (dropped, _) = tc.drop_subtree_no_host(leaf).unwrap(); + assert!(dropped); + assert!(tc.take_events().iter().all(|event| !matches!( + event, + KvCacheEvent::BlockRemoved { + medium: StorageMedium::Cpu, + .. + } + ))); + tc.reset(); + tc.remove_external_cache_storage(&hashes[1..], Some(StorageMedium::Cpu)); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockRemoved { + block_hashes: vec![crate::node::hash_str_to_int64(&hashes[1])], + medium: StorageMedium::Cpu, + extra_key: None, + }] + ); +} + #[test] fn external_linker_path_validation_is_atomic() { let mut tc = core(); @@ -2872,6 +2928,50 @@ fn external_linker_path_validation_is_atomic() { ); } +#[test] +fn external_linker_inventory_retains_lora_namespace_after_eviction() { + let mut tc = events_core(1); + let key = vec![1, 2]; + tc.insert(&insert_params_in_namespace( + &key, + &[10, 11], + Some("lora-a"), + None, + )); + let leaf = tc + .match_prefix(&match_params_in_namespace(&key, Some("lora-a"), None)) + .best_match_node_id; + let idx = tc.arena.resolve(leaf).unwrap(); + let hashes = tc.arena.node(idx).hash_value.clone().unwrap(); + tc.take_events(); + tc.record_external_cache_store(&[leaf], StorageMedium::Cpu) + .unwrap(); + let stored = tc.take_events(); + tc.drop_subtree_no_host(leaf).unwrap(); + tc.take_events(); + tc.remove_external_cache_storage(&hashes, Some(StorageMedium::Cpu)); + let removed = tc.take_events(); + let KvCacheEvent::BlockStored { + block_hashes: stored_hashes, + extra_key: stored_key, + .. + } = &stored[0] + else { + panic!("expected stored inventory"); + }; + let KvCacheEvent::BlockRemoved { + block_hashes: removed_hashes, + extra_key: removed_key, + .. + } = &removed[0] + else { + panic!("expected removed inventory"); + }; + assert_eq!(stored_key.as_deref(), Some("lora-a")); + assert_eq!(removed_key, stored_key); + assert_eq!(removed_hashes, stored_hashes); +} + #[test] fn mark_write_through_pending_stamps_the_supplied_ack() { let mut tc = core(); diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index e6d077796216..7b2e1acba3ad 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -659,6 +659,8 @@ pub struct UnifiedTreeCore { pub(crate) kv_event_queue: Vec>, /// Namespaced event hashes, seeded only by cache_salt; events omit extra_key. pub(crate) namespaced_event_hashes: HashMap>, + // Local storage outlives device tree nodes and engine cache resets. + external_cache_event_hashes: HashMap>)>, /// Hit count at which a node earns a host write-through backup. pub(crate) write_through_threshold: i64, @@ -936,6 +938,7 @@ impl UnifiedTreeCore { dynamo_kv_event_format: params.dynamo_kv_event_format, kv_event_queue: Vec::new(), namespaced_event_hashes: HashMap::new(), + external_cache_event_hashes: HashMap::new(), write_through_threshold: params.write_through_threshold, component_uuid_counters: COMPONENT_UUID_INITIAL_VALUES, device: params.device, @@ -4456,6 +4459,86 @@ impl UnifiedTreeCore { Ok(()) } + pub fn record_external_cache_store( + &mut self, + node_ids: &[NodeId], + medium: StorageMedium, + ) -> Result<(), NodeAccessError> { + if !self.enable_kv_cache_events { + return Ok(()); + } + let nodes = node_ids + .iter() + .map(|&id| self.arena.resolve(id)) + .collect::, _>>()?; + for node_idx in nodes { + self.record_store_event_(node_idx, medium, None); + let node = self.arena.node(node_idx); + let hashes = node + .hash_value + .as_ref() + .expect("store event hashes its node"); + let event_hashes: Vec = if node.namespace != KeyNamespace::default() { + self.namespaced_event_hashes[&node.id] + .iter() + .map(crate::node::hash_digest_to_int64) + .collect() + } else { + hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect() + }; + self.external_cache_event_hashes.extend( + hashes.iter().cloned().zip( + event_hashes + .into_iter() + .map(|hash| (hash, node.namespace.extra_key().map(Arc::from))), + ), + ); + } + Ok(()) + } + + pub fn remove_external_cache_storage( + &mut self, + hashes: &[String], + medium: Option, + ) { + let removed: HashSet<_> = hashes.iter().collect(); + let nodes: Vec<_> = self + .collect_all_nodes_() + .into_iter() + .filter(|&idx| { + let node = self.arena.node(idx); + node.external_cache_stored + && node + .hash_value + .as_ref() + .is_some_and(|values| values.iter().any(|key| removed.contains(key))) + }) + .collect(); + for idx in nodes { + self.arena.node_mut(idx).external_cache_stored = false; + } + let Some(medium) = medium else { + return; + }; + let mut by_namespace: HashMap>, Vec> = HashMap::new(); + for key in hashes { + if let Some((hash, extra_key)) = self.external_cache_event_hashes.remove(key) { + by_namespace.entry(extra_key).or_default().push(hash); + } + } + for (extra_key, block_hashes) in by_namespace { + self.enqueue_kv_event_(KvCacheEvent::BlockRemoved { + block_hashes, + medium, + extra_key, + }); + } + } + /// Clear the write-through-pending mark (when it matches ack_id) and record the /// host store event for each acked node. pub fn finish_write_through( diff --git a/test/registered/unit/disaggregation/test_kv_events.py b/test/registered/unit/disaggregation/test_kv_events.py index a46362a47b9e..c0447963970e 100644 --- a/test/registered/unit/disaggregation/test_kv_events.py +++ b/test/registered/unit/disaggregation/test_kv_events.py @@ -52,8 +52,12 @@ def _publisher(self, attn_dp_rank=0, **config): def test_advertised_source_delivers_events(self): for host, bind, topic in (("*", None, "kv"), ("127.0.0.1", True, "")): with self.subTest(host=host): - with zmq.Context.instance().socket(zmq.PUB) as probe: - port = probe.bind_to_random_port("tcp://127.0.0.1") + # Destroy the probe's own context before rebinding its port; + # Socket.close on the shared context only schedules close on + # its I/O thread and can race the immediately following bind. + with zmq.Context() as probe_context: + with probe_context.socket(zmq.PUB) as probe: + port = probe.bind_to_random_port("tcp://127.0.0.1") publisher = self._publisher( attn_dp_rank=4, endpoint=f"tcp://{host}:{port - 4}", @@ -333,6 +337,17 @@ def test_lora_name_is_a_named_field(self): decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) self.assertEqual(decoded["lora_name"], "adapter-a") + def test_kvcr_ownership_is_an_optional_named_field(self): + for event in ( + self._event(ownership="kvcr"), + BlockRemoved( + block_hashes=[123], medium=StorageMedium.CPU, ownership="kvcr" + ), + AllBlocksCleared(ownership="kvcr"), + ): + decoded = msgspec.msgpack.decode(msgspec.msgpack.encode(event)) + self.assertEqual(decoded["ownership"], "kvcr") + def test_one_decoder_reads_a_mixed_batch(self): batch = KVEventBatch( ts=1.0, diff --git a/test/registered/unit/mem_cache/test_unified_cache_linker.py b/test/registered/unit/mem_cache/test_unified_cache_linker.py index 69b42bca9dda..f19b34f20a45 100644 --- a/test/registered/unit/mem_cache/test_unified_cache_linker.py +++ b/test/registered/unit/mem_cache/test_unified_cache_linker.py @@ -10,19 +10,26 @@ import pytest import torch +from sglang.srt.disaggregation.kv_events import ( + AllBlocksCleared, + BlockRemoved, + BlockStored, + StorageMedium, +) from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + CacheRequestOutcome, + EvictParams, InitLoadBackParams, + InsertParams, MatchPrefixParams, MatchResult, ) from sglang.srt.mem_cache.common import release_kv_cache -from sglang.srt.mem_cache.hicache_storage import ( - PoolHitPolicy, - PoolName, - PoolTransfer, -) +from sglang.srt.mem_cache.events import KVCacheEventRecorder +from sglang.srt.mem_cache.hicache_storage import PoolHitPolicy, PoolName, PoolTransfer from sglang.srt.mem_cache.radix_cache import RadixKey from sglang.srt.mem_cache.unified_cache.cache_action import ( ReplaceWriteThroughOnNodeSplit, @@ -39,12 +46,11 @@ UnifiedCacheLinker, UnifiedCacheLinkerWrapper, ) +from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci from sglang.test.kits.unified_radix_cache import fixtures as shared_cache_suite -from sglang.test.kits.unified_radix_cache.config import ( - CacheConfig, -) +from sglang.test.kits.unified_radix_cache.config import CacheConfig from sglang.test.kits.unified_radix_cache.fixtures import ( build_fixture, node_children, @@ -114,6 +120,18 @@ def set_full_to_swa_mapping(self, full, swa): self.mapping.append((full.clone(), swa.clone())) +class _ContextLinker(_FakeLinker): + def __init__(self): + super().__init__() + self.contexts = {} + + def set_request_context(self, handle, kv_hints): + self.contexts[handle] = kv_hints + + def release_request_context(self, handle): + self.contexts.pop(handle) + + class _FakeExternalTreeCore: def __init__(self, nodes=None, offload_transfers=None): self.enable_external_cache_linker = False @@ -137,9 +155,9 @@ def mark_external_linker_offload_pending(self, node_id): def finish_external_linker_offload(self, node_ids, ack_id, success): nodes = [self.nodes[node_id] for node_id in node_ids] - assert all(node.write_through_pending_id == ack_id for node in nodes), ( - "invalid external offload state" - ) + assert all( + node.write_through_pending_id == ack_id for node in nodes + ), "invalid external offload state" for node in nodes: node.write_through_pending_id = None node.external_cache_stored |= success @@ -154,6 +172,10 @@ def _cache_for_wrapper(**kwargs): "write_through_threshold": 256, "pp_size": 1, "pp_group": None, + "attn_cp_group": None, + "attn_tp_group": None, + "tp_world_size": 1, + "_all_reduce_attn_groups": lambda value, op: None, } defaults.update(kwargs) return SimpleNamespace(**defaults) @@ -977,22 +999,100 @@ class TestUnifiedCacheLinkerTreeCoreRustBackend( tree_core_backend = "rust" -def test_restorable_prefix_intersects_sparse_rank_results(): - remote_mask = torch.tensor([0, 0, 1, 0, 0], dtype=torch.int) +@pytest.mark.parametrize("confirmed,expected", [(None, 2), ({1}, 1), (set(), 0)]) +def test_restorable_prefix_intersects_sparse_rank_results(confirmed, expected): + """Rejecting the longest shared boundary must retry a shorter shared one.""" + remote_mask = torch.tensor([0, 1, 1, 1, 0], dtype=torch.int) + probes = [] + + def confirm(pages): + probes.append(pages) + return True def intersect_remote_mask(mask, op): assert op == torch.distributed.ReduceOp.MIN - mask.copy_(torch.minimum(mask, remote_mask)) + if mask.numel() > 1: + mask.copy_(torch.minimum(mask, remote_mask)) + else: + mask[0] &= int(probes[-1] in confirmed) cache = _cache_for_wrapper(_all_reduce_attn_groups=intersect_remote_mask) wrapper = UnifiedCacheLinkerWrapper(cache, _FakeLinker()) - hit_pages = wrapper._sync_restorable_prefix([2, 4], num_pages=4, device_hit_pages=0) + hit_pages = wrapper._sync_restorable_prefix( + [1, 2, 4], + num_pages=4, + device_hit_pages=0, + confirm=None if confirmed is None else confirm, + ) + assert hit_pages == expected + assert probes == ( + [] if confirmed is None else [2, 1] + ([0] if not expected else []) + ) + + +@pytest.mark.parametrize( + "path", ["match_rejected", "allocation_failure", "device_hit", "load_refused"] +) +def test_lookup_claims_released_when_device_load_is_unnecessary( + full_linker_component, monkeypatch, path +): + """Rejected or superseded loads must not retain lookup claims or cache hits.""" + backend = _FakeLinker() + backend.requires_lookup_confirmation = True + backend.restorable = [1] + claims = {"held"} + + def discard(rid, transfers, pages): + if pages: + return False + claims.clear() + return True - assert hit_pages == 2 + backend.confirm_lookup = discard + if path == "allocation_failure": + full_linker_component.build_external_linker_transfer = lambda *args: None + wrapper = UnifiedCacheLinkerWrapper( + _cache_for_wrapper( + _components_tuple=(full_linker_component,), + page_size=2, + tree_core=SimpleNamespace( + enable_external_cache_linker=False, is_eagle=False + ), + ), + backend, + ) + key = RadixKey(array("q", [1, 2])) + if path in ("allocation_failure", "device_hit"): + wrapper.hit_markers["discard"] = ExternalCacheHitMarker(["page"], 0) + if path == "load_refused": + wrapper.hit_markers["discard"] = ExternalCacheHitMarker(["page"], 0) + backend.load = lambda *args: False + req = SimpleNamespace( + rid="discard", + last_node=0, + cache_request_handle=CacheRequestHandle("discard", 0), + kv_hints=None, + ) + if path in ("device_hit", "match_rejected"): + monkeypatch.setattr( + "sglang.srt.mem_cache.utils.get_native_hash", lambda *args: ["page"] + ) + result = MatchResult( + device_prefix_len=2 if path == "device_hit" else 0, + last_device_node=0, + last_host_node=0, + best_match_node=0, + ) + assert wrapper.match(key, req, result) == result + else: + assert wrapper.load_back(req) == (0, 0) + assert not claims + assert not wrapper.has_hit(req.rid) -def test_async_offload_pins_node_until_completion(): +@pytest.mark.parametrize("success", [False, True]) +def test_async_offload_pins_node_until_completion(success): linker = _FakeLinker() lock_params = object() locks = [] @@ -1022,15 +1122,181 @@ def inc_lock_ref(node): assert node.write_through_pending_id == node_id assert not unlocks - linker.completed_offloads.append(False) + linker.completed_offloads.append(success) completed = wrapper.take_completed_offloads(finish_count=1) wrapper.commit_completed_offloads(completed) - assert not node.external_cache_stored + assert node.external_cache_stored == success assert node.write_through_pending_id is None assert unlocks == [(node_id, lock_params)] +@pytest.mark.parametrize("success", [False, True]) +@pytest.mark.parametrize("events_enabled", [False, True]) +@pytest.mark.parametrize("remote_removal", [False, True]) +@pytest.mark.parametrize("namespace_seed", [None, b"\x01" * 32]) +def test_local_inventory_tracks_storage_removal_after_tree_eviction( + success, events_enabled, remote_removal, namespace_seed, monkeypatch +): + hashes = ["1" * 64, "2" * 64] + root = SimpleNamespace(parent=None, hash_value=[]) + node = SimpleNamespace( + id=7, + external_cache_stored=False, + write_through_pending_id=None, + key=RadixKey(array("q", [1, 2, 3, 4])), + hash_value=hashes, + parent=root, + ) + core = _FakeExternalTreeCore( + {node.id: node}, [PoolTransfer(name=PoolName.KV, keys=hashes)] + ) + core._node_arena = core.nodes + core._external_cache_event_hashes = {} + core.node_by_id = core.nodes.__getitem__ + core.kv_events = KVCacheEventRecorder(enabled=events_enabled, page_size=2) + # Removal IDs must equal the published IDs, including namespace projection. + monkeypatch.setattr(core.kv_events, "_namespace_seed", lambda node: namespace_seed) + core.record_external_cache_store = lambda node_ids, medium: ( + UnifiedTreeCore.record_external_cache_store(core, node_ids, medium) + ) + core.remove_external_cache_storage = lambda keys, medium, **kwargs: ( + UnifiedTreeCore.remove_external_cache_storage(core, keys, medium, **kwargs) + ) + backend = _FakeLinker() + backend.local_storage_medium = StorageMedium.CPU + removals = [] + backend.pop_storage_removals = lambda: list(removals) + cache = _cache_for_wrapper( + tree_core=core, + inc_lock_ref=lambda node: SimpleNamespace(to_dec_params=object), + dec_lock_ref=lambda *args: None, + ) + wrapper = UnifiedCacheLinkerWrapper(cache, backend) + wrapper.offload_nodes([node.id]) + assert core.kv_events.take() == [] + backend.completed_offloads.append(success) + completed = wrapper.take_completed_offloads(1) + assert core.kv_events.take() == [] + wrapper.commit_completed_offloads(completed) + stored = core.kv_events.take() + assert bool(stored) == (success and events_enabled) + if stored: + assert isinstance(stored[0], BlockStored) + assert stored[0].medium == StorageMedium.CPU + + if remote_removal: + cache.tp_world_size = 2 + cache.tp_group = object() + cache._all_reduce_attn_groups = lambda value, op: value.fill_(1) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2) + monkeypatch.setattr( + torch.distributed, + "all_gather_object", + lambda gathered, local, **kwargs: gathered.__setitem__( + slice(None), [local, hashes[:1]] + ), + ) + wrapper.drain_storage_removals() + assert not node.external_cache_stored + assert core.kv_events.take() == [] + assert set(core._external_cache_event_hashes) == ( + set(hashes) if stored else set() + ) + + removals.append(hashes[0]) + wrapper.drain_storage_removals() + assert not node.external_cache_stored + removed = core.kv_events.take() + if stored: + assert isinstance(removed[0], BlockRemoved) + assert removed[0].block_hashes == stored[0].block_hashes[:1] + else: + assert removed == [] + + core.nodes.clear() + removals[:] = [hashes[1]] + wrapper.drain_storage_removals() + removed = core.kv_events.take() + if stored: + assert removed[0].block_hashes == stored[0].block_hashes[1:] + assert removed[0].medium == StorageMedium.CPU + else: + assert removed == [] + + +@pytest.mark.parametrize("tree_core_backend", ["python", "rust"]) +@pytest.mark.parametrize("linker_mode", ["none", "default", "owned"]) +def test_take_events_tags_only_linker_local_storage( + tree_core_backend, linker_mode, monkeypatch +): + monkeypatch.setattr(shared_cache_suite, "get_device", lambda: "cpu") + cache, allocator, _ = build_fixture( + CacheConfig(page_size=1, kv_size=8, max_context_len=8), + enable_kv_cache_events=True, + tree_core_backend=tree_core_backend, + ) + cache.take_events() + inserted = cache.insert( + InsertParams(key=RadixKey(array("q", [1, 2])), value=allocator.alloc(2)) + ) + core = cache.tree_core + core.record_external_cache_store([inserted.last_device_node], StorageMedium.CPU) + hashes = core.get_hash_values(inserted.last_device_node) + core.remove_external_cache_storage(hashes[:1], StorageMedium.CPU) + cache.evict(EvictParams(num_tokens=2)) + core.kv_events.record_all_cleared() + if linker_mode != "none": + backend = _FakeLinker() + assert backend.local_storage_ownership is None + backend.local_storage_medium = StorageMedium.CPU + if linker_mode == "owned": + backend.local_storage_ownership = "kvcr" + cache.init_cache_linker(backend) + owner = "kvcr" if linker_mode == "owned" else None + events = cache.take_events() + assert [ + (type(event), getattr(event, "medium", None), event.ownership) + for event in events + ] == [ + (BlockStored, StorageMedium.GPU, None), + (BlockStored, StorageMedium.CPU, owner), + (BlockRemoved, StorageMedium.CPU, owner), + (BlockRemoved, StorageMedium.GPU, None), + (AllBlocksCleared, None, None), + ] + cache.reset() + assert cache.take_events() == [AllBlocksCleared()] + core.remove_external_cache_storage(hashes[1:], StorageMedium.CPU) + assert cache.take_events() == [ + BlockRemoved( + block_hashes=events[1].block_hashes[1:], + medium=StorageMedium.CPU, + ownership=owner, + ) + ] + + +@pytest.mark.parametrize("event_type", [BlockStored, BlockRemoved]) +def test_kv_event_coalescing_preserves_ownership(event_type): + recorder = KVCacheEventRecorder(enabled=True, page_size=1) + for index, owner in enumerate([None, "kvcr", "kvcr", None]): + fields = dict(block_hashes=[index], medium=StorageMedium.CPU, ownership=owner) + if event_type is BlockStored: + fields.update( + parent_block_hash=index - 1 if index else None, + token_ids=[index], + block_size=1, + lora_id=None, + ) + recorder.enqueue(event_type(**fields)) + assert [(event.block_hashes, event.ownership) for event in recorder.take()] == [ + ([0], None), + ([1, 2], "kvcr"), + ([3], None), + ] + + def test_offload_skips_node_already_stored_by_tree_core(): linker = _FakeLinker() node = SimpleNamespace( @@ -1144,7 +1410,7 @@ def test_release_request_cancels_queued_load(): linker.queued_loads["rid"] = [object()] wrapper.private_load_indices["rid"] = torch.arange(4) - wrapper.release_request("rid") + wrapper.release_request(CacheRequestHandle("rid", 0)) assert wrapper.hit_markers == {} assert wrapper.inflight_load_rids == [] @@ -1152,6 +1418,75 @@ def test_release_request_cancels_queued_load(): assert wrapper.private_load_indices == {} +@pytest.mark.parametrize("cancel_load", [False, True]) +@pytest.mark.parametrize("finish", ["drain", "reset", "close"]) +def test_request_release_notifies_backend_without_unpinning_load(cancel_load, finish): + backend = _ContextLinker() + unlocked = [] + cache = _cache_for_wrapper(dec_lock_ref=lambda *args: unlocked.append(args)) + wrapper = UnifiedCacheLinkerWrapper(cache, backend) + handle = CacheRequestHandle("rid", 1) + wrapper._request_handles[handle.rid] = handle + backend.set_request_context(handle, {"scope": "first"}) + wrapper.inflight_load_rids = [handle.rid] + wrapper.private_load_indices[handle.rid] = torch.arange(4) + wrapper._load_request_handles[handle.rid] = handle + + wrapper.release_request(handle, cancel_load=cancel_load) + + assert backend.contexts == {} + assert not unlocked + if finish == "drain": + wrapper.start_layer_wise_loading() + assert wrapper.finish_loads() == [] + else: + getattr(wrapper, finish)() + assert unlocked == [] # Request-private slots never acquire tree locks. + assert wrapper.private_load_indices == {} + assert wrapper._load_request_handles == {} + assert wrapper._request_handles == {} + + +def test_successful_request_releases_linker_context(): + released = [] + cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.linker = SimpleNamespace( + release_request=lambda handle, **kwargs: released.append((handle, kwargs)) + ) + handle = CacheRequestHandle("rid", 2) + + cache.finish(handle, CacheRequestOutcome.SUCCESS) + + assert released == [(handle, {"cancel_load": False})] + + +@pytest.mark.parametrize("finished", [False, True]) +@pytest.mark.parametrize("session_claims", [False, True]) +@pytest.mark.parametrize("checkpointed", [False, True]) +def test_normal_release_cleans_context_only_for_finished_request( + finished, session_claims, checkpointed +): + backend = _ContextLinker() + wrapper = UnifiedCacheLinkerWrapper(_cache_for_wrapper(), backend) + handle = CacheRequestHandle("rid", 0) + wrapper._request_handles[handle.rid] = handle + backend.set_request_context(handle, None) + cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.cache_controller = None + cache.linker = wrapper + cache._components_tuple = () + cache.session = SimpleNamespace(try_cache_finished_req=lambda req: session_claims) + req = SimpleNamespace( + rid="rid", cache_request_handle=handle, finished=lambda: finished + ) + + assert cache.claim_kv_row(req) == session_claims + if not session_claims: + assert backend.contexts == {handle: None} + cache.on_release(req, checkpointed=checkpointed) + assert backend.contexts == ({} if finished else {handle: None}) + + def test_failed_offload_rolls_back_split_fragments(): linker = _FakeLinker() lock_params = object() @@ -1271,6 +1606,7 @@ def test_check_hicache_events_commits_common_rank_results(): num_completed_offloads=lambda: 3, take_completed_offloads=lambda count: [True] * count, commit_completed_offloads=committed.append, + drain_storage_removals=lambda: None, ) reduce_calls = 0 @@ -1370,9 +1706,12 @@ def build_transfer(phase, node, keys): return SimpleNamespace( component_type=ComponentType.FULL, build_external_linker_transfer=MagicMock(side_effect=build_transfer), - update_external_linker_load=lambda phase, req, full_transfer, transfer, prefix_len, **kwargs: ( - transfer - ), + update_external_linker_load=lambda phase, + req, + full_transfer, + transfer, + prefix_len, + **kwargs: (transfer), ) @@ -1405,7 +1744,13 @@ def test_linker_filters_unstored_swa_from_lookup(full_linker_component): best_match_node=0, ) matched = wrapper.match( - RadixKey(array("q", [1, 2, 3, 4])), SimpleNamespace(rid="match"), result + RadixKey(array("q", [1, 2, 3, 4])), + SimpleNamespace( + rid="match", + cache_request_handle=CacheRequestHandle("match", 0), + kv_hints=None, + ), + result, ) assert matched.host_hit_length == 4 assert [c.args[0] for c in full.build_external_linker_transfer.call_args_list] == [ From 38e9d4fb0173ac94187acf9b5ac5ebfcdfa24ae4 Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Tue, 20 Oct 2026 09:07:58 -0700 Subject: [PATCH 4/5] fix(cache): use current PyO3 detach API for inventory bindings --- rust/sglang-radix-tree/src/python_bindings.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 25b6cbd45b5a..ff7081239ebd 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -2345,7 +2345,7 @@ impl TreeCoreBinding { medium: &str, ) -> PyResult<()> { let medium = parse_storage_medium(medium)?; - py.allow_threads(|| self.core().record_external_cache_store(&node_ids, medium)) + py.detach(|| self.core().record_external_cache_store(&node_ids, medium)) .map_err(node_access_error) } @@ -2356,7 +2356,7 @@ impl TreeCoreBinding { medium: Option<&str>, ) -> PyResult<()> { let medium = medium.map(parse_storage_medium).transpose()?; - py.allow_threads(|| self.core().remove_external_cache_storage(&hashes, medium)); + py.detach(|| self.core().remove_external_cache_storage(&hashes, medium)); Ok(()) } From a2f9da8c9a7d9f4b15e4ca7a4c3d1735f0bdb11d Mon Sep 17 00:00:00 2001 From: Ziqi Fan Date: Tue, 20 Oct 2026 09:10:54 -0700 Subject: [PATCH 5/5] style(cache): apply upstream formatter to linker contracts --- .../unified_cache/unified_tree_core.py | 24 +++++++++---------- .../mem_cache/test_unified_cache_linker.py | 15 +++++------- 2 files changed, 18 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 9f2d888a0b7a..60a5f17cf7b1 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -736,9 +736,9 @@ def inc_lock_ref( def _assert_receipt_anchor(node: UnifiedTreeNode, params: DecLockRefParams) -> None: """A receipt releases only the node its acquire returned; a mispaired or unanchored one would silently release (or steal) another holder's segment.""" - assert ( - params.node_id == node.id - ), f"lock receipt anchored on node {params.node_id} released on node {node.id}" + assert params.node_id == node.id, ( + f"lock receipt anchored on node {params.node_id} released on node {node.id}" + ) def _release_components( self, @@ -1695,12 +1695,12 @@ def _finish_internal_component_eviction( ) -> EvictDeviceNextNodeResult: """Resume an internal tombstone after the controller's backup attempt.""" component = self.components_by_type[component_type] - assert ( - component.is_evict_device_ongoing - ), f"{component_type.name} device eviction not started" - assert ( - self._pending_internal_evictions.get(component_type) == node_id - ), f"no matching pending internal {component_type.name} eviction" + assert component.is_evict_device_ongoing, ( + f"{component_type.name} device eviction not started" + ) + assert self._pending_internal_evictions.get(component_type) == node_id, ( + f"no matching pending internal {component_type.name} eviction" + ) del self._pending_internal_evictions[component_type] # Consuming the pending request advances the walk even if I/O changed # this victim's eligibility. No frees occur before backup completion. @@ -2893,9 +2893,9 @@ def attach_swa_window( covered = assign_end if covered >= window_end: break - assert ( - covered == window_end - ), f"attach_swa_window covered {covered} of [{window_start}, {window_end})" + assert covered == window_end, ( + f"attach_swa_window covered {covered} of [{window_start}, {window_end})" + ) return actions def _attach_swa_segment( diff --git a/test/registered/unit/mem_cache/test_unified_cache_linker.py b/test/registered/unit/mem_cache/test_unified_cache_linker.py index f19b34f20a45..1d4f29d2ef6e 100644 --- a/test/registered/unit/mem_cache/test_unified_cache_linker.py +++ b/test/registered/unit/mem_cache/test_unified_cache_linker.py @@ -155,9 +155,9 @@ def mark_external_linker_offload_pending(self, node_id): def finish_external_linker_offload(self, node_ids, ack_id, success): nodes = [self.nodes[node_id] for node_id in node_ids] - assert all( - node.write_through_pending_id == ack_id for node in nodes - ), "invalid external offload state" + assert all(node.write_through_pending_id == ack_id for node in nodes), ( + "invalid external offload state" + ) for node in nodes: node.write_through_pending_id = None node.external_cache_stored |= success @@ -1706,12 +1706,9 @@ def build_transfer(phase, node, keys): return SimpleNamespace( component_type=ComponentType.FULL, build_external_linker_transfer=MagicMock(side_effect=build_transfer), - update_external_linker_load=lambda phase, - req, - full_transfer, - transfer, - prefix_len, - **kwargs: (transfer), + update_external_linker_load=lambda phase, req, full_transfer, transfer, prefix_len, **kwargs: ( + transfer + ), )