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