diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 240ee4d5f..d3f2a6128 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -13,9 +13,11 @@ See the License for the specific language governing permissions and limitations under the License. """ + import logging import threading import time +from dataclasses import dataclass from queue import Empty, Queue from typing import TYPE_CHECKING, Callable, List, NamedTuple, Optional @@ -27,6 +29,7 @@ from sglang.srt.mem_cache.hicache_storage import ( HiCacheStorageExtraInfo, PoolName, PoolTransfer, + count_pool_hits, ) if TYPE_CHECKING: @@ -186,6 +189,32 @@ class HiCacheAck(NamedTuple): num_bytes: int = 0 +@dataclass +class PrefetchAck: + """ACK for prefetch operation. + + A sequence of PrefetchAck is sent to the scheduler thread via ack_prefetch_queue, + indicating progress or completion of the prefetch operation. + + For example, a prefetch operation may results into the following sequence of PrefetchAck: + + 1. PrefetchAck(completed_tokens = 128) + 2. PrefetchAck(completed_tokens = 256) + 3. PrefetchAck(pool_hits={INDEXER: 256}) + 4. PrefetchAck(completed_req = True) + + The last PrefetchAck always specifies completed_req = True. + """ + + rid: str + operation: PrefetchOperation + # Number of hits in KV pool. + completed_tokens: Optional[int] = None + # Number of hits in extra pools. + pool_hits: Optional[dict[str, int]] = None + completed_req: Optional[bool] = None + + class StorageOperation: counter = 0 @@ -244,19 +273,13 @@ class PrefetchOperation(StorageOperation): super().__init__(None, token_ids, last_hash, prefix_keys=prefix_keys) - def increment(self, num_tokens: int): - with self._lock: - if self._terminated_flag: - return False - self.completed_tokens += num_tokens - return True - def mark_terminate(self): with self._lock: self._terminated_flag = True def is_terminated(self) -> bool: - return self._terminated_flag + with self._lock: + return self._terminated_flag class HiCacheController: @@ -285,7 +308,8 @@ class HiCacheController: self.attn_cp_group = attn_cp_group self.attn_tp_group = attn_tp_group self.pp_group = pp_group - self.prefetch_sync_groups: List[torch.distributed.ProcessGroup] = [] + self.prefetch_hits_sync_groups: List[torch.distributed.ProcessGroup] = [] + self.prefetch_completion_sync_groups: List[torch.distributed.ProcessGroup] = [] self.mem_pool_device_allocator = token_to_kv_pool_allocator mem_pool_device = token_to_kv_pool_allocator.get_kvcache() from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool @@ -322,7 +346,10 @@ class HiCacheController: self.storage_stop_event = threading.Event() # Storage control queues, (re)created whenever the storage threads start. + self.prefetch_buffer: Optional[Queue[PrefetchOperation]] = None + self.prefetch_sync_queue: Optional[Queue[PrefetchAck]] = None self.prefetch_hit_queue: Optional[Queue[StorageOperation]] = None + self.ack_prefetch_queue = Queue[PrefetchAck]() self.ack_backup_queue: Optional[Queue[StorageOperation]] = None self.host_mem_release_queue: Optional[Queue[torch.Tensor]] = None @@ -369,16 +396,18 @@ class HiCacheController: ) return 0, 1 - def _create_prefetch_sync_groups(self) -> None: + def _create_sync_groups(self) -> List[torch.distributed.ProcessGroup]: from sglang.srt.distributed.parallel_state import create_custom_parallel_group - self.prefetch_sync_groups = [] + groups: List[torch.distributed.ProcessGroup] = [] seen_rank_sets = set() if self.attn_cp_group is not None or self.attn_tp_group is not None: base_groups = [self.attn_cp_group, self.attn_tp_group] else: base_groups = [self.tp_group] + if self.pp_group is not None: + base_groups.append(self.pp_group) for group in base_groups: if group is None or torch.distributed.get_world_size(group=group) == 1: @@ -387,22 +416,29 @@ class HiCacheController: if group_ranks in seen_rank_sets: continue seen_rank_sets.add(group_ranks) - self.prefetch_sync_groups.append( + groups.append( create_custom_parallel_group( group_ranks=list(group_ranks), backend="gloo" ) ) + return groups - def _destroy_prefetch_sync_groups(self) -> None: - for group in self.prefetch_sync_groups: + def _destroy_sync_groups( + self, groups: List[torch.distributed.ProcessGroup] + ) -> None: + for group in groups: try: torch.distributed.destroy_process_group(group) except Exception: pass - self.prefetch_sync_groups = [] - def _all_reduce_prefetch_groups(self, tensor: torch.Tensor, op) -> None: - for group in self.prefetch_sync_groups: + def _all_reduce( + self, + tensor: torch.Tensor, + op, + groups: List[torch.distributed.ProcessGroup], + ) -> None: + for group in groups: torch.distributed.all_reduce(tensor, op=op, group=group) def _start_storage_threads(self): @@ -416,17 +452,27 @@ class HiCacheController: self.prefetch_thread = threading.Thread( target=self.prefetch_thread_func, daemon=True ) + self.prefetch_io_aux_thread = threading.Thread( + target=self.prefetch_io_aux_func, daemon=True + ) + self.prefetch_sync_thread = threading.Thread( + target=self.prefetch_sync_thread_func, daemon=True + ) self.backup_thread = threading.Thread( target=self.backup_thread_func, daemon=True ) self.prefetch_queue = Queue() self.backup_queue = Queue() - + self.prefetch_buffer = Queue() + self.prefetch_sync_queue = Queue() self.prefetch_hit_queue = Queue() + self.ack_prefetch_queue = Queue() self.ack_backup_queue = Queue() self.host_mem_release_queue = Queue() self.prefetch_thread.start() + self.prefetch_io_aux_thread.start() + self.prefetch_sync_thread.start() self.backup_thread.start() def _stop_storage_threads(self): @@ -449,6 +495,8 @@ class HiCacheController: self.backup_queue.put_nowait(None) if hasattr(self, "prefetch_buffer"): self.prefetch_buffer.put_nowait(None) + if hasattr(self, "prefetch_sync_queue"): + self.prefetch_sync_queue.put_nowait(None) except Exception: pass @@ -460,6 +508,8 @@ class HiCacheController: threads.append(self.backup_thread) if hasattr(self, "prefetch_io_aux_thread"): threads.append(self.prefetch_io_aux_thread) + if hasattr(self, "prefetch_sync_thread"): + threads.append(self.prefetch_sync_thread) for t in threads: try: @@ -542,7 +592,8 @@ class HiCacheController: # Use dedicated gloo groups so storage prefetch sync is isolated # from other collectives and consistent across CPxTP participants. - self._create_prefetch_sync_groups() + self.prefetch_hits_sync_groups = self._create_sync_groups() + self.prefetch_completion_sync_groups = self._create_sync_groups() # Select the get and set functions self.page_get_func = self._generic_page_get @@ -569,7 +620,10 @@ class HiCacheController: self._stop_storage_threads() except Exception: pass - self._destroy_prefetch_sync_groups() + self._destroy_sync_groups(self.prefetch_hits_sync_groups) + self._destroy_sync_groups(self.prefetch_completion_sync_groups) + self.prefetch_hits_sync_groups = [] + self.prefetch_completion_sync_groups = [] try: if ( hasattr(self, "storage_backend") @@ -609,7 +663,11 @@ class HiCacheController: raise RuntimeError("Stop storage threads failed; detach aborted.") from e # Best-effort destroy process groups created for storage ops. - self._destroy_prefetch_sync_groups() + self._destroy_sync_groups( + self.prefetch_hits_sync_groups + self.prefetch_completion_sync_groups + ) + self.prefetch_hits_sync_groups = [] + self.prefetch_completion_sync_groups = [] # Best-effort close (some backends rely on GC/destructor). try: @@ -704,10 +762,15 @@ class HiCacheController: self.ack_load_queue.clear() if self.enable_storage: self.prefetch_thread.join() + self.prefetch_io_aux_thread.join() + self.prefetch_sync_thread.join() self.backup_thread.join() self.prefetch_queue.queue.clear() self.backup_queue.queue.clear() + self.prefetch_buffer.queue.clear() + self.prefetch_sync_queue.queue.clear() self.prefetch_hit_queue.queue.clear() + self.ack_prefetch_queue.queue.clear() self.ack_backup_queue.queue.clear() self.host_mem_release_queue.queue.clear() self.prefetch_tokens_occupied = 0 @@ -718,10 +781,18 @@ class HiCacheController: self.prefetch_thread = threading.Thread( target=self.prefetch_thread_func, daemon=True ) + self.prefetch_io_aux_thread = threading.Thread( + target=self.prefetch_io_aux_func, daemon=True + ) + self.prefetch_sync_thread = threading.Thread( + target=self.prefetch_sync_thread_func, daemon=True + ) self.backup_thread = threading.Thread( target=self.backup_thread_func, daemon=True ) self.prefetch_thread.start() + self.prefetch_io_aux_thread.start() + self.prefetch_sync_thread.start() self.backup_thread.start() def write( @@ -984,6 +1055,14 @@ class HiCacheController: return operation def terminate_prefetch(self, operation): + """ + Request to terminate a prefetch operation. + + Must be called in the scheduler thread. + + Asynchronous prefetch tasks may be running in background threads. When all prefetch + tasks are terminated, a PrefetchAck with completed_req=True will be sent to ack_prefetch_queue. + """ operation.mark_terminate() return operation.completed_tokens, operation.hash_value @@ -996,7 +1075,7 @@ class HiCacheController: def _page_get_zero_copy( self, operation, hash_values, host_indices, extra_info=None - ): + ) -> int: results = self.storage_backend.batch_get_v1( hash_values, host_indices, extra_info ) @@ -1007,61 +1086,125 @@ class HiCacheController: f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}." ) break - inc += self.page_size - operation.increment(inc) + inc += 1 + return inc # todo: deprecate - def _generic_page_get(self, operation, hash_values, host_indices, extra_info=None): + def _generic_page_get( + self, operation, hash_values, host_indices, extra_info=None + ) -> int: dummy_page_dst = [ self.mem_pool_host.get_dummy_flat_data_page() for _ in hash_values ] page_data = self.storage_backend.batch_get(hash_values, dummy_page_dst) if page_data is None: - return + return 0 + count = 0 for i in range(len(hash_values)): if page_data[i] is None: logger.warning( f"Prefetch operation {operation.request_id} failed to retrieve page {hash_values[i]}." ) break - # Must set the data before increasing the completed tokens. - # Otherwise this page may be read before being set. + if operation.is_terminated(): + break self.mem_pool_host.set_from_flat_data_page( host_indices[i * self.page_size], page_data[i], ) - if not operation.increment(self.page_size): - break # Operation terminated by controller + count += 1 + return count - def _page_transfer(self, operation): + def _page_transfer(self, operation: PrefetchOperation) -> int: # Transfer batch by batch prefix_keys = operation.prefix_keys + kv_derived_transfers = [ + transfer + for transfer in getattr(operation, "pool_transfers", None) or [] + if transfer.indices_from_pool == PoolName.KV + ] + all_success = True + completed_pages = 0 for i in range(0, len(operation.hash_value), STORAGE_BATCH_SIZE): - batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE] - batch_host_indices = operation.host_indices[ - i * self.page_size : (i + len(batch_hashes)) * self.page_size + # When an error is occurred, we should keep looping and produce the same number of + # PrefetchAck as other ranks do, because prefetch_sync_thread (i.e. consumer of + # prefetch_sync_queue) perform reduce on the results. This is so tricky. + if all_success and operation.is_terminated(): + all_success = False + if all_success: + batch_hashes = operation.hash_value[i : i + STORAGE_BATCH_SIZE] + batch_host_indices = operation.host_indices[ + i * self.page_size : (i + len(batch_hashes)) * self.page_size + ] + + # Best-effort draft L3 read before publishing target completion. + # Otherwise wait_complete can race and load back target KV before + # draft KV reaches host memory. + if self.has_draft: + self._draft_page_get(batch_hashes, batch_host_indices) + + # Get one batch token, and update the completed_tokens if succeed + extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) + + hit_pages = self._page_transfer_kv_batch( + operation, + batch_hashes, + batch_host_indices, + extra_info, + kv_derived_transfers, + ) + # Check termination + if hit_pages != len(batch_hashes): + all_success = False + if prefix_keys and len(prefix_keys) > 0: + prefix_keys += batch_hashes + completed_pages += hit_pages + ack = PrefetchAck( + rid=operation.request_id, + completed_tokens=completed_pages * self.page_size, + operation=operation, + ) + self.prefetch_sync_queue.put(ack) + return completed_pages + + def _page_transfer_kv_batch( + self, + operation: PrefetchOperation, + batch_hashes: List[str], + batch_host_indices: torch.Tensor, + extra_info: HiCacheStorageExtraInfo, + kv_derived_transfers: List[PoolTransfer], + ) -> int: + """Read a single batch from KV and KV-derived pools (e.g. indexer pool). + + Return the number of hit pages. If the hits from KV and KV-derived pools differ, + clamp to the minimal number of hits. + + Here, "batch" means a single unit of L3 read, not a "batch" in model forward. + """ + # Read from KV pool. + kv_hits = self.page_get_func( + operation, batch_hashes, batch_host_indices, extra_info + ) + + # Read from KV-derived sidecar pools, if any. + sidecar_hits: dict[str, int] = {} + if len(kv_derived_transfers) > 0: + current_kv_derived_transfers = [ + PoolTransfer( + name=transfer.name, + host_indices=batch_host_indices, + keys=batch_hashes, + ) + for transfer in kv_derived_transfers ] + sidecar_results = self.storage_backend.batch_get_v2( + current_kv_derived_transfers + ) + sidecar_hits = count_pool_hits(sidecar_results) - # Best-effort draft L3 read before publishing target completion. - # Otherwise wait_complete can race and load back target KV before - # draft KV reaches host memory. - if self.has_draft: - self._draft_page_get(batch_hashes, batch_host_indices) - - prev_completed_tokens = operation.completed_tokens - # Get one batch token, and update the completed_tokens if succeed - extra_info = HiCacheStorageExtraInfo(prefix_keys=prefix_keys) - self.page_get_func(operation, batch_hashes, batch_host_indices, extra_info) - # Check termination - if ( - operation.completed_tokens - != prev_completed_tokens + len(batch_hashes) * self.page_size - ): - operation.mark_terminate() - break # Some operations fail or operation terminated by controller - - if prefix_keys and len(prefix_keys) > 0: - prefix_keys += batch_hashes + # Clamp to minimal number of hits. + return min([kv_hits, *sidecar_hits.values()]) def prefetch_io_aux_func(self): """ @@ -1073,9 +1216,13 @@ class HiCacheController: if operation is None: continue self._page_transfer(operation) - # operation terminated by controller, release pre-allocated memory - self.append_host_mem_release( - operation.host_indices[operation.completed_tokens :] + + self.prefetch_sync_queue.put( + PrefetchAck( + rid=operation.request_id, + completed_req=True, + operation=operation, + ) ) except Empty: continue @@ -1129,11 +1276,6 @@ class HiCacheController: """ Manage prefetching operations from storage backend to host memory. """ - self.prefetch_buffer = Queue() - self.prefetch_io_aux_thread = threading.Thread( - target=self.prefetch_io_aux_func, daemon=True - ) - self.prefetch_io_aux_thread.start() while (not self.storage_stop_event.is_set()) or not self.prefetch_queue.empty(): try: operation = self.prefetch_queue.get(block=True, timeout=1) @@ -1146,8 +1288,10 @@ class HiCacheController: storage_hit_count_tensor = torch.tensor( storage_hit_count, dtype=torch.int ) - self._all_reduce_prefetch_groups( - storage_hit_count_tensor, torch.distributed.ReduceOp.MIN + self._all_reduce( + storage_hit_count_tensor, + torch.distributed.ReduceOp.MIN, + self.prefetch_hits_sync_groups, ) storage_hit_count = storage_hit_count_tensor.item() @@ -1300,3 +1444,29 @@ class HiCacheController: except Empty: continue + + def prefetch_sync_thread_func(self): + """Synchronize prefetch results across all PP and TP ranks.""" + while not self.storage_stop_event.is_set(): + try: + ack = self.prefetch_sync_queue.get(block=True, timeout=1) + if ack is None: + continue + self._reduce_prefetch_ack(ack) + self.ack_prefetch_queue.put(ack) + except Empty: + continue + + def _reduce_prefetch_ack(self, ack: PrefetchAck) -> None: + """Synchronize all ranks to agree on a PrefetchAck.""" + if ack.completed_tokens is not None: + # Determine the minimal successful prefix of tokens. + completed_tokens_tensor = torch.tensor( + ack.completed_tokens, dtype=torch.int + ) + self._all_reduce( + completed_tokens_tensor, + torch.distributed.ReduceOp.MIN, + self.prefetch_completion_sync_groups, + ) + ack.completed_tokens = completed_tokens_tensor.item() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3c0502a1f..bb2fb4941 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2875,8 +2875,6 @@ class Scheduler( if self.enable_hicache_storage: # Release prefetch events associated with the request self.tree_cache.release_aborted_request(candidate_req.rid) - elif self.enable_hierarchical_cache: - self.tree_cache.terminate_prefetch(candidate_req.rid) self.waiting_queue.pop(idx) req_to_abort = candidate_req message = "The request is aborted by a higher priority request." diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index c9829e93a..ef0759ff2 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -133,18 +133,20 @@ class PoolTransferResult: """Accumulate kv_hit_pages across batches (max = last successful batch).""" 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: + def update_extra_pool_hit_pages(self, results: dict[str, int]) -> None: """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: (rs.index(False) if False in rs else len(rs)) - for name, rs in results.items() - } - ) + self.extra_pool_hit_pages.update(results) + + +def count_pool_hits(results: dict[str, List[bool]]) -> dict[str, int]: + return { + name: (rs.index(False) if False in rs else len(rs)) + for name, rs in results.items() + } class HiCacheStorage(ABC): diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 94fb2a6b3..c91fe88af 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -5,9 +5,9 @@ import heapq import json import logging import os +import queue import threading import time -from queue import Empty from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -228,15 +228,6 @@ class HiRadixCache(RadixCache): if not reduced and self.tp_world_size > 1: torch.distributed.all_reduce(tensor, op=op, group=self.tp_group) - def _barrier_attn_groups(self): - waited = False - for group in (self.attn_cp_group, self.attn_tp_group): - if group is not None and torch.distributed.get_world_size(group=group) > 1: - torch.distributed.barrier(group=group) - waited = True - if not waited and self.tp_world_size > 1: - torch.distributed.barrier(group=self.tp_group) - def _drain_async_work(self): """ Block until all outstanding async sends are consumed, then clear. @@ -585,6 +576,7 @@ class HiRadixCache(RadixCache): """ self._drain_storage_control_queues_impl( n_storage_hit=0, + n_ack_prefetch=None, n_backup=None, n_release=None, log_metrics=False, @@ -593,21 +585,27 @@ class HiRadixCache(RadixCache): def _drain_storage_control_queues_impl( self, n_storage_hit: Optional[int], + n_ack_prefetch: Optional[int], n_backup: Optional[int], n_release: Optional[int], log_metrics: bool, ): cc = self.cache_controller - def _drain_queue(q, limit: Optional[int]): - drained = 0 - while limit is None or drained < limit: - try: - item = q.get_nowait() - except Empty: - break - drained += 1 - yield item + def _drain_queue(q: queue.Queue, n: Optional[int]): + """If n is None, consume all items from the queue. + Otherwise, consume n items from the queue. + """ + if n is None: + while not q.empty(): + item = q.get() + yield item + else: + for _ in range(n): + # Block when there is no enough elements. + # All TP/PP ranks must consume the same number of elements. + item = q.get() + yield item def _drain_and_alloc_storage_hit(): # The L3 hit count is now known, so reserve exactly that much host @@ -663,6 +661,25 @@ class HiRadixCache(RadixCache): operation.host_indices = host_indices cc.prefetch_buffer.put(operation) + def _drain_ack_prefetch(): + for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch): + operation = ack.operation + if ack.completed_tokens is not None: + if operation.request_id in self.ongoing_prefetch: + assert operation.completed_tokens <= ack.completed_tokens + operation.completed_tokens = ack.completed_tokens + if ack.pool_hits is not None: + if operation.request_id in self.ongoing_prefetch: + operation.pool_storage_result.update_extra_pool_hit_pages( + ack.pool_hits + ) + operation.pool_transfers_done = True + if ack.completed_req: + if operation.request_id in self.ongoing_prefetch: + self._handle_prefetch_result(operation) + tail = operation.host_indices[operation.completed_tokens :] + self.cache_controller.mem_pool_host.free(tail) + def _drain_backup(): for operation in _drain_queue(cc.ack_backup_queue, n_backup): ack_id = operation.id @@ -683,6 +700,7 @@ class HiRadixCache(RadixCache): cc.mem_pool_host.free(host_indices) _drain_and_alloc_storage_hit() + _drain_ack_prefetch() _drain_backup() _drain_release() @@ -1008,6 +1026,7 @@ class HiRadixCache(RadixCache): storage_queue_sizes = ( ( cache_controller.prefetch_hit_queue.qsize(), + cache_controller.ack_prefetch_queue.qsize(), cache_controller.ack_backup_queue.qsize(), cache_controller.host_mem_release_queue.qsize(), ) @@ -1541,9 +1560,12 @@ class HiRadixCache(RadixCache): self.loading_check(finish_count=load_finish_count) if self.enable_storage and storage_queue_sizes: - n_storage_hit, n_backup, n_release = storage_queue_sizes[:3] + n_storage_hit, n_ack_prefetch, n_backup, n_release = ( + storage_queue_sizes[:4] + ) self._drain_storage_control_queues_impl( n_storage_hit=n_storage_hit, + n_ack_prefetch=n_ack_prefetch, n_backup=n_backup, n_release=n_release, log_metrics=True, @@ -1563,16 +1585,18 @@ class HiRadixCache(RadixCache): qsizes = torch.tensor( [ cc.prefetch_hit_queue.qsize(), + cc.ack_prefetch_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), ], dtype=torch.int, ) - self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN) - n_storage_hit, n_backup, n_release = map(int, qsizes.tolist()) + self._all_reduce(qsizes, torch.distributed.ReduceOp.MIN) + n_storage_hit, n_ack_prefetch, n_backup, n_release = map(int, qsizes.tolist()) self._drain_storage_control_queues_impl( n_storage_hit=n_storage_hit, + n_ack_prefetch=n_ack_prefetch, n_backup=n_backup, n_release=n_release, log_metrics=True, @@ -1585,47 +1609,17 @@ class HiRadixCache(RadixCache): timeout = min(cfg.max, cfg.base + cfg.per_ki_token * num_tokens / 1024) return time.monotonic() - operation.start_time > timeout - def can_terminate_prefetch(self, operation: PrefetchOperation): - can_terminate = True - + def can_terminate_prefetch(self, operation: PrefetchOperation) -> bool: if self.prefetch_stop_policy == "best_effort": - return can_terminate - - if len(operation.hash_value) == 0: - completed = False - else: - completed = ( - operation.completed_tokens == len(operation.hash_value) * self.page_size - ) - + return True if self.prefetch_stop_policy == "wait_complete": - can_terminate = completed + return False elif self.prefetch_stop_policy == "timeout": - can_terminate = completed or self.is_prefetch_timeout(operation) + return self.is_prefetch_timeout(operation) else: # unknown prefetch stop policy, just return True 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( - [1 - int(can_terminate), int(operation_terminated)], - dtype=torch.int, - ) - self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX) - can_terminate = states[0].item() == 0 - operation_terminated = states[1].item() == 1 - # the operation should be terminated if it is already terminated on any TP worker - # or it meets the termination condition on all TP workers - 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: @@ -1642,44 +1636,56 @@ class HiRadixCache(RadixCache): # there is no ongoing prefetch for this request or it has been revoked return True - last_host_node, prefetch_key, operation = self.ongoing_prefetch[req_id] + _, _, operation = self.ongoing_prefetch[req_id] - if not self.can_terminate_prefetch(operation): + # Determine whether or not we should terminate this prefetch request. Make all + # ranks agree on the decision. When running with PP, PPn will follow PP0's decision. + should_terminate = False + if self.pp_rank == 0: + should_terminate = operation.is_terminated() or self.can_terminate_prefetch( + operation + ) + should_terminate_tensor = torch.tensor( + int(should_terminate), dtype=torch.int, device="cpu" + ) + self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX) + should_terminate = should_terminate_tensor.item() == 1 + + if not should_terminate: return False + # Terminate in-flight prefetch. + self.cache_controller.terminate_prefetch(operation) 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 + else: + # Handle partial or full completion. + self._handle_prefetch_result(operation) + return True - completed_tokens, hash_value = self.cache_controller.terminate_prefetch( - operation + def _handle_prefetch_result(self, operation: PrefetchOperation) -> None: + req_id = operation.request_id + # All PP/TP ranks will get the same `min_completed_tokens`, because `completed_tokens` + # and `pool_hits` in their operations are same. No need to sync cross-rank here. + min_completed_tokens = self._clamp_prefetch_result(operation) + logger.debug( + f"Prefetch {req_id} completed with {operation.completed_tokens} tokens" ) - logger.debug(f"Prefetch {req_id} completed with {completed_tokens} tokens") - - min_completed_tokens = self._sync_and_clamp_prefetch_result( - operation, completed_tokens - ) - + last_host_node, prefetch_key, operation = self.ongoing_prefetch.pop(req_id) + host_indices = operation.host_indices fetched_key = prefetch_key[:min_completed_tokens] - written_indices = operation.host_indices[:min_completed_tokens] + written_indices = host_indices[:min_completed_tokens] matched_length = self._insert_helper_host( last_host_node, fetched_key, written_indices, - hash_value[: min_completed_tokens // self.page_size], + operation.hash_value[: min_completed_tokens // self.page_size], ) - self.cache_controller.mem_pool_host.free( - operation.host_indices[:matched_length] - ) - self.cache_controller.append_host_mem_release( - operation.host_indices[min_completed_tokens:completed_tokens] - ) + self.cache_controller.mem_pool_host.free(host_indices[:matched_length]) last_host_node.release_host() - del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key) # Track tokens actually loaded from storage for this request (L3 hits) @@ -1688,21 +1694,19 @@ class HiRadixCache(RadixCache): if self.enable_storage_metrics: self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) + return - return True - - def _sync_and_clamp_prefetch_result( + def _clamp_prefetch_result( self, operation: PrefetchOperation, - completed_tokens: int, ) -> int: - """Sync prefetch results across ATTN groups and decide the usable prefix. + """Determine the minimal number of tokens from full KV hits and sidecar hits. 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. + this is just 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. @@ -1710,27 +1714,17 @@ class HiRadixCache(RadixCache): hit_pages = ( operation.pool_storage_result.extra_pool_hit_pages if pool_transfers else {} ) + completed_tokens = operation.completed_tokens 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 + usable_pages = 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 - - _, _, operation = self.ongoing_prefetch[req_id] - operation.mark_terminate() - def pop_prefetch_loaded_tokens(self, req_id: str) -> int: """ Pop and return the number of tokens loaded from storage for a request. @@ -2017,7 +2011,6 @@ class HiRadixCache(RadixCache): 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( 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 4d8a5e683..94b22e349 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 @@ -19,6 +19,7 @@ from sglang.srt.managers.cache_controller import ( ) from sglang.srt.managers.cache_controller import ( LayerDoneCounter, + PrefetchAck, ) from sglang.srt.managers.cache_controller import ( StorageOperation as BaseStorageOperation, @@ -29,6 +30,7 @@ from sglang.srt.mem_cache.hicache_storage import ( PoolName, PoolTransfer, PoolTransferResult, + count_pool_hits, ) from sglang.srt.mem_cache.l2_transfer import L2Transfer from sglang.srt.mem_cache.memory_pool_host import HostPoolGroup, PoolEntry @@ -78,19 +80,13 @@ class PrefetchOperation(StorageOperation): ) self.pool_transfers_done = not bool(pool_transfers) - def increment(self, num_tokens: int): - with self._lock: - if self._terminated_flag: - return False - self.completed_tokens += num_tokens - return True - def mark_terminate(self): with self._lock: self._terminated_flag = True def is_terminated(self) -> bool: - return self._terminated_flag + with self._lock: + return self._terminated_flag class HybridCacheController(BaseHiCacheController): @@ -640,26 +636,50 @@ class HybridCacheController(BaseHiCacheController): ) return host_indices, device_indices, resolved_pool_transfers - def _page_transfer(self, operation): - # KV pools first — determines actual completed page count - super()._page_transfer(operation) + def _page_transfer(self, operation: PrefetchOperation) -> bool: + # KV pools and KV-derived pools first — determines actual completed page count + kv_completed_pages = super()._page_transfer(operation) + + # Read non-KV derived sidecar pool, e.g. SWA, Mamba. + self._page_transfer_sidecar(operation, kv_completed_pages) + + def _page_transfer_sidecar( + self, operation: PrefetchOperation, kv_completed_pages: int + ) -> None: + if operation.pool_transfers is None: + return # 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 not operation.is_terminated() - and kv_completed_pages == len(operation.hash_value) + pool_hits: dict[str, int] = {} + if not operation.is_terminated() and kv_completed_pages == len( + operation.hash_value ): + # KV-derived sidecar pools are handled in CacheController._page_transfer_kv_batch. + # Only handle non-KV-derived sidecar pools here. + transfers_nonkv = [ + transfer + for transfer in operation.pool_transfers + if transfer.indices_from_pool != PoolName.KV + ] self._sync_trailing_keys( - operation.pool_transfers, operation.hash_value, kv_completed_pages + transfers_nonkv, 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) - operation.pool_transfers_done = True + self._resolve_sidecar_nonkv_derived_pool_transfers(operation) + results = self.storage_backend.batch_get_v2(transfers_nonkv) + pool_hits = count_pool_hits(results) + # Emit PrefetchAck to prefetch_sync_queue, even the operation has been canceled by the + # scheduler thread. The prefetch sync thread expects the same number of PrefetchAck objects + # to perform all_reduce. + self.prefetch_sync_queue.put( + PrefetchAck( + rid=operation.request_id, + operation=operation, + pool_hits=pool_hits, + ) + ) + return def _page_backup(self, operation): # MLA KV is replicated across TP ranks and should still be written only @@ -671,9 +691,11 @@ class HybridCacheController(BaseHiCacheController): ] if backup_transfers: - self._resolve_sidecar_derived_pool_transfers(operation) + self._resolve_sidecar_kv_derived_pool_transfers(operation) + self._resolve_sidecar_nonkv_derived_pool_transfers(operation) results = self.storage_backend.batch_set_v2(backup_transfers) - operation.pool_storage_result.update_extra_pool_hit_pages(results) + pool_hits = count_pool_hits(results) + operation.pool_storage_result.update_extra_pool_hit_pages(pool_hits) if not self.backup_skip: super()._page_backup(operation) @@ -737,7 +759,14 @@ class HybridCacheController(BaseHiCacheController): except Empty: continue - def _resolve_sidecar_derived_pool_transfers(self, operation): + def _resolve_sidecar_kv_derived_pool_transfers(self, operation): + for transfer in operation.pool_transfers: + if transfer.indices_from_pool == PoolName.KV: + transfer.host_indices = operation.host_indices + if transfer.keys is None: + transfer.keys = operation.hash_value + + def _resolve_sidecar_nonkv_derived_pool_transfers(self, operation): for transfer in operation.pool_transfers: if transfer.indices_from_pool is None: continue @@ -760,9 +789,7 @@ class HybridCacheController(BaseHiCacheController): 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 + pass def _sync_trailing_keys( self, @@ -888,3 +915,23 @@ class HybridCacheController(BaseHiCacheController): pool.host_indices = source.host_indices pool.device_indices = source.device_indices return extra_pools + + def _reduce_prefetch_ack(self, ack: PrefetchAck) -> None: + # Handle KV-derived pool. + super()._reduce_prefetch_ack(ack) + + # Handle other sidecar pools, e.g. SWA, Mamba. + if ack.pool_hits is not None: + # "for ... in PoolName" ensures the same order across all ranks. + # On prefetch failure, pool_hits may be empty dict. + packed = torch.tensor( + [ack.pool_hits.get(pool.value, 0) for pool in PoolName], + dtype=torch.int, + ) + self._all_reduce( + packed, + torch.distributed.ReduceOp.MIN, + self.prefetch_completion_sync_groups, + ) + for i, pool in enumerate(PoolName): + ack.pool_hits[pool.value] = packed[i].item() 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 53bb5a8ac..efb8edbd8 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 @@ -410,10 +410,14 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): "Mooncake package does not support ReplicateConfig.group_ids. " "Falling back to the existing batch_put_from path." ) - tp_scale_factor = 1 if storage_config is None else storage_config.tp_size + rank_scale_factor = ( + 1 + if storage_config is None + else (storage_config.tp_size * storage_config.pp_size) + ) - per_tp_global_segment_size = ( - self.config.global_segment_size // tp_scale_factor + per_rank_global_segment_size = ( + self.config.global_segment_size // rank_scale_factor ) # Use the backend tag and model name as a prefix to isolate tenants @@ -510,7 +514,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): ret_code = self.store.setup( client_hostname, self.config.metadata_server, - per_tp_global_segment_size, + per_rank_global_segment_size, DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer self.config.protocol, device_name, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 9276c4ecb..091d22d16 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -5,7 +5,7 @@ import logging import threading import time from dataclasses import replace -from queue import Empty, Queue +from queue import Queue from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, Sequence, TypeVar import torch @@ -35,7 +35,6 @@ from sglang.srt.mem_cache.buffer_mode.storage_existence_cache import ( ) from sglang.srt.mem_cache.common import RetractionBackup from sglang.srt.mem_cache.hicache_storage import ( - PoolHitPolicy, PoolName, PoolTransfer, SidecarPoolSpec, @@ -1717,64 +1716,67 @@ class UnifiedRadixCache(BasePrefixCache): def can_terminate_prefetch(self, operation: PrefetchOperation) -> bool: if self.prefetch_stop_policy == "best_effort": return True - - if len(operation.hash_value) == 0: - completed = False - else: - completed = ( - operation.completed_tokens == len(operation.hash_value) * self.page_size - ) - if self.prefetch_stop_policy == "wait_complete": - can_terminate = completed + return False elif self.prefetch_stop_policy == "timeout": - can_terminate = completed or self._prefetch_timeout_check_linear_func( - operation - ) + return self._prefetch_timeout_check_linear_func(operation) 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( - [1 - int(can_terminate), int(operation_terminated)], - dtype=torch.int, - ) - self._all_reduce_attn_groups(states, torch.distributed.ReduceOp.MAX) - can_terminate = states[0].item() == 0 - operation_terminated = states[1].item() == 1 - return can_terminate or operation_terminated @rank_consensus(same_params=True, same_results=True) def check_prefetch_progress(self, req_id: str) -> bool: if req_id not in self.ongoing_prefetch: return True + _, _, _, operation, _, _ = self.ongoing_prefetch[req_id] + + # Determine whether or not we should terminate this prefetch request. Make all + # ranks agree on the decision. When running with PP, PPn will follow PP0's decision. + should_terminate = False + if self.pp_rank == 0: + should_terminate = operation.is_terminated() or self.can_terminate_prefetch( + operation + ) + should_terminate_tensor = torch.tensor( + int(should_terminate), dtype=torch.int, device="cpu" + ) + self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX) + should_terminate = should_terminate_tensor.item() == 1 + + if not should_terminate: + return False + + self.cache_controller.terminate_prefetch(operation) + if operation.host_indices is None: + self.revoke_pending_prefetch(req_id) + else: + self._handle_prefetch_result(operation) + return True + + def _handle_prefetch_result(self, operation: PrefetchOperation) -> None: + # This function **owns**: + # - host_indices[0 : completed_tokens] + # - sidecar pool hits if operation.pool_transfers_done is true + # + # That is, when this function returns the host memory referenced must be inserted + # into the radix tree or released to pool. + + req_id = operation.request_id + completed_tokens = operation.completed_tokens + hash_value = operation.hash_value + ( last_host_node_id, prefetch_key, host_indices, - operation, + _, anchor_lock_params, comp_xfers, ) = self.ongoing_prefetch[req_id] - 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 - ) - - min_completed_tokens = self._sync_and_check_hybrid_prefetch_result( + # All PP/TP ranks will get the same `min_completed_tokens`, because `completed_tokens` + # and `pool_hits` in their operations are same. No need to sync cross-rank here. + if not self._check_hybrid_prefetch_result( req_id, operation, completed_tokens, @@ -1783,27 +1785,23 @@ class UnifiedRadixCache(BasePrefixCache): last_host_node_id, anchor_lock_params, prefetch_key, - ) - if min_completed_tokens is None: + ): # Hybrid all-or-nothing check failed; result already discarded. - return True + return if self.buffer_pipeline is not None: # No graft: release the rank-local tail beyond the synced usable # length, then park the bounce for admission-time consumption. - self.cache_controller.append_host_mem_release( - host_indices[min_completed_tokens:completed_tokens] - ) return self.buffer_pipeline.stage_completed_prefetch( - req_id, min_completed_tokens, hash_value + req_id, completed_tokens, hash_value ) - fetched_key = prefetch_key[:min_completed_tokens] + fetched_key = prefetch_key[:completed_tokens] insert_result = self.tree_core.insert_host( last_host_node_id, fetched_key, - host_indices[:min_completed_tokens], - hash_value[: min_completed_tokens // self.page_size], + host_indices[:completed_tokens], + hash_value[: completed_tokens // self.page_size], ) # Apply the host-insert walk's actions before the transfer commit. @@ -1815,7 +1813,6 @@ class UnifiedRadixCache(BasePrefixCache): extra_pools=[x for xfers in comp_xfers.values() for x in xfers], ) loaded_from_storage = 0 - released_tokens = completed_tokens else: commit_actions: list[CacheAction | ComponentAction] = [] self.tree_core.commit_hicache_transfers( @@ -1833,11 +1830,7 @@ class UnifiedRadixCache(BasePrefixCache): self.cache_controller.mem_pool_host.free( host_indices[: insert_result.prefix_len] ) - self.cache_controller.append_host_mem_release( - host_indices[min_completed_tokens:completed_tokens] - ) - loaded_from_storage = min_completed_tokens - insert_result.prefix_len - released_tokens = completed_tokens - min_completed_tokens + loaded_from_storage = completed_tokens - insert_result.prefix_len self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[req_id] @@ -1845,21 +1838,19 @@ class UnifiedRadixCache(BasePrefixCache): self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage logger.info( - "HiCache prefetch %s req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d released=%d occupied=%d", + "HiCache prefetch %s req=%s completed=%d matched=%d loaded=%d occupied=%d", "dropped" if insert_result.host_insert_dropped else "success", req_id, completed_tokens, - min_completed_tokens, insert_result.prefix_len, loaded_from_storage, - released_tokens, self.cache_controller.prefetch_tokens_occupied, ) if self.enable_storage_metrics and self.storage_metrics_collector is not None: self.storage_metrics_collector.log_prefetched_tokens(loaded_from_storage) - return True + return - def _sync_and_check_hybrid_prefetch_result( + def _check_hybrid_prefetch_result( self, req_id: str, operation: PrefetchOperation, @@ -1869,8 +1860,8 @@ class UnifiedRadixCache(BasePrefixCache): last_host_node_id: NodeId, anchor_lock_params: DecLockRefParams, prefetch_key: RadixKey, - ) -> Optional[int]: - """Sync prefetch results across ATTN groups and decide the usable prefix. + ) -> bool: + """Decide the length of usable prefix. Two strategies depending on the hybrid layout: @@ -1882,41 +1873,29 @@ class UnifiedRadixCache(BasePrefixCache): *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). + Returns true if prefetch success, or false 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 [] + # + # Skip KV-derived pools, which do not report hits in operation.pool_storage_result. + # Their hit lengths are stored in completed_tokens. + pool_transfers = [ + transfer + for transfer in operation.pool_transfers or [] + if transfer.indices_from_pool != PoolName.KV + ] 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 - + completed_tokens = operation.completed_tokens # 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( + all_succeeded = 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) ) @@ -1925,14 +1904,14 @@ class UnifiedRadixCache(BasePrefixCache): # tail (host_indices[completed_tokens:]) self.cache_controller.append_host_mem_release( host_indices=host_indices[:completed_tokens], - extra_pools=pool_transfers, + extra_pools=pool_transfers if operation.pool_transfers_done else None, ) if anchor_lock_params is not None: self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) - del self.ongoing_prefetch[req_id] if self.buffer_pipeline is not None: self.buffer_pipeline.pop_prefix_ctx(req_id) self.buffer_pipeline.release_anchor_lock(req_id) + del self.ongoing_prefetch[req_id] self.cache_controller.prefetch_tokens_occupied -= ( self._prefetch_occupied_span(prefetch_key, host_indices) ) @@ -1943,14 +1922,8 @@ class UnifiedRadixCache(BasePrefixCache): 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 - operation = self.ongoing_prefetch[req_id].operation - operation.mark_terminate() + return False + return True def pop_prefetch_loaded_tokens(self, req_id: str) -> int: return self.prefetch_loaded_tokens_by_reqid.pop(req_id, 0) @@ -1994,16 +1967,16 @@ class UnifiedRadixCache(BasePrefixCache): return completed_tokens, _ = self.cache_controller.terminate_prefetch(operation) - self._barrier_attn_groups() if anchor_lock_params is not None: self.dec_host_lock_ref(last_host_node_id, anchor_lock_params) del self.ongoing_prefetch[rid] if self.buffer_pipeline is not None: self.buffer_pipeline.pop_prefix_ctx(rid) self.buffer_pipeline.release_anchor_lock(rid) + pool_transfers = [x for xfers in comp_xfers.values() for x in xfers] self.cache_controller.append_host_mem_release( host_indices=host_indices[:completed_tokens], - extra_pools=[x for xfers in comp_xfers.values() for x in xfers], + extra_pools=pool_transfers if operation.pool_transfers_done else None, ) # Buffer mode granted occupancy at hit-alloc, sized to the bounce; # cache mode reserved the requested span at enqueue. @@ -2099,6 +2072,7 @@ class UnifiedRadixCache(BasePrefixCache): def _drain_storage_control_queues_impl( self, n_storage_hit: Optional[int], + n_ack_prefetch: Optional[int], n_backup: Optional[int], n_release: Optional[int], extra_release_counts: Optional[dict[PoolName, int]], @@ -2106,15 +2080,25 @@ class UnifiedRadixCache(BasePrefixCache): ) -> None: cc = self.cache_controller - def _drain_queue(q: Queue[T], limit: Optional[int]) -> Iterator[T]: - drained = 0 - while limit is None or drained < limit: - try: - item = q.get_nowait() - except Empty: - break - drained += 1 - yield item + def _drain_queue(q: Queue[T], n: Optional[int]) -> Iterator[T]: + """If n is None, consume all items from the queue. + Otherwise, consume n items from the queue. Blocking if there are no enough n items. + + In TP, each rank consumes the a minimal number of items of all ranks. + In PP, each rank consumes the exact number of items of PP0. Refer to _pp_sync for more details. + + This prevents TP/PP divergence. + """ + if n is None: + while not q.empty(): + item = q.get() + yield item + else: + for _ in range(n): + # Block when there are not enough elements. + # All TP/PP ranks must consume the same number of elements. + item = q.get() + yield item buffer_mode = self.host_memory_mode == "buffer_only" @@ -2202,6 +2186,33 @@ class UnifiedRadixCache(BasePrefixCache): self._prefetch_outcome_stats["declined_rate_limited"] += 1 self.buffer_pipeline.pending_hit_allocs.append(operation) + def _drain_ack_prefetch(): + for ack in _drain_queue(cc.ack_prefetch_queue, n_ack_prefetch): + operation = ack.operation + if ack.completed_tokens is not None: + if operation.request_id in self.ongoing_prefetch: + assert operation.completed_tokens <= ack.completed_tokens + operation.completed_tokens = ack.completed_tokens + if ack.pool_hits is not None: + if operation.request_id in self.ongoing_prefetch: + operation.pool_storage_result.update_extra_pool_hit_pages( + ack.pool_hits + ) + operation.pool_transfers_done = True + if ack.completed_req: + if operation.request_id in self.ongoing_prefetch: + # check_prefetch_progress() is not called for this rid yet. + # Let us insert the prefetch result into the radix tree. + self._handle_prefetch_result(operation) + cc.append_host_mem_release( + operation.host_indices[operation.completed_tokens :], + ( + operation.pool_transfers + if not operation.pool_transfers_done + else None + ), + ) + def _drain_backup(): drained = 0 for operation in _drain_queue(cc.ack_backup_queue, n_backup): @@ -2255,6 +2266,7 @@ class UnifiedRadixCache(BasePrefixCache): return drained _drain_and_alloc_storage_hit() + _drain_ack_prefetch() _drain_backup() _drain_release() _drain_extra_release() @@ -2265,6 +2277,7 @@ class UnifiedRadixCache(BasePrefixCache): extra_pool_names = list(extra_release_queues) local_qsize_list = [ cc.prefetch_hit_queue.qsize(), + cc.ack_prefetch_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), *[ @@ -2276,15 +2289,16 @@ class UnifiedRadixCache(BasePrefixCache): local_qsize_list, dtype=torch.int, ) - self._all_reduce_attn_groups(qsizes, torch.distributed.ReduceOp.MIN) + self._all_reduce(qsizes, torch.distributed.ReduceOp.MIN) qsize_list = list(map(int, qsizes.tolist())) - n_storage_hit, n_backup, n_release = qsize_list[:3] + n_storage_hit, n_ack_prefetch, 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_storage_hit=n_storage_hit, + n_ack_prefetch=n_ack_prefetch, n_backup=n_backup, n_release=n_release, extra_release_counts=extra_release_counts, @@ -2306,6 +2320,7 @@ class UnifiedRadixCache(BasePrefixCache): return self._drain_storage_control_queues_impl( n_storage_hit=0, + n_ack_prefetch=0, n_backup=None, n_release=None, extra_release_counts={ @@ -2389,6 +2404,7 @@ class UnifiedRadixCache(BasePrefixCache): storage_queue_sizes = ( ( cc.prefetch_hit_queue.qsize(), + cc.ack_prefetch_queue.qsize(), cc.ack_backup_queue.qsize(), cc.host_mem_release_queue.qsize(), *(extra_release_queues[name].qsize() for name in extra_pool_names), @@ -2612,16 +2628,19 @@ class UnifiedRadixCache(BasePrefixCache): self.loading_check(finish_count=load_finish_count) if self.enable_storage and storage_queue_sizes: - n_storage_hit, n_backup, n_release = storage_queue_sizes[:3] + n_storage_hit, n_ack_prefetch, n_backup, n_release = ( + storage_queue_sizes[:4] + ) extra_release_counts = { pool_name: count for pool_name, count in zip( extra_pool_names, - storage_queue_sizes[3:], + storage_queue_sizes[4:], ) } self._drain_storage_control_queues_impl( n_storage_hit=n_storage_hit, + n_ack_prefetch=n_ack_prefetch, n_backup=n_backup, n_release=n_release, extra_release_counts=extra_release_counts, diff --git a/test/manual/hicache/test_pp_with_hicache.py b/test/registered/hicache/test_pp_with_hicache.py similarity index 98% rename from test/manual/hicache/test_pp_with_hicache.py rename to test/registered/hicache/test_pp_with_hicache.py index 9c14d173b..57226b4dc 100644 --- a/test/manual/hicache/test_pp_with_hicache.py +++ b/test/registered/hicache/test_pp_with_hicache.py @@ -3,6 +3,10 @@ Usage: python3 -m unittest test_pp_with_hicache.TestPPWithHiCache.test_eval_accuracy """ +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-c", runner_config="4-gpu-h100") + import os import subprocess import time diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py index bea3fb06f..81249dee1 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_hicache_pp_kl.py @@ -1,7 +1,12 @@ +import os +import shutil +import tempfile import unittest from types import SimpleNamespace from urllib.parse import urlparse +from test_unified_radix_cache_kl_nightly import AccuracyTwoPassMixin + from sglang.srt.utils import kill_process_tree from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.unified_radix_cache_kit import UnifiedRadixTreeTestMixin @@ -102,5 +107,59 @@ class TestUnifiedQwen3HiCachePP(UnifiedRadixTreeTestMixin, CustomTestCase): kill_process_tree(cls.process.pid) +class TestUnifiedQwen3HiCachePPL3(AccuracyTwoPassMixin, CustomTestCase): + """Qwen3-32B + HiCache L3 (file backend) + PP + UnifiedRadixCache.""" + + gsm8k_threshold = 0.8 + + @classmethod + def setUpClass(cls): + cls.model = QWEN3_32B_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.hicache_dir = tempfile.mkdtemp(prefix="hicache_l3_pp_") + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--trust-remote-code", + "--tp-size", + "2", + "--pp-size", + "2", + "--mem-fraction-static", + "0.8", + "--cuda-graph-max-bs", + "32", + "--max-total-tokens", + "14000", + "--disable-piecewise-cuda-graph", + "--model-loader-extra-config", + '{"enable_multithread_load": true, "num_threads": 64}', + "--enable-hierarchical-cache", + "--hicache-write-policy", + "write_through", + "--hicache-storage-prefetch-policy", + "wait_complete", + "--hicache-io-backend", + "direct", + "--hicache-mem-layout", + "page_first_direct", + "--hicache-storage-backend", + "file", + ], + env={ + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR": cls.hicache_dir, + }, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + if os.path.isdir(cls.hicache_dir): + shutil.rmtree(cls.hicache_dir, ignore_errors=True) + + if __name__ == "__main__": unittest.main() 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 ff2510465..67a8a94a0 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 @@ -2627,7 +2627,6 @@ class UnifiedRadixCacheSuite: req_id, cons.root_node.id, 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)))) @@ -2647,6 +2646,282 @@ class UnifiedRadixCacheSuite: self.assertTrue(torch.equal(loaded_v, expected_v)) cons.sanity_check() + def test_release_aborted_request_l3_prefetch_io_in_progress(self): + """Test release_aborted_request while a prefetch IO is still in-progress. + 1. Fill KV and SWA to L3. + 2. Trigger L3 prefetch. + 3. Hack IO thread, blocking at _page_transfer. + 4. Call release_aborted_request. Assert that the KV and SWA buffers are not released (owned by the IO thread). + 5. Unlbock IO thread. Assert that KV and SWA buffers are eventually released. + """ + if self._skip_unsupported_hicache_test(): + return + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("SWA-only fixture required to exercise extra pool") + + # SWA prefetch is all-or-nothing over one full sliding window: size the + # request at sw_pages + 1 pages so prepare_prefetch actually materializes + # an SWA host transfer. + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + num_pages = max(4, sw_pages + 1) + seq = self._make_seq(1, num_pages) + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + # --- Producer tree: D->H backup, H->L3 offload, flush. --- + 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 = prod.resolve_node_handle(mp.last_device_node) + 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, _, _ = build_fixture(self.cfg) + self._init_hicache( + cons, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + ) + req_id = "abort-req" + + cc = cons.cache_controller + kv_pool_available_size_before = cc.mem_pool_host.get_pool( + PoolName.KV + ).available_size() + swa_pool_available_size_before = cc.mem_pool_host.get_pool( + PoolName.SWA + ).available_size() + occupied_before = cons.cache_controller.prefetch_tokens_occupied + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) + self.assertEqual( + cons.cache_controller.prefetch_tokens_occupied, + occupied_before + len(seq), + ) + self.assertIn(req_id, cons.ongoing_prefetch) + + # Block IO thread at the entry of _page_transfer. + import threading + + from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + ) + + original_sidecar = HybridCacheController._page_transfer + entered = threading.Event() + gate = threading.Event() + swa_release_q = cc.extra_host_mem_release_queues.get(PoolName.SWA) + + def _page_transfer_hook(_self, op): + # Signal that the worker has reached the sidecar step, then block + # until the main thread releases the gate. + entered.set() + gate.wait(timeout=10.0) + # Forward to the real implementation so the worker runs the + # terminate-aware release path itself. + original_sidecar(_self, op) + + # Simulate a slow prefetch IO. Hook on _page_transfer. + with mock.patch.object( + HybridCacheController, "_page_transfer", _page_transfer_hook + ): + # Pump until the IO aux thread has entered the sidecar barrier. + deadline = time.time() + 10.0 + while time.time() < deadline: + cons.drain_storage_control_queues() + op = cons.ongoing_prefetch[req_id].operation + if entered.is_set(): + break + time.sleep(0.01) + else: + self.fail("prefetch did not reach the sidecar barrier in time") + + # Now, the prefetch IO thread is stopping at _page_transfer. + # Let the scheduler thread calls release_aborted_request. + # Assert that everything will not be released. + op = cons.ongoing_prefetch[req_id].operation + self.assertFalse(op.pool_transfers_done) + self.assertEqual(cc.host_mem_release_queue.qsize(), 0) + self.assertEqual(swa_release_q.qsize(), 0) + cons.release_aborted_request(req_id) + self.assertEqual(cc.host_mem_release_queue.qsize(), 0) + self.assertEqual(swa_release_q.qsize(), 0) + + # Let the prefetch IO thread continue to run. + gate.set() + + # Wait for the IO thread has completed prefetch. + # We don't consume release queue, as we will use that later to check whether + # the host memory was released yet (as used by the prefetch IO thread). + deadline = time.time() + 10.0 + while time.time() < deadline: + cons._drain_storage_control_queues_impl( + n_storage_hit=0, + n_ack_prefetch=min(1, cc.ack_prefetch_queue.qsize()), + n_backup=0, + n_release=0, + extra_release_counts=None, + log_metrics=True, + ) + if swa_release_q.qsize() > 0: + break + time.sleep(0.01) + else: + self.fail("SWA extra pool was not released after the abort") + + # Asserts that everything is correctly released. + self.assertFalse(op.pool_transfers_done) + self.assertGreater(cc.host_mem_release_queue.qsize(), 0) + self.assertGreater(swa_release_q.qsize(), 0) + self.assertNotIn(req_id, cons.ongoing_prefetch) + self.assertNotIn(req_id, cons.prefetch_loaded_tokens_by_reqid) + self.assertEqual(cc.prefetch_tokens_occupied, occupied_before) + + cons.drain_storage_control_queues() # Drain release queue. + self.assertEqual( + cc.mem_pool_host.get_pool(PoolName.KV).available_size(), + kv_pool_available_size_before, + ) + self.assertEqual( + cc.mem_pool_host.get_pool(PoolName.SWA).available_size(), + swa_pool_available_size_before, + ) + cons.sanity_check() + + def test_release_aborted_request_l3_prefetch_io_done(self): + """Test release_aborted_request is called after the IO thread has completed + prefetch. + 1. Fill KV and SWA to L3. + 2. Trigger L3 prefetch. + 3. Wait until the completion of L3 prefetch IO. + 4. Call release_aborted_request. Assert that KV and SWA buffers are released. + """ + if self._skip_unsupported_hicache_test(): + return + if not self.cfg.has_swa or self.cfg.has_mamba: + self.skipTest("SWA-only fixture required to exercise extra pool") + + # SWA prefetch is all-or-nothing over one full sliding window: size the + # request at sw_pages + 1 pages so prepare_prefetch actually materializes + # an SWA host transfer. + sw_pages = ( + self.cfg.sliding_window_size + self.cfg.page_size - 1 + ) // self.cfg.page_size + num_pages = max(4, sw_pages + 1) + seq = self._make_seq(1, num_pages) + + storage_dir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, storage_dir, ignore_errors=True) + + # --- Producer tree: D->H backup, H->L3 offload, flush. --- + 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 = prod.resolve_node_handle(mp.last_device_node) + 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, _, _ = build_fixture(self.cfg) + self._init_hicache( + cons, + storage_backend="file", + storage_dir=storage_dir, + prefetch_threshold=1, + ) + req_id = "abort-req" + + cc = cons.cache_controller + kv_pool_available_size_before = cc.mem_pool_host.get_pool( + PoolName.KV + ).available_size() + swa_pool_available_size_before = cc.mem_pool_host.get_pool( + PoolName.SWA + ).available_size() + occupied_before = cc.prefetch_tokens_occupied + cons.prefetch_from_storage( + req_id, cons.root_node.id, array("q", seq), None, None + ) + self.assertEqual( + cons.cache_controller.prefetch_tokens_occupied, + occupied_before + len(seq), + ) + op = cons.ongoing_prefetch[req_id].operation + + swa_release_q = cc.extra_host_mem_release_queues.get(PoolName.SWA) + self.assertIsNotNone(swa_release_q) + + # Simulate polling check_hicache_events. + # There will be a sequence of events populated from queue: + # 1. a storage hit notification (from cc.prefetch_hit_queue). + # 2. a HiCacheAck, indicating the copmletion of KV pool read. + # 3. a HiCacheAck, indicating the completion of SWA pool read. + # 4. a HiCacheACk, idnicating the completion of entire prefetch request. + # We are going to stop at the exact timing-window between 3 and 4. So we have to + # consume ONE event from the queue at each iteration. + deadline = time.time() + 10.0 + while time.time() < deadline: + cons._drain_storage_control_queues_impl( + n_storage_hit=min(1, cc.prefetch_hit_queue.qsize()), + n_ack_prefetch=min(1, cc.ack_prefetch_queue.qsize()), + n_backup=0, + n_release=0, + extra_release_counts=None, + log_metrics=True, + ) + if op.pool_transfers_done: + break + time.sleep(0.01) + else: + self.fail("prefetch IO did not complete (pool_transfers_done) in time") + + self.assertIsNotNone(op.host_indices) + self.assertTrue(op.pool_transfers_done) + self.assertEqual(cc.host_mem_release_queue.qsize(), 0) + self.assertEqual(swa_release_q.qsize(), 0) + + # --- Act: abort without committing the prefetch. --- + cons.release_aborted_request(req_id) + + self.assertTrue(op.pool_transfers_done) + self.assertGreater(swa_release_q.qsize(), 0) + self.assertGreater(cc.host_mem_release_queue.qsize(), 0) + self.assertNotIn(req_id, cons.ongoing_prefetch) + self.assertNotIn(req_id, cons.prefetch_loaded_tokens_by_reqid) + self.assertEqual(cc.prefetch_tokens_occupied, occupied_before) + + cons.drain_storage_control_queues() # Drain release queue. + self.assertEqual( + cc.mem_pool_host.get_pool(PoolName.KV).available_size(), + kv_pool_available_size_before, + ) + self.assertEqual( + cc.mem_pool_host.get_pool(PoolName.SWA).available_size(), + swa_pool_available_size_before, + ) + + cons.sanity_check() + # ================================================================ # Buffer-only host memory mode (host = transient staging, L3 = cache) # ================================================================ @@ -3302,35 +3577,28 @@ class UnifiedRadixCacheSuite: # ---------- TP consistency for SWA prefetch (all-or-nothing) ---------- - def _patch_tp_all_reduce(self, cache, drop_swa: bool): - """Fake all_reduce so check_prefetch_progress runs the tp>1 path.""" + def _patch_tp_prefetch_sync(self, cache, drop_swa: bool): + """Fake all_reduce so _reduce_prefetch_ack runs the tp>1 path.""" import torch.distributed as dist - min_sizes = [] + cc = cache.cache_controller - def swa_packed_index(): - # Packed tensor is [completed_tokens, *sidecar_hits]; sidecar order - # matches comp_xfers stored in ongoing_prefetch (one live entry). - for info in cache.ongoing_prefetch.values(): - 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), 1 + len(names) - return None, None + # Fake _reduce_prefetch_ack to drop SWA pool_hits when drop_swa is True. + def fake_reduce(ack): + if drop_swa and ack.pool_hits is not None: + if PoolName.SWA.value in ack.pool_hits: + ack.pool_hits[PoolName.SWA.value] = 0 - def fake(tensor, op=None, group=None): - if op == dist.ReduceOp.MIN: - min_sizes.append(tensor.numel()) - if drop_swa: - idx, packed_numel = swa_packed_index() - if idx is not None and tensor.numel() == packed_numel: - tensor[idx] = 0 - return None + p_reduce = mock.patch.object( + cc, "_reduce_prefetch_ack", side_effect=fake_reduce + ) + p_reduce.start() + self.addCleanup(p_reduce.stop) - p = mock.patch.object(dist, "all_reduce", side_effect=fake) - p.start() - self.addCleanup(p.stop) - return min_sizes + # Make all_reduce no-op. The real all_reduce fails in unit tests with tp_world_size=2. + p_dist = mock.patch.object(dist, "all_reduce", return_value=None) + p_dist.start() + self.addCleanup(p_dist.stop) def _swa_host_on_path(self, cache, seq): m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -3368,7 +3636,6 @@ class UnifiedRadixCacheSuite: req_id, cons.root_node.id, array("q", seq), None, None ) self._run_prefetch_to_completion(cons, req_id) - cons.drain_storage_control_queues() def _setup_swa_tp_prefetch(self): """Skip non-SWA fixtures; produce one full SWA window+1 page to L3. @@ -3407,7 +3674,7 @@ class UnifiedRadixCacheSuite: cons = self._l3_consumer(storage_dir) cons.tp_world_size = 2 - min_sizes = self._patch_tp_all_reduce(cons, drop_swa=True) + self._patch_tp_prefetch_sync(cons, drop_swa=True) self._consume_prefetch(cons, seq, "drop") m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -3415,10 +3682,6 @@ class UnifiedRadixCacheSuite: self.assertFalse( self._swa_host_on_path(cons, seq), "SWA must be dropped when a peer misses" ) - # Full + sidecars must be synced through a packed MIN all_reduce. The - # poll loop may observe more than one completed check, so do not pin the - # exact number of reductions. - self.assertIn(2, min_sizes) cons.sanity_check() def test_tp_swa_prefetch_adopted_when_peer_present(self): @@ -3431,7 +3694,7 @@ class UnifiedRadixCacheSuite: cons = self._l3_consumer(storage_dir) cons.tp_world_size = 2 - min_sizes = self._patch_tp_all_reduce(cons, drop_swa=False) # peer == local + self._patch_tp_prefetch_sync(cons, drop_swa=False) # peer == local self._consume_prefetch(cons, seq, "keep") m = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -3440,7 +3703,6 @@ class UnifiedRadixCacheSuite: self._swa_host_on_path(cons, seq), "SWA must be adopted when all ranks have it", ) - self.assertIn(2, min_sizes) cons.sanity_check() def test_tp_swa_prefetch_drop_frees_host_pool(self): @@ -3453,7 +3715,7 @@ class UnifiedRadixCacheSuite: cons = self._l3_consumer(storage_dir) cons.tp_world_size = 2 - self._patch_tp_all_reduce(cons, drop_swa=True) + self._patch_tp_prefetch_sync(cons, drop_swa=True) avail_before = cons.swa_kv_pool_host.available_size() self._consume_prefetch(cons, seq, "drop") @@ -3464,6 +3726,7 @@ class UnifiedRadixCacheSuite: 0, ) # Whole window dropped -> its host buffer is fully released back. + cons.drain_storage_control_queues() # Drain the release queue. self.assertEqual(cons.swa_kv_pool_host.available_size(), avail_before) def test_hicache_write_back_evict_drops_unbacked_leaf_when_host_full(self): @@ -6965,12 +7228,14 @@ class TestPrefetchCommitOrdering(CustomTestCase): insert_result.prefix_len = 4 insert_result.host_insert_dropped = False cache.tree_core.insert_host.return_value = insert_result + operation = mock.MagicMock() + operation.request_id = "req" cache.ongoing_prefetch = { - "req": ( + operation.request_id: ( 7, list(range(8)), list(range(100, 108)), - mock.MagicMock(), + operation, None, {}, ) @@ -6979,9 +7244,11 @@ class TestPrefetchCommitOrdering(CustomTestCase): 8, [f"h{i}" for i in range(8)], ) - cache._sync_and_check_hybrid_prefetch_result.return_value = 8 + cache._check_hybrid_prefetch_result.return_value = 8 cache.cache_controller.prefetch_tokens_occupied = 100 cache.prefetch_loaded_tokens_by_reqid = {} + cache.can_terminate_prefetch.return_value = True + cache.pp_rank = 0 order = mock.MagicMock() applied = [] @@ -6994,6 +7261,11 @@ class TestPrefetchCommitOrdering(CustomTestCase): cache._apply_cache_actions = order.apply cache.tree_core.commit_hicache_transfers = order.commit + def _handle_prefetch_result(operation): + UnifiedRadixCache._handle_prefetch_result(cache, operation) + + cache._handle_prefetch_result = _handle_prefetch_result + self.assertTrue(UnifiedRadixCache.check_prefetch_progress(cache, "req")) self.assertEqual([c[0] for c in order.mock_calls], ["apply", "commit", "apply"]) @@ -7129,6 +7401,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): operation = mock.Mock() operation.host_indices = host_indices + operation.completed_tokens = completed_tokens operation.pool_storage_result = PoolTransferResult( kv_hit_pages=completed_tokens // self.ps, extra_pool_hit_pages={ @@ -7138,6 +7411,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): ) anchor_lock_params = cache.inc_host_lock_ref(parent_id).to_dec_params() req_id = "drop-all-resources" + operation.request_id = req_id cache.ongoing_prefetch[req_id] = _OngoingPrefetch( parent_id, prefetch_key, @@ -7148,6 +7422,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): ) cache.cache_controller.prefetch_tokens_occupied = completed_tokens hashes = [f"h{i}" for i in range(completed_tokens // self.ps)] + operation.hash_value = hashes with ( mock.patch.object(cache, "can_terminate_prefetch", return_value=True), @@ -7155,7 +7430,7 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): # step: treat the whole fetched prefix as usable so the insert runs. mock.patch.object( cache, - "_sync_and_check_hybrid_prefetch_result", + "_check_hybrid_prefetch_result", return_value=completed_tokens, ), mock.patch.object( @@ -7168,6 +7443,11 @@ class TestUnifiedRadixPrefetchCorruption(CustomTestCase): mock.patch.object( cache.cache_controller, "append_host_mem_release" ) as release, + mock.patch.object( + operation, + "is_terminated", + return_value=False, + ), ): self.assertTrue(cache.check_prefetch_progress(req_id))