From 7cd55c6818727124b49a2f9bb8ae09fc27cb66f1 Mon Sep 17 00:00:00 2001 From: Bob Chen Date: Fri, 17 Jul 2026 06:00:53 +0800 Subject: [PATCH] [HiCache] Optimize L2 mem allocation when cache miss in L3 (#19320) Co-authored-by: Zhiqiang Xie --- benchmark/hf3fs/bench_zerocopy.py | 2 +- .../sglang/srt/managers/cache_controller.py | 29 ++-- .../srt/mem_cache/hi_mamba_radix_cache.py | 117 +++++++++----- python/sglang/srt/mem_cache/hiradix_cache.py | 145 +++++++++++------- .../hybrid_cache/hybrid_cache_controller.py | 6 +- .../srt/mem_cache/unified_radix_cache.py | 126 +++++++++------ .../test_unified_radix_cache_unittest.py | 16 +- 7 files changed, 274 insertions(+), 167 deletions(-) diff --git a/benchmark/hf3fs/bench_zerocopy.py b/benchmark/hf3fs/bench_zerocopy.py index 800299d47..0f9c5e150 100644 --- a/benchmark/hf3fs/bench_zerocopy.py +++ b/benchmark/hf3fs/bench_zerocopy.py @@ -114,7 +114,6 @@ print(f"{tok-tik:.6f} s") operations = [ PrefetchOperation( f"{i}", - torch.tensor(list(range(i, i + op_size))), list(range(i, i + op_size)), f"{i}", ) @@ -122,6 +121,7 @@ operations = [ ] for operation in operations: + operation.host_indices = torch.tensor(operation.token_ids) operation.hash_value = [ f"{j}" for j in range( diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 2cc807e91..0f3234b3b 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -170,7 +170,7 @@ class StorageOperation: def __init__( self, - host_indices: torch.Tensor, + host_indices: Optional[torch.Tensor], token_ids: List[int], last_hash: Optional[str] = None, hash_value: Optional[List[str]] = None, @@ -194,7 +194,6 @@ class PrefetchOperation(StorageOperation): def __init__( self, request_id: str, - host_indices: torch.Tensor, token_ids: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, @@ -203,9 +202,10 @@ class PrefetchOperation(StorageOperation): self._lock = threading.Lock() self._terminated_flag = False + self.storage_hit_count = 0 self.start_time = time.monotonic() - super().__init__(host_indices, token_ids, last_hash, prefix_keys=prefix_keys) + super().__init__(None, token_ids, last_hash, prefix_keys=prefix_keys) def increment(self, num_tokens: int): with self._lock: @@ -374,6 +374,7 @@ class HiCacheController: self.prefetch_queue = Queue() self.backup_queue = Queue() + self.prefetch_hit_queue: Queue[StorageOperation] = Queue() self.prefetch_revoke_queue: Queue[str] = Queue() self.ack_backup_queue: Queue[StorageOperation] = Queue() self.host_mem_release_queue: Queue[torch.Tensor] = Queue() @@ -652,6 +653,7 @@ class HiCacheController: self.prefetch_queue.queue.clear() self.backup_queue.queue.clear() self.prefetch_revoke_queue.queue.clear() + self.prefetch_hit_queue.queue.clear() self.ack_backup_queue.queue.clear() self.host_mem_release_queue.queue.clear() self.prefetch_tokens_occupied = 0 @@ -898,7 +900,6 @@ class HiCacheController: def prefetch( self, request_id: str, - host_indices: torch.Tensor, new_input_tokens: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, @@ -907,7 +908,7 @@ class HiCacheController: Prefetch KV caches from storage backend to host memory. """ operation = PrefetchOperation( - request_id, host_indices, new_input_tokens, last_hash, prefix_keys + request_id, new_input_tokens, last_hash, prefix_keys ) self.prefetch_queue.put(operation) return operation @@ -1057,7 +1058,10 @@ class HiCacheController: operation = self.prefetch_queue.get(block=True, timeout=1) if operation is None: continue - hash_value, storage_hit_count = self._storage_hit_query(operation) + if operation.is_terminated(): + hash_value, storage_hit_count = [], 0 + else: + hash_value, storage_hit_count = self._storage_hit_query(operation) storage_hit_count_tensor = torch.tensor( storage_hit_count, dtype=torch.int ) @@ -1069,23 +1073,16 @@ class HiCacheController: if storage_hit_count < self.prefetch_threshold: # not to prefetch if not enough benefits self.prefetch_revoke_queue.put(operation.request_id) - self.append_host_mem_release(operation.host_indices) logger.debug( f"Revoking prefetch for request {operation.request_id} due to insufficient hits ({storage_hit_count})." ) else: + # Record hit count, so the scheduler thread will know the exact memory to allocate operation.hash_value = hash_value[ : (storage_hit_count // self.page_size) ] - # free the pre-allocated memory for pages that are not hit - self.append_host_mem_release( - operation.host_indices[storage_hit_count:] - ) - operation.host_indices = operation.host_indices[:storage_hit_count] - logger.debug( - f"Prefetching {len(operation.hash_value)} pages for request {operation.request_id}." - ) - self.prefetch_buffer.put(operation) + operation.storage_hit_count = storage_hit_count + self.prefetch_hit_queue.put(operation) except Empty: continue diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py index 565fe4475..249c970cc 100644 --- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py @@ -1476,14 +1476,28 @@ class HiMambaRadixCache(MambaRadixCache): def _drain_storage_control_queues_local(self): self._drain_storage_control_queues_impl( n_revoke=None, + n_storage_hit=0, n_backup=None, n_release=None, log_metrics=False, ) + def _revoke_pending_prefetch(self, req_id: str): + info = self.ongoing_prefetch.pop(req_id, None) + if info is None: + return + last_host_node, token_ids, _host_indices, operation = info + self.prefetch_abort(operation.pool_transfers) + self._release_host_node(last_host_node) + cc = self.cache_controller + cc.prefetch_tokens_occupied = max( + 0, cc.prefetch_tokens_occupied - len(token_ids) + ) + def _drain_storage_control_queues_impl( self, n_revoke: Optional[int], + n_storage_hit: Optional[int], n_backup: Optional[int], n_release: Optional[int], log_metrics: bool, @@ -1502,14 +1516,57 @@ class HiMambaRadixCache(MambaRadixCache): def _drain_revoke(): for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): - info = self.ongoing_prefetch.pop(req_id, None) - if info is not None: - last_host_node, token_ids, _, operation = info - self.prefetch_abort(operation.pool_transfers) - self._release_host_node(last_host_node) - cc.prefetch_tokens_occupied -= len(token_ids) - if cc.prefetch_tokens_occupied < 0: - cc.prefetch_tokens_occupied = 0 + self._revoke_pending_prefetch(req_id) + + def _drain_and_alloc_storage_hit(): + # The L3 hit count is now known, so reserve exactly that much host + # KV memory. NOTE: alloc/evict here is rank-local but deterministic + # across TP ranks without extra synchronization: host pool + # mutations only happen on the scheduler thread at lockstep points + # (releases are page-granular and drained by the TP-min count), so + # every rank reaches the same success / fallback / revoke decision. + for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit): + req_id = operation.request_id + info = self.ongoing_prefetch.get(req_id) + if info is None: + # request already aborted/cleaned up, skip + continue + if operation.is_terminated(): + # request was aborted while the storage query was in flight + self._revoke_pending_prefetch(req_id) + continue + + alloc_len = operation.storage_hit_count + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self.evict_host(alloc_len) + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + # Memory-pressure fallback: a shorter page-aligned prefix. + available_size = cc.mem_pool_host.available_size() + alloc_len = min( + operation.storage_hit_count, + available_size - (available_size % self.page_size), + ) + if alloc_len >= self.prefetch_threshold: + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self._revoke_pending_prefetch(req_id) + continue + + operation.storage_hit_count = alloc_len + operation.hash_value = operation.hash_value[ + : alloc_len // self.page_size + ] + operation.host_indices = host_indices + last_host_node, token_ids, _, op = info + self.ongoing_prefetch[req_id] = ( + last_host_node, + token_ids, + host_indices, + op, + ) + cc.prefetch_buffer.put(operation) def _drain_backup(): for operation in _drain_queue(cc.ack_backup_queue, n_backup): @@ -1532,6 +1589,7 @@ class HiMambaRadixCache(MambaRadixCache): cc.mem_pool_host.free(host_indices) _drain_revoke() + _drain_and_alloc_storage_hit() _drain_backup() _drain_release() @@ -1642,6 +1700,7 @@ class HiMambaRadixCache(MambaRadixCache): qsizes = torch.tensor( [ cc.prefetch_revoke_queue.qsize(), + cc.prefetch_hit_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), ], @@ -1652,9 +1711,10 @@ class HiMambaRadixCache(MambaRadixCache): qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group ) - n_revoke, n_backup, n_release = map(int, qsizes.tolist()) + n_revoke, n_storage_hit, n_backup, n_release = map(int, qsizes.tolist()) self._drain_storage_control_queues_impl( n_revoke=n_revoke, + n_storage_hit=n_storage_hit, n_backup=n_backup, n_release=n_release, log_metrics=True, @@ -1710,8 +1770,6 @@ class HiMambaRadixCache(MambaRadixCache): return _, _, _, operation = self.ongoing_prefetch[req_id] - if operation.host_indices is None: - return operation.mark_terminate() def pop_prefetch_loaded_tokens(self, req_id: str) -> int: @@ -1756,30 +1814,9 @@ class HiMambaRadixCache(MambaRadixCache): self._protect_host_node(last_host_node, protect_mamba=False) - # Allocate host KV memory - host_indices = self._alloc_with_evict( - self.cache_controller.mem_pool_host, - prefetch_length, - self.evict_host, - ) - if host_indices is None: - # truncate the prefetch length to the page-aligned available host size - available_size = self.cache_controller.mem_pool_host.available_size() - prefetch_length = available_size - (available_size % self.page_size) - if prefetch_length < self.prefetch_threshold: - self._release_host_node(last_host_node, release_mamba=False) - return - new_input_tokens = new_input_tokens[:prefetch_length] - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - - if host_indices is None: - self._release_host_node(last_host_node, release_mamba=False) - return - # Allocate host mamba slot extra_pools = self.mamba_prefetch_alloc(new_input_tokens, last_hash) if extra_pools is None: - self.cache_controller.mem_pool_host.free(host_indices) self._release_host_node(last_host_node, release_mamba=False) return @@ -1790,7 +1827,6 @@ class HiMambaRadixCache(MambaRadixCache): operation = self.cache_controller.prefetch( req_id, - host_indices, new_input_tokens, last_hash, prefix_keys, @@ -1799,7 +1835,7 @@ class HiMambaRadixCache(MambaRadixCache): self.ongoing_prefetch[req_id] = ( last_host_node, new_input_tokens, - host_indices, + None, operation, ) self.cache_controller.prefetch_tokens_occupied += len(new_input_tokens) @@ -1812,12 +1848,17 @@ class HiMambaRadixCache(MambaRadixCache): req_id ] - if operation.host_indices is None: - return True - if not self.can_terminate_prefetch(operation): return False + if operation.host_indices is None: + # Stopping before host memory was committed (best_effort, timeout, + # or still mid-query): signal the worker to stop, then release the request. + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(req_id) + return True + + host_indices = operation.host_indices completed_tokens, hash_value = self.cache_controller.terminate_prefetch( operation ) @@ -1957,6 +1998,8 @@ class HiMambaRadixCache(MambaRadixCache): last_host_node, token_ids, host_indices, operation = self.ongoing_prefetch[rid] if operation.host_indices is None: + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(rid) return completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 9391ad944..392802451 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -524,15 +524,15 @@ class HiRadixCache(RadixCache): try: for req_id, info in list(self.ongoing_prefetch.items()): try: - last_host_node, token_ids, host_indices, _operation = info + last_host_node, prefetch_key, _operation = info except Exception: # Unexpected shape; just drop it. self.ongoing_prefetch.pop(req_id, None) continue try: - if host_indices is not None: - cc.mem_pool_host.free(host_indices) + if _operation.host_indices is not None: + cc.mem_pool_host.free(_operation.host_indices) except Exception: logger.exception( "Failed to free host indices for prefetch %s", req_id @@ -546,7 +546,7 @@ class HiRadixCache(RadixCache): ) try: - cc.prefetch_tokens_occupied -= len(token_ids) + cc.prefetch_tokens_occupied -= len(prefetch_key) if cc.prefetch_tokens_occupied < 0: cc.prefetch_tokens_occupied = 0 except Exception: @@ -577,6 +577,7 @@ class HiRadixCache(RadixCache): """ self._drain_storage_control_queues_impl( n_revoke=None, + n_storage_hit=0, n_backup=None, n_release=None, log_metrics=False, @@ -585,6 +586,7 @@ class HiRadixCache(RadixCache): def _drain_storage_control_queues_impl( self, n_revoke: Optional[int], + n_storage_hit: Optional[int], n_backup: Optional[int], n_release: Optional[int], log_metrics: bool, @@ -603,13 +605,54 @@ class HiRadixCache(RadixCache): def _drain_revoke(): for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): - info = self.ongoing_prefetch.pop(req_id, None) - if info is not None: - last_host_node, token_ids, _, _ = info - last_host_node.release_host() - cc.prefetch_tokens_occupied -= len(token_ids) - if cc.prefetch_tokens_occupied < 0: - cc.prefetch_tokens_occupied = 0 + self._revoke_pending_prefetch(req_id) + + def _drain_and_alloc_storage_hit(): + # The L3 hit count is now known, so reserve exactly that much host + # memory (this is the whole point: no over-allocation up front). + # NOTE: alloc/evict here is rank-local but deterministic across TP + # ranks without extra synchronization: host pool mutations only + # happen on the scheduler thread at lockstep points (releases are + # page-granular and drained by the TP-min count), so every rank + # reaches the same success / fallback / revoke decision. + for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit): + req_id = operation.request_id + info = self.ongoing_prefetch.get(req_id) + if info is None: + # request already aborted/cleaned up, skip + continue + if operation.is_terminated(): + # request was aborted while the storage query was in flight + self._revoke_pending_prefetch(req_id) + continue + + alloc_len = operation.storage_hit_count + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self.evict_host(alloc_len) + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + # Memory-pressure fallback: a shorter page-aligned prefix. + available_size = cc.mem_pool_host.available_size() + alloc_len = min( + operation.storage_hit_count, + available_size - (available_size % self.page_size), + ) + if alloc_len >= self.prefetch_threshold: + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self._revoke_pending_prefetch(req_id) + logger.debug( + f"Revoking prefetch for request {req_id} due to host memory allocation failure." + ) + continue + + operation.storage_hit_count = alloc_len + operation.hash_value = operation.hash_value[ + : alloc_len // self.page_size + ] + operation.host_indices = host_indices + cc.prefetch_buffer.put(operation) def _drain_backup(): for operation in _drain_queue(cc.ack_backup_queue, n_backup): @@ -631,6 +674,7 @@ class HiRadixCache(RadixCache): cc.mem_pool_host.free(host_indices) _drain_revoke() + _drain_and_alloc_storage_hit() _drain_backup() _drain_release() @@ -1370,7 +1414,6 @@ class HiRadixCache(RadixCache): extra_kwargs["pool_transfers"] = self._get_extra_pools().get("extra_pools") operation = prefetch_op_cls( "__storage_hit_query__", - self.cache_controller.mem_pool_host.get_dummy_flat_data_page()[:0], prefetch_key, last_hash, prefix_keys, @@ -1419,6 +1462,7 @@ class HiRadixCache(RadixCache): qsizes = torch.tensor( [ cc.prefetch_revoke_queue.qsize(), + cc.prefetch_hit_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), ], @@ -1426,9 +1470,10 @@ class HiRadixCache(RadixCache): ) self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN) - n_revoke, n_backup, n_release = map(int, qsizes.tolist()) + n_revoke, n_storage_hit, n_backup, n_release = map(int, qsizes.tolist()) self._drain_storage_control_queues_impl( n_revoke=n_revoke, + n_storage_hit=n_storage_hit, n_backup=n_backup, n_release=n_release, log_metrics=True, @@ -1482,24 +1527,34 @@ class HiRadixCache(RadixCache): can_terminate = can_terminate or operation_terminated return can_terminate + def _revoke_pending_prefetch(self, req_id: str): + info = self.ongoing_prefetch.pop(req_id, None) + if info is None: + return + last_host_node, prefetch_key, _ = info + last_host_node.release_host() + cc = self.cache_controller + cc.prefetch_tokens_occupied = max( + 0, cc.prefetch_tokens_occupied - len(prefetch_key) + ) + def check_prefetch_progress(self, req_id: str) -> bool: if req_id not in self.ongoing_prefetch: # there is no ongoing prefetch for this request or it has been revoked return True - # todo: more policies for prefetch progress such as timeout - # the current policy is to prefetch with best effort and terminate when queuing is over - last_host_node, prefetch_key, host_indices, operation = self.ongoing_prefetch[ - req_id - ] - - if operation.host_indices is None: - # prefetch has not been issued due to insufficient host memory - return True + last_host_node, prefetch_key, operation = self.ongoing_prefetch[req_id] if not self.can_terminate_prefetch(operation): return False + if operation.host_indices is None: + # Stopping before host memory was committed (best_effort, timeout, or + # still mid-query): signal the worker to stop, then release the request. + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(req_id) + return True + completed_tokens, hash_value = self.cache_controller.terminate_prefetch( operation ) @@ -1513,7 +1568,7 @@ class HiRadixCache(RadixCache): ) min_completed_tokens = completed_tokens_tensor.item() fetched_key = prefetch_key[:min_completed_tokens] - written_indices = host_indices[:min_completed_tokens] + written_indices = operation.host_indices[:min_completed_tokens] matched_length = self._insert_helper_host( last_host_node, fetched_key, @@ -1521,9 +1576,11 @@ class HiRadixCache(RadixCache): hash_value[: min_completed_tokens // self.page_size], ) - self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) + self.cache_controller.mem_pool_host.free( + operation.host_indices[:matched_length] + ) self.cache_controller.append_host_mem_release( - host_indices[min_completed_tokens:completed_tokens] + operation.host_indices[min_completed_tokens:completed_tokens] ) last_host_node.release_host() del self.ongoing_prefetch[req_id] @@ -1542,9 +1599,7 @@ class HiRadixCache(RadixCache): if req_id not in self.ongoing_prefetch: return - _, _, _, operation = self.ongoing_prefetch[req_id] - if operation.host_indices is None: - return + _, _, operation = self.ongoing_prefetch[req_id] operation.mark_terminate() def pop_prefetch_loaded_tokens(self, req_id: str) -> int: @@ -1612,28 +1667,11 @@ class HiRadixCache(RadixCache): return last_host_node.protect_host() - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - if host_indices is None: - self.evict_host(prefetch_length) - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - if host_indices is None: - available_size = self.cache_controller.mem_pool_host.available_size() - prefetch_length = available_size - (available_size % self.page_size) - if prefetch_length >= self.prefetch_threshold: - prefetch_key = prefetch_key[:prefetch_length] - host_indices = self.cache_controller.mem_pool_host.alloc( - prefetch_length - ) - if host_indices is None: - last_host_node.release_host() - return - else: - last_host_node.release_host() - # no sufficient host memory for prefetch - return + # NOTE: host_indices is no longer pre-allocated here. It is allocated + # lazily in _drain_and_alloc_storage_hit() once the L3 storage hit count is known, + # so we only reserve host memory for pages that actually hit. operation = self.cache_controller.prefetch( req_id, - host_indices, prefetch_key, last_hash, prefix_keys, @@ -1642,7 +1680,6 @@ class HiRadixCache(RadixCache): self.ongoing_prefetch[req_id] = ( last_host_node, prefetch_key, - host_indices, operation, ) self.cache_controller.prefetch_tokens_occupied += len(prefetch_key) @@ -1838,15 +1875,17 @@ class HiRadixCache(RadixCache): if rid not in self.ongoing_prefetch: return - last_host_node, prefetch_key, host_indices, operation = self.ongoing_prefetch[ - rid - ] + last_host_node, prefetch_key, operation = self.ongoing_prefetch[rid] if operation.host_indices is None: + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(rid) return completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) self._barrier_attn_groups() last_host_node.release_host() del self.ongoing_prefetch[rid] - self.cache_controller.append_host_mem_release(host_indices[:completed_tokens]) + self.cache_controller.append_host_mem_release( + operation.host_indices[:completed_tokens] + ) self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) 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 c46c2d566..d6c20b821 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 @@ -122,7 +122,6 @@ class PrefetchOperation(StorageOperation): def __init__( self, request_id: str, - host_indices: torch.Tensor, token_ids: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, @@ -131,9 +130,10 @@ class PrefetchOperation(StorageOperation): self.request_id = request_id self._lock = threading.Lock() self._terminated_flag = False + self.storage_hit_count = 0 self.start_time = time.monotonic() super().__init__( - host_indices, + None, token_ids, last_hash, prefix_keys=prefix_keys, @@ -559,7 +559,6 @@ class HybridCacheController(BaseHiCacheController): def prefetch( self, request_id: str, - host_indices: torch.Tensor, new_input_tokens: List[int], last_hash: Optional[str] = None, prefix_keys: Optional[List[str]] = None, @@ -567,7 +566,6 @@ class HybridCacheController(BaseHiCacheController): ) -> PrefetchOperation: operation = PrefetchOperation( request_id, - host_indices, new_input_tokens, last_hash, prefix_keys=prefix_keys, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 1b2bfc7a2..6d3c750b3 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -1789,7 +1789,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if phase == CacheTransferPhase.BACKUP_HOST else indices_source.host_indices ) - if indices is None or len(indices) == 0: + defer_kv_sidecar = ( + phase == CacheTransferPhase.PREFETCH + and spec.indices_from_pool == PoolName.KV + ) + if (indices is None or len(indices) == 0) and not defer_kv_sidecar: continue transfers.append( PoolTransfer( @@ -1889,25 +1893,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return anchor_lock_params = self.inc_host_lock_ref(last_host_node).to_dec_params() - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - if host_indices is None: - self.evict_host(prefetch_length) - host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length) - if host_indices is None: - available_size = self.cache_controller.mem_pool_host.available_size() - prefetch_length = available_size - (available_size % self.page_size) - if prefetch_length >= self.prefetch_threshold: - prefetch_key = prefetch_key[:prefetch_length] - host_indices = self.cache_controller.mem_pool_host.alloc( - prefetch_length - ) - else: - self.dec_host_lock_ref(last_host_node, anchor_lock_params) - return - if host_indices is None: - self.dec_host_lock_ref(last_host_node, anchor_lock_params) - return - comp_xfers: dict[ComponentType, list[PoolTransfer]] = {} alloc_failed = False for comp in self._components_tuple: @@ -1925,13 +1910,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): break if transfers: comp_xfers[comp.component_type] = transfers - kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=host_indices) + kv_xfer = PoolTransfer(name=PoolName.KV, host_indices=None) sidecar_xfers = self._build_sidecar_transfers( CacheTransferPhase.PREFETCH, kv_xfer, comp_xfers ) if alloc_failed: self.cache_controller.append_host_mem_release( - host_indices=host_indices, extra_pools=[x for xfers in comp_xfers.values() for x in xfers], ) self.dec_host_lock_ref(last_host_node, anchor_lock_params) @@ -1941,7 +1925,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): aux_xfers.extend(sidecar_xfers) operation = self.cache_controller.prefetch( req_id, - host_indices, prefetch_key, last_hash, prefix_keys, @@ -1950,7 +1933,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): self.ongoing_prefetch[req_id] = _OngoingPrefetch( last_host_node, prefetch_key, - host_indices, + None, operation, anchor_lock_params, comp_xfers, @@ -2012,10 +1995,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): anchor_lock_params, comp_xfers, ) = self.ongoing_prefetch[req_id] - if operation.host_indices is None: - return True if not self.can_terminate_prefetch(operation): return False + if operation.host_indices is None: + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(req_id) + return True completed_tokens, hash_value = self.cache_controller.terminate_prefetch( operation @@ -2082,8 +2067,6 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): if req_id not in self.ongoing_prefetch: return operation = self.ongoing_prefetch[req_id].operation - if operation.host_indices is None: - return operation.mark_terminate() def pop_prefetch_loaded_tokens(self, req_id: str) -> int: @@ -2103,6 +2086,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): comp_xfers, ) = self.ongoing_prefetch[rid] if operation.host_indices is None: + self.cache_controller.terminate_prefetch(operation) + self._revoke_pending_prefetch(rid) return completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) @@ -2115,9 +2100,31 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): ) self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) + def _revoke_pending_prefetch(self, req_id: str) -> None: + info = self.ongoing_prefetch.pop(req_id, None) + if info is None: + return + ( + last_host_node, + prefetch_key, + _host_indices, + _operation, + anchor_lock_params, + comp_xfers, + ) = info + cc = self.cache_controller + cc.append_host_mem_release( + extra_pools=[x for xfers in comp_xfers.values() for x in xfers] + ) + self.dec_host_lock_ref(last_host_node, anchor_lock_params) + cc.prefetch_tokens_occupied = max( + 0, cc.prefetch_tokens_occupied - len(prefetch_key) + ) + def _drain_storage_control_queues_impl( self, n_revoke: Optional[int], + n_storage_hit: Optional[int], n_backup: Optional[int], n_release: Optional[int], extra_release_counts: Optional[dict[PoolName, int]], @@ -2136,28 +2143,46 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): yield item def _drain_revoke(): - drained = 0 for req_id in _drain_queue(cc.prefetch_revoke_queue, n_revoke): - info = self.ongoing_prefetch.pop(req_id, None) + self._revoke_pending_prefetch(req_id) + + def _drain_and_alloc_storage_hit(): + for operation in _drain_queue(cc.prefetch_hit_queue, n_storage_hit): + req_id = operation.request_id + info = self.ongoing_prefetch.get(req_id) if info is None: + # request already aborted/cleaned up, skip continue - drained += 1 - ( - last_host_node, - prefetch_key, - _host_indices, - _operation, - anchor_lock_params, - comp_xfers, - ) = info - cc.append_host_mem_release( - extra_pools=[x for xfers in comp_xfers.values() for x in xfers] - ) - self.dec_host_lock_ref(last_host_node, anchor_lock_params) - cc.prefetch_tokens_occupied -= len(prefetch_key) - if cc.prefetch_tokens_occupied < 0: - cc.prefetch_tokens_occupied = 0 - return drained + if operation.is_terminated(): + # request was aborted while the storage query was in flight + self._revoke_pending_prefetch(req_id) + continue + + alloc_len = operation.storage_hit_count + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self.evict_host(alloc_len) + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + # Memory-pressure fallback: a shorter page-aligned prefix. + available_size = cc.mem_pool_host.available_size() + alloc_len = min( + operation.storage_hit_count, + available_size - (available_size % self.page_size), + ) + if alloc_len >= self.prefetch_threshold: + host_indices = cc.mem_pool_host.alloc(alloc_len) + if host_indices is None: + self._revoke_pending_prefetch(req_id) + continue + + operation.storage_hit_count = alloc_len + operation.hash_value = operation.hash_value[ + : alloc_len // self.page_size + ] + operation.host_indices = host_indices + self.ongoing_prefetch[req_id] = info._replace(host_indices=host_indices) + cc.prefetch_buffer.put(operation) def _drain_backup(): drained = 0 @@ -2208,6 +2233,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): return drained _drain_revoke() + _drain_and_alloc_storage_hit() _drain_backup() _drain_release() _drain_extra_release() @@ -2218,6 +2244,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): extra_pool_names = list(extra_release_queues) local_qsize_list = [ cc.prefetch_revoke_queue.qsize(), + cc.prefetch_hit_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), *[ @@ -2231,13 +2258,14 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache): ) self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN) qsize_list = list(map(int, qsizes.tolist())) - n_revoke, n_backup, n_release = qsize_list[:3] + n_revoke, n_storage_hit, n_backup, n_release = qsize_list[:4] extra_release_counts = { pool_name: count - for pool_name, count in zip(extra_pool_names, qsize_list[3:]) + for pool_name, count in zip(extra_pool_names, qsize_list[4:]) } self._drain_storage_control_queues_impl( n_revoke=n_revoke, + n_storage_hit=n_storage_hit, n_backup=n_backup, n_release=n_release, extra_release_counts=extra_release_counts, 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 c46487251..d613275dd 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 @@ -380,7 +380,6 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): def prefetch( self, request_id, - host_indices, new_input_tokens, last_hash=None, prefix_keys=None, @@ -388,7 +387,6 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): ): self.prefetch_args = ( request_id, - host_indices, new_input_tokens, last_hash, prefix_keys, @@ -400,7 +398,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): cache.cache_controller = controller cache.prefetch_from_storage("req", cache.root_node, tokens) - _, _, storage_key, _, _, _ = controller.prefetch_args + _, storage_key, _, _, _ = controller.prefetch_args self.assertIsInstance(storage_key, RadixKey) self.assertTrue(storage_key.is_bigram) self.assertEqual(len(storage_key), len(tokens) - 1) @@ -2251,6 +2249,10 @@ class UnifiedRadixCacheSuite: def _run_prefetch_to_completion(self, cache, req_id, timeout: float = 10.0): deadline = time.time() + timeout while time.time() < deadline: + # Host memory is reserved (and IO started) by the scheduler-thread + # drain once the L3 hit count is known, so pump it like the real + # scheduler loop does (check_hicache_events before progress checks). + cache.drain_storage_control_queues() if cache.check_prefetch_progress(req_id): return time.sleep(0.01) @@ -2385,15 +2387,15 @@ class UnifiedRadixCacheSuite: comp_xfers = info[-1] names = [t.name for xfers in comp_xfers.values() for t in xfers] if PoolName.SWA in names: - return 1 + names.index(PoolName.SWA) - return None + return 1 + names.index(PoolName.SWA), 1 + len(names) + return None, None def fake(tensor, op=None, group=None): if op == dist.ReduceOp.MIN: min_sizes.append(tensor.numel()) if drop_swa: - idx = swa_packed_index() - if idx is not None and idx < tensor.numel(): + idx, packed_numel = swa_packed_index() + if idx is not None and tensor.numel() == packed_numel: tensor[idx] = 0 return None