diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py index 08489c172..83b2dcd30 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py @@ -1078,6 +1078,10 @@ class TokenizedGenerateReqInput(BaseReq, kw_only=True): # Cache namespace used to isolate otherwise-identical prefixes. cache_salt: Optional[str] = None + # Internal PP control bit, set by PP0 before forwarding the request. + # Keep at the end to preserve the positional Rust wire schema. + pp_prefetch_ticketed: bool = False + def wrap_pickle_fields(self): self.time_stats = wrap_as_pickle(self.time_stats) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 4e1a41e24..f633e64b9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2923,6 +2923,8 @@ class Scheduler( self._add_request_to_queue(req) return + if recv_req.pp_prefetch_ticketed is True: + self.tree_cache.bind_prefetch_ticket(req.rid) self._maybe_namespace_elastic_radix_cache(req) if mm_input_error is not None: @@ -3114,6 +3116,12 @@ class Scheduler( self._add_request_to_queue(req) return + if self.ps.pp_rank == 0 and getattr( + self.tree_cache.cache_controller, "pp_prefetch_command_group", None + ): + recv_req.pp_prefetch_ticketed = bool(self._prefetch_kvcache(req)) + self.tree_cache.bind_prefetch_ticket(req.rid, recv_req.pp_prefetch_ticketed) + added_to_grammar_queue = self.grammar_manager.process_req_with_grammar(req) if not added_to_grammar_queue: self._add_request_to_queue(req) @@ -3168,7 +3176,7 @@ class Scheduler( if tree_cache.hicache_storage_pass_prefix_keys else None ) - tree_cache.prefetch_from_storage( + return tree_cache.prefetch_from_storage( req.cache_request_handle, last_host_node, new_input_tokens, @@ -3268,6 +3276,7 @@ class Scheduler( def _add_request_to_queue(self, req: Req, is_retracted: bool = False): if not self._set_or_validate_priority(req): + self._release_aborted_request(req) return if is_retracted: req.storage_prefetch_retry_attempts = 0 @@ -3275,6 +3284,7 @@ class Scheduler( req.staged_prefetch_plan = None if self.disaggregation_mode == DisaggregationMode.NULL: if self._abort_on_queued_limit(req): + self._release_aborted_request(req) return self._prefetch_kvcache(req) self.waiting_queue.append(req) 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 e6409fe2d..eb19e6e3e 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 @@ -5,7 +5,9 @@ import logging import os import threading import time -from dataclasses import replace +from array import array +from dataclasses import dataclass, field, replace +from enum import Enum, auto from queue import Empty, Queue from typing import TYPE_CHECKING, Any, Callable, List, Optional @@ -36,15 +38,67 @@ from sglang.srt.mem_cache.hicache_storage import ( from sglang.srt.mem_cache.l2_transfer import L2Transfer from sglang.srt.mem_cache.pool_host import HostPoolGroup, PoolEntry from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost +from sglang.srt.mem_cache.radix_cache import RadixKey if TYPE_CHECKING: from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.utils import get_storage_hash_str +from sglang.srt.utils import broadcast_pyobj logger = logging.getLogger(__name__) +@dataclass(frozen=True) +class PPPrefetchPoolSpec: + name: PoolName + num_slots: int + keys: Optional[List[str]] = None + hit_policy: PoolHitPolicy = PoolHitPolicy.ALL_PAGES + indices_from_pool: Optional[PoolName] = None + + @classmethod + def from_transfer(cls, transfer: PoolTransfer) -> PPPrefetchPoolSpec: + return cls( + name=transfer.name, + num_slots=( + int(transfer.host_indices.numel()) + if transfer.host_indices is not None + and transfer.indices_from_pool is None + else 0 + ), + keys=list(transfer.keys) if transfer.keys is not None else None, + hit_policy=transfer.hit_policy, + indices_from_pool=transfer.indices_from_pool, + ) + + +@dataclass +class PPPrefetchTicket: + handle: CacheRequestHandle + prefetch_key: RadixKey + last_hash: Optional[str] + prefix_keys: Optional[List[str]] + matched_prefix_tokens: List[int] + pool_specs: tuple[PPPrefetchPoolSpec, ...] + storage_hit_count: int = 0 + + +class PPPrefetchDecision(Enum): + SKIPPED = auto() # Initial storage miss or skipped prefetch. + TICKETED = auto() # Ticket issued, awaiting consumption. + CANCELLED = auto() # Cancelled before local ticket registration. + + +@dataclass +class PPPrefetchState: + ticket: PPPrefetchTicket + operation: PrefetchOperation + release_requested: bool = False + ready_event: threading.Event = field(default_factory=threading.Event, repr=False) + consumed: bool = False + + class StorageOperation(BaseStorageOperation): def __init__( self, @@ -104,6 +158,12 @@ class PrefetchOperation(StorageOperation): return self._terminated_flag +@dataclass(frozen=True) +class PrefetchSubmission: + operation: Optional[PrefetchOperation] = None + decision: Optional[bool] = None + + class HybridCacheController(BaseHiCacheController): def __init__( self, @@ -127,6 +187,12 @@ class HybridCacheController(BaseHiCacheController): ): startup_storage_backend = storage_backend self.extra_host_mem_release_queues: dict[PoolName, Queue[torch.Tensor]] = {} + self.pp_prefetch_command_group = None + self.pp_prefetch_command_thread = None + self.pp_prefetch_command_queue: Queue[Optional[PPPrefetchTicket]] = Queue() + self.pp_prefetch_state_lock = threading.Lock() + self.pp_prefetch_states: dict[str, PPPrefetchState] = {} + self.pp_prefetch_decisions: dict[str, PPPrefetchDecision] = {} super().__init__( token_to_kv_pool_allocator=token_to_kv_pool_allocator, mem_pool_host=mem_pool_host, @@ -164,6 +230,27 @@ class HybridCacheController(BaseHiCacheController): def _start_storage_threads(self): super()._start_storage_threads() self._init_extra_host_mem_release_queues() + if self.pp_prefetch_command_group is not None: + self.pp_prefetch_command_queue = Queue() + self.pp_prefetch_command_thread = threading.Thread( + target=self.pp_prefetch_command_thread_func, daemon=True + ) + self.pp_prefetch_command_thread.start() + + def _stop_pp_prefetch_thread(self) -> None: + thread = self.pp_prefetch_command_thread + if thread is None: + return + if thread.is_alive() and getattr(self, "pp_rank", 0) == 0: + self.pp_prefetch_command_queue.put(None) + thread.join(timeout=10) + if thread.is_alive(): + raise RuntimeError("Failed to stop PP HiCache ticket thread cleanly.") + self.pp_prefetch_command_thread = None + + def _stop_storage_threads(self): + self._stop_pp_prefetch_thread() + super()._stop_storage_threads() def attach_storage_backend( self, @@ -173,16 +260,47 @@ class HybridCacheController(BaseHiCacheController): storage_backend_extra_config: Optional[dict] = None, host_pools: Optional[list[PoolEntry]] = None, ): - super().attach_storage_backend( - storage_backend=storage_backend, - prefetch_threshold=prefetch_threshold, - model_name=model_name, - storage_backend_extra_config=storage_backend_extra_config, + enable_pp_ticket = ( + self.host_memory_mode == "buffer_only" + and storage_backend == "mooncake" + and self.pp_group is not None + and torch.distributed.get_world_size(group=self.pp_group) > 1 ) + if enable_pp_ticket and self.pp_prefetch_command_group is None: + from sglang.srt.distributed.parallel_state import ( + create_custom_parallel_group, + ) + + self.pp_prefetch_command_group = create_custom_parallel_group( + group_ranks=torch.distributed.get_process_group_ranks(self.pp_group), + backend="gloo", + ) + + try: + super().attach_storage_backend( + storage_backend=storage_backend, + prefetch_threshold=prefetch_threshold, + model_name=model_name, + storage_backend_extra_config=storage_backend_extra_config, + ) + except Exception: + if self.pp_prefetch_command_group is not None: + torch.distributed.destroy_process_group(self.pp_prefetch_command_group) + self.pp_prefetch_command_group = None + raise for entry in host_pools or []: self.storage_backend.register_mem_host_pool_v2(entry.host_pool, entry.name) + def detach_storage_backend(self): + super().detach_storage_backend() + if self.pp_prefetch_command_group is not None: + torch.distributed.destroy_process_group(self.pp_prefetch_command_group) + self.pp_prefetch_command_group = None + with self.pp_prefetch_state_lock: + self.pp_prefetch_states.clear() + self.pp_prefetch_decisions.clear() + def register_host_pool_entry(self, entry: PoolEntry) -> None: if not isinstance(self.mem_pool_host, HostPoolGroup): raise TypeError("Dynamic HiCache sidecars require HostPoolGroup.") @@ -314,7 +432,17 @@ class HybridCacheController(BaseHiCacheController): ) def reset(self): + self._stop_pp_prefetch_thread() super().reset() + with self.pp_prefetch_state_lock: + self.pp_prefetch_states.clear() + self.pp_prefetch_decisions.clear() + if self.enable_storage and self.pp_prefetch_command_group is not None: + self.pp_prefetch_command_queue = Queue() + self.pp_prefetch_command_thread = threading.Thread( + target=self.pp_prefetch_command_thread_func, daemon=True + ) + self.pp_prefetch_command_thread.start() if self.enable_storage: self.host_mem_release_queue.queue.clear() for release_queue in self.extra_host_mem_release_queues.values(): @@ -578,6 +706,315 @@ class HybridCacheController(BaseHiCacheController): self.prefetch_queue.put(operation) return operation + def get_prefetch_submission(self, rid: str) -> Optional[PrefetchSubmission]: + if self.pp_prefetch_command_group is None: + return None + if self.pp_rank != 0: + return PrefetchSubmission() + with self.pp_prefetch_state_lock: + state = self.pp_prefetch_states.get(rid) + if state is not None: + return PrefetchSubmission(decision=not state.consumed) + if rid not in self.pp_prefetch_decisions: + return None + decision = self.pp_prefetch_decisions[rid] + return PrefetchSubmission(decision=decision is PPPrefetchDecision.TICKETED) + + def submit_prefetch( + self, + handle: CacheRequestHandle, + prefetch_key: RadixKey, + last_hash: Optional[str], + prefix_keys: Optional[List[str]], + matched_prefix_tokens: Optional[List[int]], + pool_transfers: Optional[list[PoolTransfer]], + assume_stored: bool = False, + ) -> PrefetchSubmission: + if self.pp_prefetch_command_group is None: + return PrefetchSubmission( + operation=self.prefetch( + handle, + prefetch_key, + last_hash, + prefix_keys, + extra_pools=pool_transfers, + assume_stored=assume_stored, + ) + ) + + if self.pp_rank != 0: + raise RuntimeError("Only PP0 can submit a PP prefetch ticket.") + + rid = handle.rid + ticket = PPPrefetchTicket( + handle=handle, + prefetch_key=RadixKey( + array("q", prefetch_key.token_ids), + extra_key=prefetch_key.extra_key, + cache_salt=prefetch_key.cache_salt, + is_bigram=prefetch_key.is_bigram, + ), + last_hash=last_hash, + prefix_keys=list(prefix_keys) if prefix_keys else None, + matched_prefix_tokens=list(matched_prefix_tokens or []), + pool_specs=tuple( + PPPrefetchPoolSpec.from_transfer(transfer) + for transfer in pool_transfers or [] + ), + ) + operation = PrefetchOperation( + handle, + ticket.prefetch_key, + last_hash, + prefix_keys=ticket.prefix_keys, + pool_transfers=pool_transfers, + ) + + storage_hit_count = len(ticket.prefetch_key.token_ids) + try: + for pp_rank in range(self.tp_rank, self.pp_size, self.tp_size): + _, rank_hit_count = self._storage_hit_query(operation, pp_rank=pp_rank) + storage_hit_count = min(storage_hit_count, rank_hit_count) + except Exception: + logger.exception("PP HiCache hit query failed for req=%s.", rid) + storage_hit_count = 0 + + pp_group_ranks = set(torch.distributed.get_process_group_ranks(self.pp_group)) + local_groups = [ + group + for group in self.prefetch_hits_sync_groups + if set(torch.distributed.get_process_group_ranks(group)) != pp_group_ranks + ] + hit_tensor = torch.tensor(storage_hit_count, dtype=torch.int) + self._all_reduce(hit_tensor, torch.distributed.ReduceOp.MIN, local_groups) + storage_hit_count = int(hit_tensor.item()) + storage_hit_count -= storage_hit_count % self.page_size + if storage_hit_count < self.prefetch_threshold: + with self.pp_prefetch_state_lock: + self.pp_prefetch_decisions[rid] = PPPrefetchDecision.SKIPPED + return PrefetchSubmission(decision=False) + + ticket.storage_hit_count = storage_hit_count + operation.is_pp_broadcast = True + state = PPPrefetchState(ticket=ticket, operation=operation) + with self.pp_prefetch_state_lock: + if rid in self.pp_prefetch_states: + raise RuntimeError(f"Duplicate PP prefetch request id: {rid}") + self.pp_prefetch_states[rid] = state + self.pp_prefetch_decisions[rid] = PPPrefetchDecision.TICKETED + + self.pp_prefetch_command_queue.put(ticket) + return PrefetchSubmission(operation=operation, decision=True) + + def _allocate_pp_prefetch_buffers( + self, ticket: PPPrefetchTicket, operation: PrefetchOperation + ) -> bool: + allocated: list[tuple[PoolName, torch.Tensor]] = [] + + def alloc( + name: PoolName, size: int, num_pages: int = 0 + ) -> Optional[torch.Tensor]: + try: + if size == 0 and num_pages: + size = num_pages * self.mem_pool_host.get_pool(name).page_size + indices = self.mem_pool_host.alloc(size, pool=name) + except Exception: + for pool, owned in reversed(allocated): + self.mem_pool_host.free(owned, pool=pool) + raise + if indices is not None: + allocated.append((name, indices)) + return indices + + host_indices = alloc(PoolName.KV, ticket.storage_hit_count) + if host_indices is None: + return False + + existing = { + transfer.name: transfer for transfer in operation.pool_transfers or [] + } + sources: dict[PoolName, torch.Tensor] = {PoolName.KV: host_indices} + transfers: list[PoolTransfer] = [] + for spec in ticket.pool_specs: + transfer = existing.get(spec.name) + if spec.indices_from_pool is None: + indices = transfer.host_indices if transfer is not None else None + if indices is None: + indices = alloc( + spec.name, + spec.num_slots, + min( + len(spec.keys or []), + ticket.storage_hit_count // self.page_size, + ), + ) + if indices is None: + for name, owned in reversed(allocated): + self.mem_pool_host.free(owned, pool=name) + return False + sources[spec.name] = indices + else: + indices = sources.get(spec.indices_from_pool) + if indices is None: + for name, owned in reversed(allocated): + self.mem_pool_host.free(owned, pool=name) + return False + transfers.append( + PoolTransfer( + name=spec.name, + host_indices=indices, + keys=list(spec.keys) if spec.keys is not None else None, + hit_policy=spec.hit_policy, + indices_from_pool=spec.indices_from_pool, + ) + ) + + operation.host_indices = host_indices + operation.pool_transfers = transfers or None + operation.pool_transfers_done = not bool(transfers) + self.prefetch_tokens_occupied += len(host_indices) + return True + + def _free_pp_prefetch_state(self, state: PPPrefetchState) -> None: + operation = state.operation + if operation.host_indices is not None: + self.mem_pool_host.free(operation.host_indices, pool=PoolName.KV) + self.prefetch_tokens_occupied -= len(operation.host_indices) + operation.host_indices = None + for transfer in operation.pool_transfers or []: + if transfer.indices_from_pool is None and transfer.host_indices is not None: + self.mem_pool_host.free(transfer.host_indices, pool=transfer.name) + transfer.host_indices = None + + def is_pp_prefetch_ready(self, rid: str) -> bool: + with self.pp_prefetch_state_lock: + state = self.pp_prefetch_states.get(rid) + if state is None or state.consumed: + return False + thread = self.pp_prefetch_command_thread + if thread is not None and not thread.is_alive(): + raise RuntimeError(f"PP HiCache ticket thread exited: req={rid}") + return state.ready_event.is_set() + + def take_ready_pp_prefetch(self, rid: str) -> Optional[PPPrefetchState]: + with self.pp_prefetch_state_lock: + state = self.pp_prefetch_states.get(rid) + if state is None or state.consumed: + return None + state.ready_event.wait() + with self.pp_prefetch_state_lock: + if self.pp_prefetch_states.get(rid) is not state or state.consumed: + return None + # Keep the ticket for deduplication until the request finishes. + state.consumed = True + self.pp_prefetch_decisions.pop(rid, None) + if ( + state.operation.host_indices is None + or state.operation.completed_tokens == 0 + ): + self._free_pp_prefetch_state(state) + return state + + def release_pp_prefetch(self, rid: str) -> bool: + """Retire a ticket; free only buffers not handed to buffer mode.""" + with self.pp_prefetch_state_lock: + decision = self.pp_prefetch_decisions.pop(rid, None) + state = self.pp_prefetch_states.get(rid) + if state is None: + if decision not in (None, PPPrefetchDecision.SKIPPED): + # Request rejection may precede local ticket registration. + self.pp_prefetch_decisions[rid] = PPPrefetchDecision.CANCELLED + return True + return False + if state.consumed: + self.pp_prefetch_states.pop(rid) + return False # Buffer mode owns any remaining staging buffers. + state.release_requested = True + if state.ready_event.is_set(): + self._free_pp_prefetch_state(state) + self.pp_prefetch_states.pop(rid, None) + return True + + def pp_prefetch_command_thread_func(self) -> None: + group = self.pp_prefetch_command_group + assert group is not None + rank = torch.distributed.get_rank() + source = torch.distributed.get_process_group_ranks(group)[0] + is_source = self.pp_rank == 0 + + while True: + objects = [] + if is_source: + try: + objects = [self.pp_prefetch_command_queue.get(timeout=60)] + except Empty: + pass # Complete idle broadcasts before the group times out. + operation = None + try: + objects = broadcast_pyobj(objects, rank, group, src=source) + if not objects: + continue + ticket = objects[0] + if ticket is None: + return + + rid = ticket.handle.rid + with self.pp_prefetch_state_lock: + state = self.pp_prefetch_states.get(rid) + if state is None: + operation = PrefetchOperation( + ticket.handle, + ticket.prefetch_key, + ticket.last_hash, + prefix_keys=ticket.prefix_keys, + # Keep sidecar ACKs aligned even if allocation fails. + pool_transfers=[ + PoolTransfer( + name=spec.name, + keys=spec.keys, + hit_policy=spec.hit_policy, + indices_from_pool=spec.indices_from_pool, + ) + for spec in ticket.pool_specs + ] + or None, + ) + operation.is_pp_broadcast = True + state = PPPrefetchState(ticket=ticket, operation=operation) + self.pp_prefetch_states[rid] = state + decision = self.pp_prefetch_decisions.get(rid) + if decision is PPPrefetchDecision.CANCELLED: + state.release_requested = True + self.pp_prefetch_decisions.pop(rid) + operation = state.operation + + operation.hash_value = get_storage_hash_str( + ticket.prefetch_key, + ticket.last_hash, + page_size=self.page_size, + )[: ticket.storage_hit_count // self.page_size] + operation.all_hash_values = list(operation.hash_value) + operation.storage_hit_count = ticket.storage_hit_count + operation.storage_start = len(ticket.matched_prefix_tokens) + operation.pool_storage_result = PoolTransferResult.empty() + if not self._allocate_pp_prefetch_buffers(ticket, operation): + operation.host_indices = torch.empty(0, dtype=torch.int64) + operation.mark_terminate() + self.prefetch_buffer.put(operation) + except Exception: + logger.exception("PP HiCache ticket processing failed.") + if operation is None or not operation.hash_value: + # Without a received ticket/ACK schedule, continuing is unsafe. + return + if operation.host_indices is None: + operation.host_indices = torch.empty(0, dtype=torch.int64) + operation.mark_terminate() + # Preserve every KV/sidecar ACK even when local allocation fails. + self.prefetch_buffer.put(operation) + finally: + if is_source and objects: + self.pp_prefetch_command_queue.task_done() + def write_storage( self, host_indices: torch.Tensor, @@ -596,7 +1033,9 @@ class HybridCacheController(BaseHiCacheController): self.backup_queue.put(operation) return operation.id - def _storage_hit_query(self, operation) -> tuple[list[str], int]: + def _storage_hit_query( + self, operation, pp_rank: Optional[int] = None + ) -> tuple[list[str], int]: hash_value = get_storage_hash_str( operation.token_ids, operation.last_hash, page_size=self.page_size ) @@ -610,7 +1049,8 @@ class HybridCacheController(BaseHiCacheController): return hash_value, kv_hit_pages * self.page_size extra_info = HiCacheStorageExtraInfo( - prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None + prefix_keys=operation.prefix_keys.copy() if operation.prefix_keys else None, + extra_info={"pp_rank": pp_rank} if pp_rank is not None else None, ) if operation.pool_transfers: hit_result = self.storage_backend.batch_exists_v2( @@ -966,3 +1406,41 @@ class HybridCacheController(BaseHiCacheController): ) for i, pool in enumerate(PoolName): ack.pool_hits[pool.value] = packed[i].item() + + def prefetch_sync_thread_func(self) -> None: + """Synchronize progressive ACKs and publish PP ticket completion.""" + 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) + if not getattr(ack.operation, "is_pp_broadcast", False): + self.ack_prefetch_queue.put(ack) + continue + + operation = ack.operation + if ack.completed_tokens is not None: + operation.completed_tokens = ack.completed_tokens + if ack.pool_hits is not None: + operation.pool_storage_result.update_extra_pool_hit_pages( + ack.pool_hits + ) + operation.pool_transfers_done = True + if not ack.completed_req: + continue + + with self.pp_prefetch_state_lock: + state = self.pp_prefetch_states.get(operation.request_id) + if state is None: + continue + operation.hash_value = operation.hash_value[ + : operation.completed_tokens // self.page_size + ] + operation.storage_hit_count = operation.completed_tokens + state.ready_event.set() + if state.release_requested: + self._free_pp_prefetch_state(state) + self.pp_prefetch_states.pop(operation.request_id, None) + except Empty: + continue 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 a9d96cf78..a30e4325b 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 @@ -881,7 +881,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): keys, transfer ) component_keys = self._tag_keys(component_keys) - ex = self._batch_exist(component_keys) + ex = self._batch_exist(component_keys, extra_info) if key_multiplier > 0: page_exists = [ all( @@ -1322,7 +1322,7 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): query_keys.append(f"{key}_{self.mha_suffix}_v") key_multiplier = 2 - exist_result = self._batch_exist(query_keys) + exist_result = self._batch_exist(query_keys, extra_info) for i in range(len(query_keys)): if exist_result[i] != 1: return i // key_multiplier @@ -1374,7 +1374,21 @@ class MooncakeStore(HiCacheStorage, MooncakeBaseStore): ) return self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes) - def _batch_exist(self, key_strs: List[str]) -> List[int]: + def _batch_exist( + self, key_strs: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None + ) -> List[int]: + pp_rank = ( + (extra_info.extra_info or {}).get("pp_rank") + if extra_info is not None + else None + ) + if pp_rank is not None: + # PP is the last rank field before the pool suffix. Replace from + # the right so an identical TP rank or backend tag stays unchanged. + key_strs = [ + f"_{pp_rank}_".join(key.rsplit(f"_{self.pp_rank}_", 1)) + for key in key_strs + ] return self.store.batch_is_exist(key_strs) def get_stats(self): diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index b27536bb0..d2fa662cf 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -41,6 +41,7 @@ from sglang.srt.mem_cache.common import RetractionBackup from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer, SidecarPoolSpec from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( HybridCacheController, + PPPrefetchDecision, PrefetchOperation, ) from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool @@ -961,6 +962,13 @@ class UnifiedRadixCache(BasePrefixCache): def cache_finished_req( self, req: Req, is_insert: bool = True, *, owned_kv_len: int, **kwargs ) -> None: + # Retraction also enters here: retain its ticket until actual finish. + if ( + self.cache_controller is not None + and self.cache_controller.pp_prefetch_command_group is not None + and req.finished() + ): + self.cache_controller.release_pp_prefetch(req.rid) if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs): return @@ -1929,12 +1937,16 @@ class UnifiedRadixCache(BasePrefixCache): extra_key: Optional[str] = None, cache_salt: Optional[str] = None, storage_hit_end: Optional[int] = None, - ) -> None: + ) -> Optional[bool]: if not self.enable_storage or self.cache_controller is None: return req_id = request.rid self.storage_prefetch_retries.cancel(req_id) + submission = self.cache_controller.get_prefetch_submission(req_id) + if submission is not None: + return submission.decision + buffer_mode = self.host_memory_mode == "buffer_only" if request in self.ongoing_prefetch or ( self.buffer_pipeline is not None @@ -2029,21 +2041,31 @@ class UnifiedRadixCache(BasePrefixCache): aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) - operation = self.cache_controller.prefetch( + submission = self.cache_controller.submit_prefetch( request, prefetch_key, last_hash, prefix_keys, - extra_pools=aux_xfers or None, + matched_prefix_tokens, + aux_xfers or None, assume_stored=assume_stored, ) + operation = submission.operation + if operation is None: + assert submission.decision is not None + self.cache_controller.append_host_mem_release(extra_pools=aux_xfers or None) + return submission.decision + stats["issued"] += 1 - if assume_stored: + if operation.assume_stored: stats["issued_assumed_stored"] += 1 # Snapshot the requested span for L3 miss-token accounting at the # rank-synchronized query outcome. operation.stats_requested_tokens = prefetch_length operation.storage_start = storage_start + if submission.decision is not None: + return submission.decision + self.ongoing_prefetch[request] = _OngoingPrefetch( last_host_node_id, prefetch_key, @@ -2097,8 +2119,70 @@ class UnifiedRadixCache(BasePrefixCache): def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool: return handle in self.ongoing_prefetch + def bind_prefetch_ticket(self, req_id: str, decision: bool = True) -> None: + # PP0 also records misses/skips; only hits travel with the request relay. + with self.cache_controller.pp_prefetch_state_lock: + self.cache_controller.pp_prefetch_decisions[req_id] = ( + PPPrefetchDecision.TICKETED if decision else PPPrefetchDecision.SKIPPED + ) + + def _check_pp_prefetch_progress(self, request: CacheRequestHandle) -> bool: + req_id = request.rid + ready = self.pp_rank == 0 and self.cache_controller.is_pp_prefetch_ready(req_id) + ready_tensor = torch.tensor(int(ready), dtype=torch.int, device="cpu") + self._all_reduce(ready_tensor, torch.distributed.ReduceOp.MAX) + if ready_tensor.item() == 0: + return False + + state = self.cache_controller.take_ready_pp_prefetch(req_id) + if state is None: + raise RuntimeError( + f"PP prefetch became ready before local state existed: {req_id}" + ) + + operation = state.operation + if operation.host_indices is None or operation.completed_tokens == 0: + self.prefetch_loaded_tokens_by_reqid[request] = 0 + return True + + ticket = state.ticket + prefetch_key = ticket.prefetch_key + self.ongoing_prefetch[request] = _OngoingPrefetch( + self.root_node_handle(prefetch_key.extra_key), + prefetch_key, + operation.host_indices, + operation, + None, + { + BASE_COMPONENT_TYPE: [ + transfer + for transfer in operation.pool_transfers or [] + if transfer.indices_from_pool is None + ] + }, + ) + self.buffer_pipeline.set_prefix_ctx( + request, + ticket.matched_prefix_tokens, + extra_key=prefetch_key.extra_key, + cache_salt=prefetch_key.cache_salt, + ) + self.buffer_pipeline.try_lock_anchor(request, operation.completed_tokens) + self.cache_controller.append_host_mem_release( + operation.host_indices[operation.completed_tokens :] + ) + self._handle_prefetch_result(operation) + return True + @rank_consensus(same_params=True, same_results=True) def check_prefetch_progress(self, request: CacheRequestHandle) -> bool: + if ( + self.cache_controller is not None + and self.cache_controller.pp_prefetch_decisions.get(request.rid) + is PPPrefetchDecision.TICKETED + ): + return self._check_pp_prefetch_progress(request) + if request not in self.ongoing_prefetch: return True @@ -2462,6 +2546,12 @@ class UnifiedRadixCache(BasePrefixCache): self.prefetch_loaded_tokens_by_reqid.pop(request, None) self.prefetch_loaded_storage_start_by_reqid.pop(request, None) self.storage_prefetch_retries.cancel(rid) + if ( + self.buffer_pipeline is not None + and self.cache_controller.pp_prefetch_command_group is not None + and self.cache_controller.release_pp_prefetch(rid) + ): + return if ( self.buffer_pipeline is not None and self.buffer_pipeline.release_staged_hold(request) @@ -3079,7 +3169,7 @@ class UnifiedRadixCache(BasePrefixCache): cc = self.cache_controller extra_release_queues = getattr(cc, "extra_host_mem_release_queues", {}) extra_pool_names = tuple(extra_release_queues) if self.enable_storage else () - if cc is None or self.pp_rank > 0: + if cc is None or (self.pp_rank > 0 and self.host_memory_mode != "buffer_only"): write_acks = 0 load_acks = 0 # Zero placeholders shaped like PP0's slots: _pp_sync hands the @@ -3117,7 +3207,10 @@ class UnifiedRadixCache(BasePrefixCache): dtype=torch.int64, device="cpu", ) - self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN) + if self.host_memory_mode == "buffer_only" and self.pp_size > 1: + self._all_reduce_attn_groups(ready_counts, torch.distributed.ReduceOp.MIN) + else: + self._all_reduce(ready_counts, torch.distributed.ReduceOp.MIN) count_values = list(map(int, ready_counts.tolist())) assert digest == count_values[-2] and digest == -count_values[-1], ( diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index 3c533ff86..d8a3e4cbc 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -20,6 +20,7 @@ from sglang.srt.managers.schedule_batch import ( # noqa: E402 ReqKvInfo, ) from sglang.srt.managers.scheduler import Scheduler # noqa: E402 +from sglang.srt.mem_cache.base_prefix_cache import CacheRequestOutcome from sglang.srt.runtime_context import get_context, publish, reset_context # noqa: E402 from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci @@ -43,8 +44,13 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase): scheduler.enable_priority_scheduling = True scheduler.schedule_low_priority_values_first = False scheduler.abort_on_priority_when_disabled = False + scheduler.enable_hierarchical_cache = False + scheduler.enable_hicache_storage = False + scheduler.enable_unified_cache_external_linker = False + scheduler.processed_tokens_counter = 0 scheduler.waiting_queue = [] scheduler._prefetch_kvcache = MagicMock() + scheduler.tree_cache = MagicMock() scheduler._abort_on_queued_limit = MagicMock(return_value=False) scheduler.model_config = SimpleNamespace(num_key_value_heads=8) scheduler.disagg_prefill_bootstrap_queue = MagicMock() @@ -100,6 +106,61 @@ class TestDisaggregationPriorityQueueing(unittest.TestCase): scheduler.ipc_channels.send_to_tokenizer.send_output.assert_called_once() req.time_stats.trace_ctx.abort.assert_called_once() + def test_priority_rejection_releases_early_prefetch(self): + for mode in DisaggregationMode: + with self.subTest(mode=mode): + scheduler = self._new_scheduler(mode) + scheduler.enable_priority_scheduling = False + scheduler.abort_on_priority_when_disabled = True + scheduler.enable_hicache_storage = True + scheduler.tree_cache = MagicMock(spec=["finish"]) + req = self._new_req(priority=10) + + with patch( + "sglang.srt.managers.scheduler.get_serving", + return_value=SimpleNamespace(weight_version="v0"), + ): + scheduler._add_request_to_queue(req) + + scheduler.tree_cache.finish.assert_called_once_with( + req.cache_request_handle, CacheRequestOutcome.ABORT + ) + scheduler._prefetch_kvcache.assert_not_called() + self.assertEqual(scheduler.waiting_queue, []) + scheduler.disagg_prefill_bootstrap_queue.add.assert_not_called() + scheduler.disagg_decode_prealloc_queue.add.assert_not_called() + + def test_full_queue_releases_only_the_rejected_request(self): + for priority in (0, 2): + with self.subTest(incoming_priority=priority): + scheduler = self._new_scheduler(DisaggregationMode.NULL) + del scheduler._abort_on_queued_limit # Exercise the real queue check. + scheduler.max_queued_requests = 1 + scheduler.enable_hicache_storage = True + scheduler.tree_cache = MagicMock(spec=["finish"]) + scheduler.beam_coordinator = MagicMock() + queued = self._new_req(priority=1) + queued.rid = "queued" + scheduler.waiting_queue = [queued] + incoming = self._new_req(priority=priority) + + with patch( + "sglang.srt.managers.scheduler.get_serving", + return_value=SimpleNamespace(weight_version="v0"), + ): + scheduler._add_request_to_queue(incoming) + + rejected = incoming if priority == 0 else queued + scheduler.tree_cache.finish.assert_called_once_with( + rejected.cache_request_handle, CacheRequestOutcome.ABORT + ) + if priority == 0: + self.assertEqual(scheduler.waiting_queue, [queued]) + scheduler._prefetch_kvcache.assert_not_called() + else: + self.assertEqual(scheduler.waiting_queue, [incoming]) + scheduler._prefetch_kvcache.assert_called_once_with(incoming) + class TestDecodePreallocQueuePriority(unittest.TestCase): def setUp(self): diff --git a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py b/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py index c3600af32..05bf95168 100644 --- a/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py +++ b/test/registered/unit/mem_cache/test_hiradix_pp_sync_drain.py @@ -1,6 +1,7 @@ """Unit tests for HiCache PP synchronization.""" import unittest +from queue import Queue from types import SimpleNamespace from unittest.mock import MagicMock @@ -59,6 +60,7 @@ class TestUnifiedPPSyncBatching(unittest.TestCase): ) cache.pp_rank = pp_rank cache.pp_size = 2 + cache.host_memory_mode = "cache" cache.enable_storage_metrics = False cache.storage_metrics_collector = None cache.buffer_pipeline = None @@ -110,6 +112,42 @@ class TestUnifiedPPSyncBatching(unittest.TestCase): follower.writing_check.assert_called_once_with(finish_count=1) follower.loading_check.assert_called_once_with(finish_count=1) + def test_buffer_mode_follower_drains_rank_local_completions(self): + cache = self._make_cache(1, [True, False], [True, True]) + cache.host_memory_mode = "buffer_only" + cache.tree_core.enable_storage = True + cache._all_reduce_attn_groups = MagicMock() + cache._drain_storage_control_queues_impl = MagicMock() + cc = cache.cache_controller + for size, name in enumerate( + ( + "prefetch_hit_queue", + "ack_prefetch_queue", + "ack_backup_queue", + "host_mem_release_queue", + ), + start=1, + ): + queue = Queue() + for _ in range(size): + queue.put(object()) + setattr(cc, name, queue) + + cache.check_hicache_events() + + cache._all_reduce.assert_not_called() + cache._all_reduce_attn_groups.assert_called_once() + cache.writing_check.assert_called_once_with(finish_count=1) + cache.loading_check.assert_called_once_with(finish_count=2) + cache._drain_storage_control_queues_impl.assert_called_once_with( + n_storage_hit=1, + n_ack_prefetch=2, + n_backup=3, + n_release=4, + extra_release_counts={}, + log_metrics=True, + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py b/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py new file mode 100644 index 000000000..6802af04e --- /dev/null +++ b/test/registered/unit/mem_cache/test_pp_prefetch_ticket.py @@ -0,0 +1,414 @@ +"""CPU regressions for PP ticket ordering, admission, and buffer ownership.""" + +import pickle +import threading +import unittest +from array import array +from queue import Empty, Queue +from unittest.mock import Mock, call, patch + +import torch + +from sglang.srt.managers.cache_controller import PrefetchAck +from sglang.srt.mem_cache.base_prefix_cache import ( + CacheRequestHandle, + CacheRequestOutcome, +) +from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline +from sglang.srt.mem_cache.hicache_storage import PoolName, PoolTransfer +from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + HybridCacheController, + PPPrefetchDecision, +) +from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache +from sglang.srt.mem_cache.utils import get_storage_hash_str +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestPPPrefetchTicket(unittest.TestCase): + def setUp(self): + self.c = c = HybridCacheController.__new__(HybridCacheController) + c.page_size = c.prefetch_threshold = 4 + c.pp_rank = c.tp_rank = 0 + c.pp_size, c.tp_size = 4, 2 + c.pp_group, c.pp_prefetch_command_group = "pp", "command" + c.pp_prefetch_command_thread = None + c.prefetch_hits_sync_groups = c.prefetch_completion_sync_groups = ["pp", "tp"] + c.pp_prefetch_states, c.pp_prefetch_decisions = {}, {} + c.pp_prefetch_state_lock = threading.Lock() + c.pp_prefetch_command_queue = Queue() + c.prefetch_queue = Queue() + c.prefetch_sync_queue = Queue() + c.ack_prefetch_queue = Queue() + c.host_mem_release_queue = Queue() + c.prefetch_buffer = Queue() + c.storage_stop_event = threading.Event() + c.prefetch_tokens_occupied = 0 + c.mem_pool_host = Mock(page_size=4) + c.mem_pool_host.alloc.side_effect = lambda size, **_: torch.arange(size) + c._storage_hit_query = Mock(return_value=(["h0", "h1"], 8)) + c._all_reduce = Mock() + ranks = {"pp": [0, 2, 4, 6], "tp": [0, 1], "command": [0, 2, 4, 6]} + patcher = patch.object( + torch.distributed, "get_process_group_ranks", side_effect=ranks.__getitem__ + ) + patcher.start() + self.addCleanup(patcher.stop) + patcher = patch.object( + torch.distributed, "get_rank", side_effect=lambda: c.pp_rank * 2 + c.tp_rank + ) + patcher.start() + self.addCleanup(patcher.stop) + self.cache = cache = UnifiedRadixCache.__new__(UnifiedRadixCache) + cache.pp_rank = 1 + cache.tree_core = Mock(enable_storage=True) + cache.cache_controller = c + cache._all_reduce = Mock() + cache.ongoing_prefetch = {} + cache.prefetch_loaded_tokens_by_reqid = {} + cache.prefetch_loaded_storage_start_by_reqid = {} + cache.storage_prefetch_retries = StoragePrefetchRetries() + cache.linker = None + cache.root_node_handle = Mock(return_value=0) + cache.buffer_pipeline = Mock(spec=BufferModePipeline) + cache._handle_prefetch_result = Mock() + + def submit(self, rid="hit", pools=None, attempt_id=0, assume_stored=False): + key = RadixKey( + array("q", range(9)), "adapter", is_bigram=True, cache_salt="tenant" + ) + return self.c.submit_prefetch( + CacheRequestHandle(rid, attempt_id), + key, + "ab" * 32, + ["prefix"], + [10, 11, 12, 13], + pools, + assume_stored=assume_stored, + ) + + def run_commands(self, *tickets): + commands = iter((*tickets, None)) + self.c.pp_prefetch_command_queue.put(None) + + def broadcast(objects, *args, **kwargs): + return [next(commands)] if self.c.pp_rank else objects + + with patch(f"{HybridCacheController.__module__}.broadcast_pyobj", broadcast): + self.c.pp_prefetch_command_thread_func() + + def sync_acks(self, *acks): + for ack in acks: + self.c.prefetch_sync_queue.put(ack) + with patch.object( + self.c.storage_stop_event, + "is_set", + side_effect=[False] * len(acks) + [True], + ): + self.c.prefetch_sync_thread_func() + + def test_hit_and_miss_are_reused_across_enqueue_retract_and_retry(self): + c = self.c + for hit in (0, 8): + with self.subTest(hit=hit): + rid = str(hit) + handle = CacheRequestHandle(rid, 0) + c._storage_hit_query.return_value = ([], hit) + self.assertEqual(self.submit(rid).decision, bool(hit)) + queries = c._storage_hit_query.call_count + for _ in range(2): + self.assertEqual( + self.cache.prefetch_from_storage(handle, 0, []), bool(hit) + ) + self.assertEqual(c._storage_hit_query.call_count, queries) + if hit: + state = c.pp_prefetch_states[rid] + state.ready_event.set() + c.take_ready_pp_prefetch(rid) + self.assertFalse( + self.cache.prefetch_from_storage( + CacheRequestHandle(rid, 1), 0, [] + ) + ) + self.assertTrue(self.cache.check_prefetch_progress(handle)) + self.assertFalse(c.release_pp_prefetch(rid)) + self.assertIsNone(c.get_prefetch_submission(rid)) + self.assertEqual(c.pp_prefetch_command_queue.qsize(), 1) + self.cache._all_reduce.assert_not_called() + + def test_submission_does_not_wait_for_worker_and_uses_only_tp_consensus(self): + c = self.c + with patch.object( + c.pp_prefetch_command_queue, + "join", + side_effect=AssertionError("scheduler blocked"), + ): + self.assertTrue(self.submit(assume_stored=True).decision) + self.assertTrue(self.submit("next").decision) + self.assertEqual(c.pp_prefetch_command_queue.qsize(), 2) + self.assertFalse(c.is_pp_prefetch_ready("hit")) + c.mem_pool_host.alloc.assert_not_called() + self.assertFalse(c.pp_prefetch_states["hit"].operation.assume_stored) + self.assertEqual( + c._all_reduce.call_args.args[1:], (torch.distributed.ReduceOp.MIN, ["tp"]) + ) + self.assertEqual( + [call.kwargs["pp_rank"] for call in c._storage_hit_query.call_args_list], + [0, 2, 0, 2], + ) + + def test_tp_only_preserves_normal_prefetch(self): + c = self.c + c.pp_prefetch_command_group = None + result = self.submit(attempt_id=3, assume_stored=True) + self.assertIsNone(result.decision) + self.assertEqual(result.operation.handle, CacheRequestHandle("hit", 3)) + self.assertTrue(result.operation.assume_stored) + self.assertIs(c.prefetch_queue.get_nowait(), result.operation) + c._storage_hit_query.assert_not_called() + + def test_downstream_request_and_ticket_order_preserves_pp0_admission(self): + c, cache = self.c, self.cache + self.submit(attempt_id=3) + handle = CacheRequestHandle("hit", 3) + ticket = pickle.loads(pickle.dumps(c.pp_prefetch_states["hit"].ticket)) + ticket.last_hash = None + c.pp_rank = 1 + c.pp_prefetch_states.clear() + cache.bind_prefetch_ticket("hit") + self.assertFalse(cache.check_prefetch_progress(handle)) + c._storage_hit_query.reset_mock() + self.run_commands(ticket) + operation = c.prefetch_buffer.get_nowait() + self.assertEqual(operation.handle, handle) + self.assertEqual( + operation.hash_value, get_storage_hash_str(ticket.prefetch_key, page_size=4) + ) + self.assertTrue(operation.token_ids.is_bigram) + self.assertEqual(len(operation.token_ids), 8) + c._storage_hit_query.assert_not_called() + self.sync_acks(PrefetchAck("hit", operation, completed_tokens=8)) + self.assertFalse(c.is_pp_prefetch_ready("hit")) + self.sync_acks(PrefetchAck("hit", operation, completed_req=True)) + self.assertFalse(cache.check_prefetch_progress(handle)) # PP0 has not admitted. + cache._all_reduce.side_effect = lambda tensor, _: tensor.fill_(1) + self.assertTrue(cache.check_prefetch_progress(handle)) + cache._handle_prefetch_result.assert_called_once_with(operation) + cache.buffer_pipeline.try_lock_anchor.assert_called_once_with(handle, 8) + key = cache.ongoing_prefetch[handle].prefetch_key + self.assertEqual((key.extra_key, key.cache_salt), ("adapter", "tenant")) + self.assertTrue(key.is_bigram) + self.assertFalse(c.is_pp_prefetch_ready("hit")) + + def test_allocation_error_preserves_ack_sequence_and_next_ticket(self): + c = self.c + self.submit( + pools=[ + PoolTransfer(PoolName.SWA, host_indices=torch.arange(4), keys=["h1"]) + ] + ) + ticket = c.pp_prefetch_states.pop("hit").ticket + following = pickle.loads(pickle.dumps(ticket)) + following.handle = CacheRequestHandle("next", 0) + c.pp_rank = 1 + kv = torch.arange(8) + c.mem_pool_host.alloc.side_effect = [ + kv, + KeyError(PoolName.SWA), + torch.arange(8), + torch.arange(4), + ] + with self.assertLogs(level="ERROR"): + self.run_commands(ticket, following) + failed = c.pp_prefetch_states["hit"].operation + self.assertTrue(failed.is_terminated()) + c.mem_pool_host.free.assert_called_once_with(kv, pool=PoolName.KV) + c.page_get_func = Mock(return_value=1) + c.storage_backend = Mock() + c.storage_backend.batch_get_v2.return_value = {"swa": [True]} + with ( + patch("sglang.srt.managers.cache_controller.STORAGE_BATCH_SIZE", 1), + patch.object( + c.storage_stop_event, "is_set", side_effect=[False, False, True] + ), + ): + c.prefetch_io_aux_func() + acks = [c.prefetch_sync_queue.get_nowait() for _ in range(8)] + self.assertEqual([a.rid for a in acks], ["hit"] * 4 + ["next"] * 4) + self.assertEqual( + [a.completed_tokens for a in acks], [0, 0, None, None, 4, 8, None, None] + ) + self.assertEqual((acks[2].pool_hits, acks[6].pool_hits), ({}, {"swa": 1})) + self.assertEqual( + c.page_get_func.call_count, 2 + ) # No reads for the failed ticket. + self.sync_acks(*acks) + self.assertTrue(c.is_pp_prefetch_ready("hit")) + self.assertEqual(c.pp_prefetch_states["next"].operation.completed_tokens, 8) + self.assertEqual(c.prefetch_tokens_occupied, 8) + + def test_cancel_before_ticket_defers_free_until_final_ack(self): + c, cache = self.c, self.cache + self.submit(pools=[PoolTransfer(PoolName.SWA, host_indices=torch.arange(4))]) + ticket = c.pp_prefetch_states.pop("hit").ticket + c.pp_rank = 1 + cache.bind_prefetch_ticket("hit") + cache.finish(ticket.handle, CacheRequestOutcome.ABORT) + cache.finish(ticket.handle, CacheRequestOutcome.ABORT) + self.assertIs(c.pp_prefetch_decisions["hit"], PPPrefetchDecision.CANCELLED) + self.run_commands(ticket) + operation = c.prefetch_buffer.get_nowait() + self.sync_acks(PrefetchAck("hit", operation, completed_tokens=4)) + c.mem_pool_host.free.assert_not_called() + self.sync_acks(PrefetchAck("hit", operation, completed_req=True)) + self.assertEqual( + [call.kwargs["pool"] for call in c.mem_pool_host.free.call_args_list], + [PoolName.KV, PoolName.SWA], + ) + self.assertEqual(c.prefetch_tokens_occupied, 0) + self.assertEqual(c.pp_prefetch_states, {}) + self.assertEqual(c.pp_prefetch_decisions, {}) + + def test_lazy_sidecars_use_hit_pages_and_pool_page_size(self): + c = self.c + c.mem_pool_host.get_pool.side_effect = lambda name: Mock( + page_size=1 if name == PoolName.MAMBA else 4 + ) + self.submit( + pools=[ + PoolTransfer(PoolName.SWA, keys=["pending"] * 3), + PoolTransfer(PoolName.MAMBA, keys=["pending"]), + PoolTransfer(PoolName.DRAFT_SWA, indices_from_pool=PoolName.SWA), + ] + ) + ticket = pickle.loads(pickle.dumps(c.pp_prefetch_states["hit"].ticket)) + for rank in (0, 1): + with self.subTest(rank=rank): + c.pp_rank = rank + if rank: + c.pp_prefetch_states.clear() + c.mem_pool_host.alloc.reset_mock() + self.run_commands(ticket) + operation = c.prefetch_buffer.get_nowait() + self.assertEqual( + c.mem_pool_host.alloc.call_args_list, + [ + call(8, pool=PoolName.KV), + call(8, pool=PoolName.SWA), + call(1, pool=PoolName.MAMBA), + ], + ) + swa, mamba, draft_swa = operation.pool_transfers + self.assertIs(draft_swa.host_indices, swa.host_indices) + self.assertEqual(mamba.host_indices.numel(), 1) + + # Missing pool metadata also rolls back KV, before a failed-ticket ACK. + c.pp_prefetch_states.clear() + c.mem_pool_host.get_pool.side_effect = KeyError(PoolName.SWA) + with self.assertLogs(level="ERROR"): + self.run_commands(ticket) + self.assertTrue(c.prefetch_buffer.get_nowait().is_terminated()) + self.assertEqual(c.mem_pool_host.free.call_count, 1) + self.assertEqual(c.mem_pool_host.free.call_args.kwargs["pool"], PoolName.KV) + + def test_failed_source_allocation_does_not_free_borrowed_sidecar_early(self): + c = self.c + borrowed = PoolTransfer(PoolName.SWA, host_indices=torch.arange(4)) + operation = self.submit(pools=[borrowed]).operation + c.mem_pool_host.alloc.side_effect = lambda *args, **kwargs: None + self.run_commands() + self.assertTrue(operation.is_terminated()) + c.mem_pool_host.free.assert_not_called() + self.sync_acks( + PrefetchAck("hit", operation, completed_tokens=0, completed_req=True) + ) + c.take_ready_pp_prefetch("hit") + self.assertIsNone(borrowed.host_indices) + self.assertEqual(c.mem_pool_host.free.call_args.kwargs["pool"], PoolName.SWA) + + def test_command_failure_finishes_queue_task_and_is_reported_on_poll(self): + c = self.c + with ( + patch( + f"{HybridCacheController.__module__}.broadcast_pyobj", + side_effect=RuntimeError("broken"), + ), + self.assertLogs(level="ERROR"), + ): + worker = threading.Thread( + target=c.pp_prefetch_command_thread_func, daemon=True + ) + c.pp_prefetch_command_thread = worker + worker.start() + self.submit() + worker.join(timeout=1) + self.assertFalse(worker.is_alive()) + self.assertEqual(c.pp_prefetch_command_queue.unfinished_tasks, 0) + for rid in ("hit", "next"): + if rid == "next": + self.submit(rid) + with self.assertRaisesRegex(RuntimeError, "ticket thread exited"): + c.is_pp_prefetch_ready(rid) + + def test_idle_source_broadcasts_empty_then_processes_ticket_and_stop(self): + c = self.c + c.tp_rank = 1 # The source is a global rank, not pp_rank=0. + operation = self.submit().operation + c.pp_prefetch_command_queue.put(None) + get = c.pp_prefetch_command_queue.get + idle_polls = [True, True] + + def get_after_idle(*, timeout): + self.assertEqual(timeout, 60) + if idle_polls: + idle_polls.pop() + raise Empty + return get(block=False) + + with ( + patch.object(c.pp_prefetch_command_queue, "get", get_after_idle), + patch.object( + torch.distributed, "get_process_group_ranks", return_value=[1, 3, 5, 7] + ), + patch( + f"{HybridCacheController.__module__}.broadcast_pyobj", + side_effect=lambda objects, *args, **kwargs: objects, + ) as broadcast, + ): + c.pp_prefetch_command_thread_func() + calls = broadcast.call_args_list + self.assertEqual([len(c.args[0]) for c in calls], [0, 0, 1, 1]) + for invocation in calls: + self.assertEqual(invocation.args[1:], (1, "command")) + self.assertEqual(invocation.kwargs, {"src": 1}) + self.assertEqual(c.pp_prefetch_command_queue.unfinished_tasks, 0) + self.assertTrue(c.pp_prefetch_command_queue.empty()) + self.assertIs(c.prefetch_buffer.get_nowait(), operation) + self.assertTrue(c.prefetch_buffer.empty()) + c.mem_pool_host.alloc.assert_called_once() + + def test_idle_downstream_skips_empty_broadcasts_before_ticket_and_stop(self): + c = self.c + self.submit() + ticket = c.pp_prefetch_states.pop("hit").ticket + c.pp_rank = 1 + with ( + patch( + f"{HybridCacheController.__module__}.broadcast_pyobj", + side_effect=[[], [], [ticket], [None]], + ), + patch.object(c.pp_prefetch_command_queue, "task_done") as task_done, + ): + c.pp_prefetch_command_thread_func() + self.assertEqual(c.prefetch_buffer.get_nowait().request_id, "hit") + self.assertTrue(c.prefetch_buffer.empty()) + c.mem_pool_host.alloc.assert_called_once() + task_done.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py index eb58879c9..29331c03f 100644 --- a/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py +++ b/test/registered/unit/mem_cache/test_storage_prefetch_lifecycle.py @@ -71,6 +71,7 @@ def _staged_fixture(full_match=2): full_available_size=Mock(return_value=100) ) cc = HybridCacheController.__new__(HybridCacheController) + cc.pp_prefetch_command_group = None cc.page_size = 2 cc.get_hash_str = get_hash_str cc.prefetch_queue = Queue() 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 c755d3bdd..50ca52f1a 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 @@ -861,6 +861,9 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): ) def test_l3_prefetch_uses_bigram_radix_key(self): + from sglang.srt.mem_cache.hybrid_cache.hybrid_cache_controller import ( + PrefetchSubmission, + ) from sglang.srt.mem_cache.utils import get_hash_str cache, allocator, _ = build_fixture(self.cfg) @@ -891,23 +894,30 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): def prefetch_rate_limited(self): return False - def prefetch( + def get_prefetch_submission(self, rid): + return None + + def submit_prefetch( self, - request_id, - new_input_tokens, - last_hash=None, - prefix_keys=None, - extra_pools=None, + handle, + prefetch_key, + last_hash, + prefix_keys, + matched_prefix_tokens, + pool_transfers, assume_stored=False, ): self.prefetch_args = ( - request_id, - new_input_tokens, + handle, + prefetch_key, last_hash, prefix_keys, - extra_pools, + matched_prefix_tokens, + pool_transfers, + ) + return PrefetchSubmission( + operation=mock.Mock(assume_stored=assume_stored) ) - return mock.Mock() controller = FakeCacheController() cache.cache_controller = controller @@ -915,7 +925,7 @@ class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase): CacheRequestHandle("req", 0), cache.root_node_handle(), 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) @@ -9664,6 +9674,7 @@ class TestPrefetchCommitOrdering(CustomTestCase): cache.page_size = 1 cache.enable_storage_metrics = False cache.buffer_pipeline = None # cache-mode commit path + cache.cache_controller.pp_prefetch_decisions = {} walk_action = object() insert_result = mock.MagicMock() insert_result.cache_actions = [walk_action]