[PP + HiCache] Add PP Prefetch Tickets for eager cross-stage storage prefetch (#36700)

Co-authored-by: Chao Shi <stepinto@live.com>

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)
 

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)

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

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):

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], (

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):

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()

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()

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()

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]