diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 56057ade6..ec50932e0 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -625,7 +625,15 @@ class HiCacheController: self.dp_rank = 0 # Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool. - is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool) + # DeepSeekV4TokenToKVPool has compressed MLA-style rank-replicated cache + # data. storage only needs rank 0 to write it back. + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool + + is_mla_model = isinstance(self.mem_pool_device, MLATokenToKVPool) + is_compressed_mla_model = isinstance( + self.mem_pool_device, DeepSeekV4TokenToKVPool + ) + is_rank_replicated = is_mla_model or is_compressed_mla_model # Least Common Multiple among heterogeneous tp size tp_lcm_size = storage_backend_extra_config.pop("tp_lcm_size", None) should_split_heads = False @@ -635,7 +643,7 @@ class HiCacheController: tp_lcm_size % self.tp_size == 0 ), "tp_lcm_size must be divisible by tp_size." should_split_heads = ( - not is_mla_backend + not is_rank_replicated and self.mem_pool_host.layout == "page_head" and tp_lcm_size > self.tp_size ) @@ -649,7 +657,8 @@ class HiCacheController: pp_size=self.pp_size, attn_cp_rank=attn_cp_rank, attn_cp_size=attn_cp_size, - is_mla_model=is_mla_backend, + # TODO(hzh): Rename is_mla_model to is_rank_replicated. + is_mla_model=is_rank_replicated, enable_storage_metrics=self.enable_storage_metrics, is_page_first_layout=self.mem_pool_host.layout == "page_first", model_name=model_name, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 1514e601d..62a0252b2 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -74,6 +74,7 @@ class InsertResult: """Result of an insert operation""" prefix_len: int + total_len: int = 0 mamba_exist: bool = False inserted_host_node: Any = None diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index cf14fb6c3..58e620f2f 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -136,6 +136,7 @@ class PrefetchOperation(StorageOperation): prefix_keys=prefix_keys, pool_transfers=pool_transfers, ) + self.pool_transfers_done = not bool(pool_transfers) def increment(self, num_tokens: int): with self._lock: @@ -581,9 +582,6 @@ class HybridCacheController(BaseHiCacheController): kv_hit_pages = hit_result.kv_hit_pages operation.pool_storage_result.update_kv_hit_pages(kv_hit_pages) - if kv_hit_pages > 0 and operation.pool_transfers: - self._sync_trailing_keys(operation.pool_transfers, hash_value, kv_hit_pages) - return ( hash_value[:kv_hit_pages], kv_hit_pages * self.page_size, @@ -618,14 +616,21 @@ class HybridCacheController(BaseHiCacheController): return host_indices, device_indices, resolved_pool_transfers def _page_transfer(self, operation): - # Transfer extra pools - if operation.pool_transfers and not operation.is_terminated(): + # KV pools first — determines actual completed page count + super()._page_transfer(operation) + + # Extra pools only after KV fully completes. If KV terminated early + # (IO failure, timeout, TP mismatch), skip extra IO entirely to avoid + # data misalignment. + kv_completed_pages = operation.completed_tokens // self.page_size + if operation.pool_transfers and kv_completed_pages == len(operation.hash_value): + self._sync_trailing_keys( + operation.pool_transfers, operation.hash_value, kv_completed_pages + ) self._resolve_sidecar_derived_pool_transfers(operation) results = self.storage_backend.batch_get_v2(operation.pool_transfers) operation.pool_storage_result.update_extra_pool_hit_pages(results) - - # Transfer kv pools - super()._page_transfer(operation) + operation.pool_transfers_done = True def _page_backup(self, operation): # Backup extra pools @@ -642,14 +647,27 @@ class HybridCacheController(BaseHiCacheController): if transfer.indices_from_pool is None: continue if transfer.indices_from_pool != PoolName.KV: - # TODO(hzh): Support storage sidecar derived pools from other sources - raise AssertionError( - "Storage sidecar derived pool currently only supports KV-shared " - f"indices, got {transfer.name} from {transfer.indices_from_pool}." + source = next( + ( + t + for t in operation.pool_transfers + if t.indices_from_pool is None + and t.name == transfer.indices_from_pool + ), + None, ) - transfer.host_indices = operation.host_indices - if transfer.keys is None: - transfer.keys = operation.hash_value + if source is None: + raise AssertionError( + "Storage sidecar derived pool source missing: " + f"{transfer.name} from {transfer.indices_from_pool}." + ) + transfer.host_indices = source.host_indices + if transfer.keys is None: + transfer.keys = source.keys + else: + transfer.host_indices = operation.host_indices + if transfer.keys is None: + transfer.keys = operation.hash_value def _sync_trailing_keys( self, diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py index e12c9d350..63abbd73c 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_pool_assembler.py @@ -4,7 +4,11 @@ import logging from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Callable, Optional -from sglang.srt.mem_cache.hicache_storage import PoolName, SidecarPoolSpec +from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, + PoolName, + SidecarPoolSpec, +) from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, ) @@ -737,7 +741,15 @@ class _DeepSeekV4Strategy(StackStrategy): enable_storage_metrics=enable_storage_metrics, ) sidecars = [ - SidecarPoolSpec(pool_name=name, indices_from_pool=src) + SidecarPoolSpec( + pool_name=name, + indices_from_pool=src, + hit_policy=( + PoolHitPolicy.TRAILING_PAGES + if src == PoolName.SWA + else PoolHitPolicy.ALL_PAGES + ), + ) for name, src in ( (PoolName.DEEPSEEK_V4_C4, PoolName.KV), (PoolName.DEEPSEEK_V4_C4_INDEXER, PoolName.KV), diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 913918e16..69e65f4e4 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -2078,6 +2078,7 @@ class DeepSeekV4PagedHostPool(HostKVCache): return host_rows = self._to_page_indices(host_indices) device_rows = self._to_page_indices(device_indices) + if io_backend == "kernel" and self.layout == "layer_first": transfer_kv_per_layer_mla( src=self.data_refs[layer_id], diff --git a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py index 6124a60fe..71281bed9 100644 --- a/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py +++ b/python/sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py @@ -601,6 +601,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): def register_mem_pool_host(self, mem_pool_host: HostKVCache): super().register_mem_pool_host(mem_pool_host) + if getattr(self.mem_pool_host, "kv_buffer", None) is None: + # Hybrid logical anchors only own allocation indices. Their physical + # tensors are registered through register_mem_host_pool_v2(). + return assert self.mem_pool_host.layout in [ "page_first", "page_first_direct", @@ -631,37 +635,41 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): # the corresponding host pool implementation at runtime. self.registered_pools[host_pool_name] = host_pool - # Hybrid pools expose the tensors that Mooncake needs for zero-copy I/O. - # The storage backend only depends on this accessor, not concrete fields. - buf_list = host_pool.get_hybrid_pool_buffer() - for buf in buf_list: + # Non-anchor pools are either sidecar-specific pools with their own + # accessor, or ordinary KV-like host pools used as SWA side pools. + get_buffers = getattr( + host_pool, + "get_hybrid_pool_buffer", + lambda: [getattr(host_pool, "kv_buffer", None)], + ) + for buf in get_buffers(): + if buf is None: + continue super().register_buffer(buf) def _tag_keys(self, keys: List[str]) -> List[str]: if self.extra_backend_tag is None: return keys - return [f"{ self.extra_backend_tag}_{key}" for key in keys] + return [f"{self.extra_backend_tag}_{key}" for key in keys] def _get_hybrid_page_component_keys( self, page_keys: List[str], transfer: PoolTransfer ) -> Tuple[List[str], int]: - # A logical "page" may map to multiple physical objects in storage. - # - INDEXER: one key per page - # - MAMBA : one temporal key + N conv keys per page - # key_multiplier records how many component keys are generated per page. - name = transfer.name + host_pool = getattr(self, "registered_pools", {}).get(transfer.name) + if host_pool is None: + raise ValueError(f"Unregistered Mooncake hybrid pool: {transfer.name}") + + # Suffix order must match get_page_buffer_meta() for one page, because + # Mooncake zips object keys with registered buffer pointers. + pool_name = transfer.name suffixes = [] - if name == PoolName.INDEXER: - suffixes = [f"_{self.mla_suffix}_{PoolName.INDEXER}"] - elif name == PoolName.MAMBA: - pools = getattr(self, "registered_pools", {}) - mamba_pool = pools.get(PoolName.MAMBA) - conv_num = len(getattr(mamba_pool, "conv_buffer", None) or []) - base_suffix = f"_{self.mha_suffix}" - suffixes = [f"{base_suffix}_temporal"] + [ - f"{base_suffix}_conv_{i}" for i in range(conv_num) + if pool_name == PoolName.MAMBA: + # Mamba stores one temporal object plus one object per conv state. + conv_num = len(getattr(host_pool, "conv_buffer", None) or []) + suffixes = [f"_{self.mha_suffix}_temporal"] + [ + f"_{self.mha_suffix}_conv_{i}" for i in range(conv_num) ] - elif name == PoolName.DRAFT: + elif pool_name == PoolName.DRAFT: # Draft pool's MLA/MHA layout is independent from the target # (e.g. EAGLE-MHA draft on top of an MLA target), so pick the # suffix scheme from the draft pool's own class. The `_draft` @@ -675,6 +683,33 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): f"_{self.mha_suffix}_{PoolName.DRAFT}_k", f"_{self.mha_suffix}_{PoolName.DRAFT}_v", ] + elif pool_name in ( + PoolName.INDEXER, + PoolName.DEEPSEEK_V4_C4, + PoolName.DEEPSEEK_V4_C4_INDEXER, + PoolName.DEEPSEEK_V4_C128, + PoolName.DEEPSEEK_V4_C4_STATE, + PoolName.DEEPSEEK_V4_C4_INDEXER_STATE, + PoolName.DEEPSEEK_V4_C128_STATE, + ): + # DSA indexer and DeepSeek V4 side pools are page-packed + # single-object pools. + suffixes = [f"_{self.mla_suffix}_{pool_name}"] + elif pool_name == PoolName.SWA: + if not self.is_mla_backend and hasattr(host_pool, "v_buffer"): + # Ordinary MHA SWA mirrors a K/V pool. + suffixes = [ + f"_{self.mha_suffix}_{pool_name}_k", + f"_{self.mha_suffix}_{pool_name}_v", + ] + elif self.is_mla_backend: + suffixes = [f"_{self.mla_suffix}_{pool_name}"] + + if not suffixes: + raise ValueError( + f"Unsupported Mooncake hybrid pool name: {pool_name}, " + f"host_pool={type(host_pool)}" + ) key_multiplier = len(suffixes) component_keys = [ f"{page_key}{suffix}" for page_key in page_keys for suffix in suffixes @@ -687,7 +722,12 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): pool_transfers: Optional[List[PoolTransfer]] = None, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> PoolTransferResult: - kv_pages = self.batch_exists(keys, extra_info) + if self.mem_pool_host.kv_buffer is None: + # Logical anchor: no physical KV object exists in Mooncake, so the + # usable prefix is determined entirely by required sidecar objects. + kv_pages = len(keys) + else: + kv_pages = self.batch_exists(keys, extra_info) hit_count: dict = {PoolName.KV: kv_pages} if kv_pages else {} final_pages = kv_pages @@ -863,6 +903,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): host_indices: torch.Tensor, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: + if self.mem_pool_host.kv_buffer is None: + # DeepSeek V4's KV anchor is logical only; v2 side pools carry data. + return [True] * len(keys) + # Apply extra_backend_tag prefix if available keys = self._tag_keys(keys) @@ -888,6 +932,10 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): host_indices: torch.Tensor, extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: + if self.mem_pool_host.kv_buffer is None: + # DeepSeek V4's KV anchor is logical only; v2 side pools carry data. + return [True] * len(keys) + # Apply extra_backend_tag prefix if available keys = self._tag_keys(keys) diff --git a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py index 4f9c1519a..d6b7cb6c3 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache_components/swa_component.py @@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( MatchResult, ) from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, PoolName, PoolTransfer, PoolTransferResult, @@ -415,7 +416,9 @@ class SWAComponent(TreeComponent): root = self.cache.root_node sliding_window_size = self.sliding_window_size swa_lock_size = 0 - swa_uuid_for_lock = None + swa_uuid = None + uuid_key = "host_uuid" if lock_host else "uuid" + lru = self.cache.host_lru_lists[ct] if lock_host else self.cache.lru_lists[ct] # Tombstoned nodes (cd.value is None) have no SWA chunk to protect # skip them and keep walking up. This path is hit when HiCache @@ -423,23 +426,36 @@ class SWAComponent(TreeComponent): cur = node while cur != root and swa_lock_size < sliding_window_size: comp = cur.component_data[ct] - if comp.value is None: + value = comp.host_value if lock_host else comp.value + if value is None: result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id) cur = cur.parent continue - if comp.lock_ref == 0: - key_len = len(cur.key) - self.cache.component_evictable_size_[ct] -= key_len - self.cache.component_protected_size_[ct] += key_len - comp.lock_ref += 1 - swa_lock_size += len(cur.key) + + ref = comp.host_lock_ref if lock_host else comp.lock_ref + if ref == 0: + if lock_host: + if lru.in_list(cur): + lru.remove_node(cur) + else: + key_len = len(cur.key) + self.cache.component_evictable_size_[ct] -= key_len + self.cache.component_protected_size_[ct] += key_len + if lock_host: + comp.host_lock_ref = ref + 1 + else: + comp.lock_ref = ref + 1 + swa_lock_size += len(value) if swa_lock_size >= sliding_window_size: - if comp.metadata.get("uuid") is None: - comp.metadata["uuid"] = next_component_uuid() - swa_uuid_for_lock = comp.metadata["uuid"] + if comp.metadata.get(uuid_key) is None: + comp.metadata[uuid_key] = next_component_uuid() + swa_uuid = comp.metadata[uuid_key] cur = cur.parent - result.swa_uuid_for_lock = swa_uuid_for_lock + if lock_host: + result.swa_uuid_for_host_lock = swa_uuid + else: + result.swa_uuid_for_lock = swa_uuid return result def release_component_lock( @@ -450,9 +466,14 @@ class SWAComponent(TreeComponent): ) -> None: ct = self.component_type root = self.cache.root_node - swa_uuid_for_lock = params.swa_uuid_for_lock if params else None + swa_uuid_for_lock = ( + (params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock) + if params + else None + ) skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else () dec_swa = True + uuid_key = "host_uuid" if lock_host else "uuid" # A node in skip_lock_node_ids was a tombstone when this lock was acquired. cur = node @@ -461,15 +482,25 @@ class SWAComponent(TreeComponent): if cur.id in skip_lock_node_ids: cur = cur.parent continue - if comp.lock_ref == 0: + ref = comp.host_lock_ref if lock_host else comp.lock_ref + if ref == 0: cur = cur.parent continue - if comp.lock_ref == 1: - key_len = len(cur.key) - self.cache.component_evictable_size_[ct] += key_len - self.cache.component_protected_size_[ct] -= key_len - comp.lock_ref -= 1 - if swa_uuid_for_lock and comp.metadata.get("uuid") == swa_uuid_for_lock: + if ref == 1: + if lock_host: + if comp.value is None and comp.host_value is not None: + host_lru = self.cache.host_lru_lists[ct] + if not host_lru.in_list(cur): + host_lru.insert_mru(cur) + else: + key_len = len(comp.value) + self.cache.component_evictable_size_[ct] += key_len + self.cache.component_protected_size_[ct] -= key_len + if lock_host: + comp.host_lock_ref = ref - 1 + else: + comp.lock_ref = ref - 1 + if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock: dec_swa = False cur = cur.parent @@ -546,6 +577,46 @@ class SWAComponent(TreeComponent): ) ] + if phase == CacheTransferPhase.BACKUP_STORAGE: + cd = node.component_data[ct] + if cd.host_value is None or not node.hash_value: + return None + num_pages = len(cd.host_value) // self.cache.page_size + if num_pages == 0: + return None + return [ + PoolTransfer( + name=PoolName.SWA, + host_indices=cd.host_value[-num_pages * self.cache.page_size :], + keys=node.hash_value[-num_pages:], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + + if phase == CacheTransferPhase.PREFETCH: + num_pages = min( + prefetch_tokens // self.cache.page_size, + (self.sliding_window_size + self.cache.page_size - 1) + // self.cache.page_size, + ) + if num_pages == 0: + return None + num_tokens = num_pages * self.cache.page_size + host_indices = self._swa_kv_pool_host.alloc(num_tokens) + if host_indices is None: + self.cache.evict_host(num_tokens, ComponentType.SWA) + host_indices = self._swa_kv_pool_host.alloc(num_tokens) + if host_indices is None: + return [] + return [ + PoolTransfer( + name=PoolName.SWA, + host_indices=host_indices, + keys=["__placeholder__"] * num_pages, + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + ] + return None def commit_hicache_transfer( @@ -586,6 +657,98 @@ class SWAComponent(TreeComponent): assert offset == len(xfer.host_indices) return + if phase == CacheTransferPhase.PREFETCH: + self._commit_prefetch( + node, + transfers, + insert_result=insert_result, + pool_storage_result=pool_storage_result, + ) + return + + def _release_swa_host(self, host_indices: torch.Tensor) -> None: + if host_indices is not None and host_indices.numel() > 0: + self.cache.cache_controller.append_host_mem_release( + extra_pools=[PoolTransfer(name=PoolName.SWA, host_indices=host_indices)] + ) + + def _attach_swa_host_value( + self, node: "UnifiedTreeNode", host_indices: torch.Tensor + ) -> None: + """Write host_indices into node's SWA host_value and refresh tree state.""" + ct = self.component_type + cd = node.component_data[ct] + cd.host_value = host_indices.clone() + host_lru = self.cache.host_lru_lists[ct] + if cd.value is None and not host_lru.in_list(node): + host_lru.insert_mru(node) + self.cache._update_evictable_leaf_sets(node) + if node.parent: + self.cache._update_evictable_leaf_sets(node.parent) + + def _commit_prefetch( + self, + anchor, + transfers: list[PoolTransfer], + *, + insert_result: Optional[InsertResult] = None, + pool_storage_result: Optional[PoolTransferResult] = None, + ) -> None: + """Distribute the prefetched SWA buffer onto the leaf→anchor path. + + The buffer holds the trailing ``loaded_pages`` of the completed KV + prefix, mapped to token range ``[loaded_start, total_len)``. We walk + upward from ``inserted_host_node`` to ``anchor`` and, for each node + whose token range overlaps the buffer: + - SWA tombstone (host_value is None) → fill from buffer (split if + the node only partially overlaps at the buffer's left edge) + - already has SWA host_value → release the corresponding slice + Any leftover buffer beyond the walked range is also released. + """ + if not transfers: + return + ct = self.component_type + host_indices = transfers[0].host_indices + loaded_pages = ( + pool_storage_result.extra_pool_hit_pages.get(PoolName.SWA, 0) + if pool_storage_result + else 0 + ) + target = insert_result.inserted_host_node if insert_result else None + if not loaded_pages or target is None: + self._release_swa_host(host_indices) + return + + # Buffer covers token range [loaded_start, total_len). + loaded_start = insert_result.total_len - loaded_pages * self.cache.page_size + + # Walk leaf → anchor; ``pos`` is the right edge of ``cur`` in tokens. + pos, cur = insert_result.total_len, target + while cur is not anchor and pos > loaded_start: + node_start = pos - len(cur.key) + # Intersection of cur's range and the buffer. + fill_start = max(node_start, loaded_start) + fill_len = pos - fill_start + buf_off = fill_start - loaded_start + slice_ = host_indices[buf_off : buf_off + fill_len] + + cd = cur.component_data[ct] + if cd.host_value is None and fill_len > 0: + # Tombstone: split off the in-buffer tail if needed, then fill. + if fill_start > node_start: + self.cache._split_node(cur.key, cur, fill_start - node_start) + self._attach_swa_host_value(cur, slice_) + else: + # Already has SWA (or empty overlap): drop this slice. + self._release_swa_host(slice_) + + pos = node_start + cur = cur.parent + + # Buffer prefix that fell outside the anchor→leaf path. + if pos > loaded_start: + self._release_swa_host(host_indices[: pos - loaded_start]) + def drive_host_eviction( self, num_tokens: int, tracker: dict[ComponentType, int] ) -> None: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 36406875f..5017ab3af 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1121,9 +1121,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if len(key): child_key = key.child_key(self.page_size) - result = InsertResult( - prefix_len=matched_length, - ) + result = InsertResult(prefix_len=matched_length, total_len=total_len) if len(key) == 0: if ( node is not self.root_node @@ -1662,6 +1660,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): transfers.append( PoolTransfer( name=spec.pool_name, + keys=indices_source.keys, hit_policy=spec.hit_policy, indices_from_pool=spec.indices_from_pool, ) @@ -1850,6 +1849,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): ) else: return True + if ( + completed + and getattr(operation, "pool_transfers", None) + and not getattr(operation, "pool_transfers_done", True) + ): + can_terminate = False operation_terminated = operation.is_terminated() states = torch.tensor( diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index f590ada12..5edfbdb6e 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1,5 +1,9 @@ """Unit tests for UnifiedRadixCache""" +import json +import shutil +import tempfile +import time import unittest from array import array from dataclasses import dataclass, replace @@ -1840,6 +1844,146 @@ class UnifiedRadixCacheSuite: # HiCache Unit Tests (real cache_controller D<->H backup/load) # ================================================================ + # ---------- L3 storage (file backend) helpers ---------- + + def _path_chain(self, tree, node): + """Return root->node node chain (excluding root).""" + chain = [] + cur = node + while cur is not tree.root_node: + chain.append(cur) + cur = cur.parent + chain.reverse() + return chain + + def _write_path_to_l3(self, tree, node): + """Offload every node on root->node path from host to L3 storage.""" + for n in self._path_chain(tree, node): + tree.write_backup_storage(n) + + def _flush_l3_backups(self, tree, timeout: float = 10.0): + """Wait for backup threads to finish, then drain acks (release locks).""" + deadline = time.time() + timeout + while tree.ongoing_backup and time.time() < deadline: + tree.drain_storage_control_queues() + if tree.ongoing_backup: + time.sleep(0.01) + tree.drain_storage_control_queues() + self.assertFalse(tree.ongoing_backup, "L3 backups did not complete in time") + + def _run_prefetch_to_completion(self, tree, req_id, timeout: float = 10.0): + deadline = time.time() + timeout + while time.time() < deadline: + if tree.check_prefetch_progress(req_id): + return + time.sleep(0.01) + self.fail(f"prefetch {req_id} did not complete in time") + + def _all_page_hashes(self, tree, node): + hashes = [] + for n in self._path_chain(tree, node): + hashes.extend(list(n.hash_value)) + return hashes + + def test_hicache_l3_write_storage(self): + """D->H->L3 offload: every KV page lands in the file storage backend.""" + if self._skip_unsupported_hicache_test(): + return + if self.cfg.has_mamba: + self.skipTest("mamba L3 offload is out of scope for this unit fixture") + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + self._init_hicache( + tree, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + ) + + seq = self._make_seq(1, 4) + self._insert(tree, allocator, req_to_token_pool, seq) + m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + leaf = m.last_device_node + + # D->H first, then H->L3. + self._backup_node(tree, leaf) + self.assertTrue(leaf.hash_value) + self._write_path_to_l3(tree, leaf) + self._flush_l3_backups(tree) + + # Every KV page hash on the path must now exist in storage. + backend = tree.cache_controller.storage_backend + page_hashes = self._all_page_hashes(tree, leaf) + self.assertEqual(len(page_hashes), len(seq) // self.cfg.page_size) + self.assertEqual(backend.batch_exists(page_hashes), len(page_hashes)) + tree.sanity_check() + + def test_hicache_l3_prefetch(self): + """L3 round trip: write with one tree, prefetch into a fresh tree. + + Uses two independent trees that share the same file storage dir so the + prefetch path genuinely reloads from L3 (no host/device residue). + """ + if self._skip_unsupported_hicache_test(): + return + if self.cfg.has_mamba: + self.skipTest("mamba L3 prefetch is out of scope for this unit fixture") + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + seq = self._make_seq(1, 4) + + # --- Producer tree: fill KV, backup D->H, offload H->L3. --- + prod, prod_alloc, prod_rtp = build_fixture(self.cfg) + self._init_hicache( + prod, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + ) + self._insert(prod, prod_alloc, prod_rtp, seq) + mp = prod.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + prod_leaf = mp.last_device_node + self._fill_full_kv(prod_alloc, mp.device_indices, marker=7) + expected_k, expected_v = self._snapshot_full_kv(prod_alloc, mp.device_indices) + self._backup_node(prod, prod_leaf) + self._write_path_to_l3(prod, prod_leaf) + self._flush_l3_backups(prod) + + # --- Consumer tree: prefetch the same tokens straight from L3. --- + cons, cons_alloc, cons_rtp = build_fixture(self.cfg) + self._init_hicache( + cons, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + ) + req_id = "l3-prefetch-req" + cons.prefetch_from_storage(req_id, cons.root_node, array("q", seq), None, None) + self._run_prefetch_to_completion(cons, req_id) + cons.drain_storage_control_queues() + + # The full prefix must now be a host hit (loaded from L3). + mc = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) + self.assertEqual(mc.host_hit_length, len(seq)) + host_node = mc.last_host_node + self.assertIsNot(host_node, cons.root_node) + self.assertTrue(host_node.evicted) + + # Load the reloaded host prefix back to device and verify KV bytes. + self._load_back_node(cons, host_node) + loaded_indices = cons.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).device_indices + self.assertEqual(len(loaded_indices), len(seq)) + loaded_k, loaded_v = self._snapshot_full_kv(cons_alloc, loaded_indices) + self.assertTrue(torch.equal(loaded_k, expected_k)) + self.assertTrue(torch.equal(loaded_v, expected_v)) + cons.sanity_check() + def _skip_unsupported_hicache_test(self): if self.cfg.has_swa and self.cfg.has_mamba: self.skipTest("HiCache unit fixture does not support SWA + Mamba stacks") @@ -1869,7 +2013,16 @@ class UnifiedRadixCacheSuite: self._simulate_backup(tree, node) stack.extend(node.children.values()) - def _init_hicache(self, tree, *, write_policy: str = "write_through"): + def _init_hicache( + self, + tree, + *, + write_policy: str = "write_through", + storage_backend: Optional[str] = None, + storage_dir: Optional[str] = None, + prefetch_threshold: Optional[int] = None, + prefetch_policy: str = "wait_complete", + ): import sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler as assembler orig_kv_host_pool = assembler.MHATokenToKVPoolHost @@ -1899,11 +2052,42 @@ class UnifiedRadixCacheSuite: patcher.start() self.addCleanup(patcher.stop) + storage_extra_config = None + if storage_backend == "file": + import sglang.srt.managers.cache_controller as cache_controller + + # The file-backend storage config records TP rank/size. These unit + # fixtures run without initializing distributed parallel state, so + # provide the local single-rank values that the fixture represents. + tp_rank_patcher = mock.patch.object( + cache_controller, "get_tensor_model_parallel_rank", return_value=0 + ) + tp_size_patcher = mock.patch.object( + cache_controller, "get_tensor_model_parallel_world_size", return_value=1 + ) + tp_rank_patcher.start() + tp_size_patcher.start() + self.addCleanup(tp_rank_patcher.stop) + self.addCleanup(tp_size_patcher.stop) + + assert storage_dir is not None, "file backend needs a storage_dir" + # HiCacheFile reads the directory from this env var. + cm = envs.SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR.override(storage_dir) + cm.__enter__() + self.addCleanup(cm.__exit__, None, None, None) + extra = {} + if prefetch_threshold is not None: + extra["prefetch_threshold"] = prefetch_threshold + storage_extra_config = json.dumps(extra) if extra else None + server_args = ServerArgs( model_path="dummy", page_size=self.cfg.page_size, hicache_io_backend="direct", hicache_write_policy=write_policy, + hicache_storage_backend=storage_backend, + hicache_storage_backend_extra_config=storage_extra_config, + hicache_storage_prefetch_policy=prefetch_policy, ) # See build_fixture for why _mamba_cache_chunk_size is preset. server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, self.cfg.page_size) @@ -1911,6 +2095,17 @@ class UnifiedRadixCacheSuite: tree.init_hicache(server_args, tree.cache_init_params) tree.write_through_threshold = 1 << 30 tree.load_back_threshold = 0 + if storage_backend is not None: + # Unit fixtures size host/device pools equally, which makes the + # production prefetch capacity limit (host - device) zero. Keep the + # L3 tests focused on storage round trips by allowing one fixture + # worth of prefetch tokens. + tree.cache_controller.prefetch_capacity_limit = max( + tree.cache_controller.prefetch_capacity_limit, + tree.cache_controller.mem_pool_host.size, + ) + # Background prefetch/backup threads are daemon; stop them per-test. + self.addCleanup(tree.cache_controller._stop_storage_threads) def _build_hicache_fixture(self): fixture = build_fixture(self.cfg)