From 1a53f231d62d373cbc35ecbeaad727ff75834fe8 Mon Sep 17 00:00:00 2001 From: Zhangheng Date: Tue, 21 Jul 2026 18:59:51 +0800 Subject: [PATCH] [HiCache]: Optimize hybrid/DSA L3 prefetch result sync and usable-prefix clamping (#31443) Co-authored-by: AlejandroParedesLT Co-authored-by: Kevin Flansburg Co-authored-by: Chao Shi --- .../sglang/srt/mem_cache/hicache_storage.py | 11 +- python/sglang/srt/mem_cache/hiradix_cache.py | 42 ++++++- .../hybrid_cache/hybrid_cache_controller.py | 6 +- .../srt/mem_cache/unified_radix_cache.py | 110 +++++++++++++++--- 4 files changed, 146 insertions(+), 23 deletions(-) diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index fe9e0dd2c..28de214ab 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -132,9 +132,16 @@ class PoolTransferResult: self.kv_hit_pages = max(self.kv_hit_pages, kv_hit_pages) def update_extra_pool_hit_pages(self, results: dict[str, List[bool]]) -> None: - """Record actual load/write success counts per extra pool.""" + """Record actual load/write success counts per extra pool. + + Every extra pool contributes a prefix that must be contiguous from the + start, so count the leading run of successes + """ self.extra_pool_hit_pages.update( - {name: sum(rs) for name, rs in results.items()} + { + name: (rs.index(False) if False in rs else len(rs)) + for name, rs in results.items() + } ) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 7455ca09e..bedf2f740 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1556,13 +1556,10 @@ class HiRadixCache(RadixCache): ) logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens") - min_completed_tokens = completed_tokens - # Synchronize workers before mutating host cache tree state. - completed_tokens_tensor = torch.tensor(min_completed_tokens, dtype=torch.int) - self._all_reduce_attn_groups( - completed_tokens_tensor, torch.distributed.ReduceOp.MIN + min_completed_tokens = self._sync_and_clamp_prefetch_result( + operation, completed_tokens ) - min_completed_tokens = completed_tokens_tensor.item() + fetched_key = prefetch_key[:min_completed_tokens] written_indices = operation.host_indices[:min_completed_tokens] matched_length = self._insert_helper_host( @@ -1591,6 +1588,39 @@ class HiRadixCache(RadixCache): return True + def _sync_and_clamp_prefetch_result( + self, + operation: PrefetchOperation, + completed_tokens: int, + ) -> int: + """Sync prefetch results across ATTN groups and decide the usable prefix. + + HiRadixCache only wires DSA-style stacks (Full attention + a KV-derived + ALL_PAGES sidecar such as the DSA / MiniMax indexer); For the DSA case we *clamp* + to the minimum fetched prefix shared by the Full KV pool and every + sidecar rather than discarding everything. With no sidecar (FULL-only) + this is just the synced Full KV completion. + """ + # Sync completed tokens and per-pool hit pages across ATTN groups, taking + # the minimum so every rank agrees on the same usable prefix length. + pool_transfers = getattr(operation, "pool_transfers", None) or [] + hit_pages = ( + operation.pool_storage_result.extra_pool_hit_pages if pool_transfers else {} + ) + pool_hit_pages = [hit_pages.get(t.name, 0) for t in pool_transfers] + packed = torch.tensor([completed_tokens, *pool_hit_pages], dtype=torch.int) + self._all_reduce_attn_groups(packed, torch.distributed.ReduceOp.MIN) + min_completed_tokens = int(packed[0].item()) + pool_hit_pages = list(map(int, packed[1:].tolist())) + + # Clamp to the shared minimum prefix of the Full KV completion and each + # KV-derived ALL_PAGES sidecar (e.g. the DSA indexer). FULL-only has no + # sidecar, so the usable prefix is just the Full KV completion. + usable_pages = min_completed_tokens // self.page_size + if pool_transfers: + usable_pages = min(usable_pages, *pool_hit_pages) + return usable_pages * self.page_size + def terminate_prefetch(self, req_id: str): if req_id not in self.ongoing_prefetch: return 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 d6c20b821..af35e964a 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 @@ -654,7 +654,11 @@ class HybridCacheController(BaseHiCacheController): # (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): + if ( + operation.pool_transfers + and not operation.is_terminated() + and kv_completed_pages == len(operation.hash_value) + ): self._sync_trailing_keys( operation.pool_transfers, operation.hash_value, kv_completed_pages ) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 71042a21e..24adcb48b 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -30,6 +30,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( ) from sglang.srt.mem_cache.events import KVCacheEventMixin from sglang.srt.mem_cache.hicache_storage import ( + PoolHitPolicy, PoolName, PoolTransfer, SidecarPoolSpec, @@ -2005,20 +2006,20 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): completed_tokens, hash_value = self.cache_controller.terminate_prefetch( operation ) - min_completed_tokens = completed_tokens - hit_pages = operation.pool_storage_result.extra_pool_hit_pages - if self.tp_world_size > 1: - # Reduce full completed tokens together with the sidecar pools that - # this prefetch actually transferred, in one all_reduce. - sidecar_pools = [t.name for xfers in comp_xfers.values() for t in xfers] - packed = torch.tensor( - [completed_tokens] + [hit_pages.get(p, 0) for p in sidecar_pools], - dtype=torch.int, - ) - self._all_reduce_attn_groups(packed, torch.distributed.ReduceOp.MIN) - min_completed_tokens = int(packed[0].item()) - for i, p in enumerate(sidecar_pools, start=1): - hit_pages[p] = int(packed[i].item()) + + min_completed_tokens = self._sync_and_check_hybrid_prefetch_result( + req_id, + operation, + completed_tokens, + hash_value, + host_indices, + last_host_node, + anchor_lock_params, + prefetch_key, + ) + if min_completed_tokens is None: + # Hybrid all-or-nothing check failed; result already discarded. + return True fetched_key = prefetch_key[:min_completed_tokens] insert_result = self._insert_helper_host( @@ -2063,6 +2064,87 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) return True + def _sync_and_check_hybrid_prefetch_result( + self, + req_id: str, + operation: PrefetchOperation, + completed_tokens: int, + hash_value: list[str], + host_indices: torch.Tensor, + last_host_node: UnifiedTreeNode, + anchor_lock_params: DecLockRefParams, + prefetch_key: RadixKey, + ) -> Optional[int]: + """Sync prefetch results across ATTN groups and decide the usable prefix. + + Two strategies depending on the hybrid layout: + + * DSA-style (Full attention + KV-derived ALL_PAGES sidecar such as the + DSA / MiniMax indexer): *clamp* to the minimum fetched prefix shared by + the Full KV pool and every sidecar. A partial prefix is still usable + because the sidecar is page-aligned with KV and required for every page. + * Everything else (SWA / Mamba components, mixed DeepSeekV4 stacks): + *all-or-nothing*. Their pools only cover a window / tail and cannot be + truncated page by page, so any shortfall discards the whole prefetch. + + Returns the synced usable token count (possibly clamped, possibly 0), or + ``None`` when an all-or-nothing prefetch was discarded (the caller should + then treat the prefetch as finished). + """ + # Sync completed tokens and per-pool hit pages across ATTN groups, taking + # the minimum so every rank agrees on the same usable prefix length. + pool_transfers = operation.pool_transfers or [] + hit_pages = ( + operation.pool_storage_result.extra_pool_hit_pages if pool_transfers else {} + ) + pool_hit_pages = [hit_pages.get(t.name, 0) for t in pool_transfers] + packed = torch.tensor([completed_tokens, *pool_hit_pages], dtype=torch.int) + self._all_reduce_attn_groups(packed, torch.distributed.ReduceOp.MIN) + min_completed_tokens = int(packed[0].item()) + pool_hit_pages = list(map(int, packed[1:].tolist())) + for transfer, count in zip(pool_transfers, pool_hit_pages): + hit_pages[transfer.name] = count + + # DSA-style clamp: every sidecar is KV-derived and required for the whole + # prefix (ALL_PAGES), so the usable length is simply the shared minimum of + # the Full KV completion and each sidecar hit. + clampable = bool(pool_transfers) and all( + t.hit_policy == PoolHitPolicy.ALL_PAGES + and t.indices_from_pool == PoolName.KV + for t in pool_transfers + ) + if clampable: + usable_pages = min(min_completed_tokens // self.page_size, *pool_hit_pages) + return usable_pages * self.page_size + + # Hybrid cache state is all-or-nothing: every extra pool (SWA / Mamba / ...) + # must cover the same fetched prefix. If any pool falls short the whole + # prefetch result is unusable, so discard it and release everything. + expected_tokens = len(hash_value) * self.page_size + all_succeeded = min_completed_tokens == expected_tokens and all( + transfer.keys is not None and count == len(transfer.keys) + for transfer, count in zip(pool_transfers, pool_hit_pages) + ) + if pool_transfers and not all_succeeded: + # The controller's prefetch IO thread already releases the untransferred + # tail (host_indices[completed_tokens:]) + self.cache_controller.append_host_mem_release( + host_indices=host_indices[:completed_tokens], + extra_pools=pool_transfers, + ) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + del self.ongoing_prefetch[req_id] + self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + self.prefetch_loaded_tokens_by_reqid[req_id] = 0 + logger.warning( + "HiCache hybrid prefetch discarded req=%s completed=%d requested=%d", + req_id, + completed_tokens, + expected_tokens, + ) + return None + return min_completed_tokens + def terminate_prefetch(self, req_id: str) -> None: if req_id not in self.ongoing_prefetch: return