[Fix] hicache: vote the prefetch timeout once per tick, not per admission poll

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -4,6 +4,7 @@
 import logging
 import threading
 import time
+import zlib
 from array import array
 from contextlib import nullcontext
 from dataclasses import replace
@@ -190,6 +191,7 @@ class _OngoingLoadBack(NamedTuple):
 class _StoragePrefetchWait(msgspec.Struct):
     timer: PollTimer = msgspec.field(default_factory=PollTimer)
     timeout_s: float = 0.0
+    # Rank-agreed verdict of the per-tick vote; admission polls only read it.
     expired: bool = False
 
 
@@ -2643,18 +2645,65 @@ def _can_terminate_prefetch(self, operation: PrefetchOperation) -> bool:
             # progress, rank-min queue drain); a vote here only adds collectives.
             return False
         if self.prefetch_stop_policy == "timeout":
-            # Wall clocks differ across ranks: PP0 decides, TP takes MAX (any
-            # rank timed out) and _all_reduce broadcasts the verdict along PP.
-            should_terminate = False
-            if self.pp_rank == 0:
-                should_terminate = self._prefetch_deadline_expired(operation)
-            should_terminate_tensor = torch.tensor(
-                int(should_terminate), dtype=torch.int, device="cpu"
-            )
-            self._all_reduce(should_terminate_tensor, torch.distributed.ReduceOp.MAX)
-            return should_terminate_tensor.item() == 1
+            # The verdict is agreed once per tick by _vote_prefetch_deadlines,
+            # which also terminates the operation (the poll then takes the
+            # is_terminated path); a poll only reads it, so ranks may poll
+            # different requests (batch_is_full, allocatable-request exits)
+            # without a collective.
+            wait = self._storage_prefetch_wait.get(operation.handle)
+            return wait is not None and wait.expired
         return True
 
+    def _vote_prefetch_deadlines(self) -> None:
+        """One MAX vote per tick over every live prefetch's `timeout` budget.
+
+        Wall clocks differ across ranks: PP0 decides per operation, TP takes
+        MAX (any rank timed out) and _all_reduce relays the verdict along PP.
+        The packed tensor has one slot per `ongoing_prefetch` entry in
+        insertion order; this tick's ready-count sync asserted that set is
+        rank-identical (_hicache_sync_digest), and the drain between them
+        only retires in lockstep. With no live prefetch no rank issues a
+        collective. A wait that has not started yet (no blocked poll on PP0)
+        contributes 0 until the next tick. An expired operation is terminated
+        here so the transfer worker and the hit drain stop spending on it;
+        its retirement stays with the admission poll or the drain.
+        """
+        if self.prefetch_stop_policy != "timeout" or not self.ongoing_prefetch:
+            return
+        operations = [info.operation for info in self.ongoing_prefetch.values()]
+        expired = [0] * len(operations)
+        if self.pp_rank == 0:
+            expired = [
+                int(self._prefetch_deadline_expired(operation))
+                for operation in operations
+            ]
+        verdicts = torch.tensor(expired, dtype=torch.int, device="cpu")
+        self._all_reduce(verdicts, torch.distributed.ReduceOp.MAX)
+        for operation, verdict in zip(operations, verdicts.tolist()):
+            if verdict:
+                self._mark_prefetch_deadline_expired(operation)
+                self.cache_controller.terminate_prefetch(operation)
+
+    def _hicache_sync_digest(self) -> int:
+        """Value piggybacked on the ready-count MIN sync: equal on every rank
+        iff the write_back reclaim victim order matched and, when the deadline
+        vote reads it, so did the live prefetch set (`ongoing_prefetch` in
+        insertion order). The admission polls retire per request, so a rank
+        that left the loop earlier than its peers (#30760) is caught here, in
+        the same tick on every rank, instead of hanging in a vote of a
+        different length. Fits int64: 42-bit digest * 1000003 + crc32."""
+        digest = self.tree_core.write_back_duplicate_reclaim_digest
+        if (
+            self.enable_storage
+            and self.prefetch_stop_policy == "timeout"
+            and self.ongoing_prefetch
+        ):
+            live = "\0".join(
+                f"{handle.rid}\0{handle.attempt_id}" for handle in self.ongoing_prefetch
+            )
+            digest = digest * 1000003 + zlib.crc32(live.encode()) + 1
+        return digest
+
     def has_ongoing_prefetch(self, handle: CacheRequestHandle) -> bool:
         return handle in self.ongoing_prefetch
 
@@ -3396,6 +3445,9 @@ def _try_alloc_storage_hit(operation) -> bool:
             if info is None:
                 return True  # aborted/cleaned; nothing to retry
             if operation.is_terminated():
+                if operation.deadline_expired:
+                    # Cut by the tick vote while parked: same reason as the poll.
+                    self._finish_storage_prefetch(request, 0, reason="deadline")
                 self.revoke_pending_prefetch(request)
                 return True
 
@@ -3791,8 +3843,8 @@ def _sync_hicache_ready_counts(
             )
 
         # Piggybacked TP check: [digest, -digest] MIN-reduces to [min, -max],
-        # equal iff reclaim victim order matched on every rank.
-        digest = self.tree_core.write_back_duplicate_reclaim_digest
+        # equal iff reclaim victim order and live prefetch set matched on every rank.
+        digest = self._hicache_sync_digest()
         exchange = self._shm_count_exchange()
         if exchange is not None:
             values = [write_acks, load_acks, *storage_queue_sizes, *fold_u64(digest)]
@@ -3801,7 +3853,8 @@ def _sync_hicache_ready_counts(
             if len(values) <= exchange.width:
                 rows = exchange.allgather(values)
                 assert all(row[-2:] == rows[0][-2:] for row in rows), (
-                    "write_back duplicate-reclaim victims diverged across TP ranks"
+                    "write_back duplicate-reclaim victims or live storage "
+                    "prefetches diverged across TP ranks"
                 )
                 mins = [min(col) for col in zip(*rows)]
                 return mins[0], mins[1], tuple(mins[2:-2]), extra_pool_names
@@ -3823,7 +3876,8 @@ def _sync_hicache_ready_counts(
 
         count_values = list(map(int, ready_counts.tolist()))
         assert digest == count_values[-2] and digest == -count_values[-1], (
-            "write_back duplicate-reclaim victims diverged across PP/TP ranks"
+            "write_back duplicate-reclaim victims or live storage prefetches "
+            "diverged across PP/TP ranks"
         )
         return (
             count_values[0],
@@ -4105,6 +4159,9 @@ def check_hicache_events(self) -> None:
                 extra_release_counts=extra_release_counts,
                 log_metrics=True,
             )
+        if self.enable_storage:
+            # After the drain: a prefetch it retired is not voted on.
+            self._vote_prefetch_deadlines()
         if self.buffer_pipeline is not None:
             self.buffer_pipeline.flush_pending_writes()
         if self.enable_storage_metrics and self.storage_metrics_collector is not None:

test/registered/unit/managers/scheduler_components/test_shm_tag_exchange.py

@@ -50,6 +50,9 @@ def _hicache_tp_rank(rank: int, extra_pools: int) -> UnifiedRadixCache:
     )
     cache.pp_rank, cache.pp_size, cache.pp_group = 0, 1, None
     cache.host_memory_mode = "cache"
+    # The tick reads the stop policy and the live prefetch set (deadline vote).
+    cache.prefetch_stop_policy = "wait_complete"
+    cache.ongoing_prefetch = {}
     cache.attn_cp_group = cache.attn_tp_group = None
     cache.tp_group, cache.tp_world_size = torch.distributed.group.WORLD, 2
     cache.cache_controller = SimpleNamespace(

test/registered/unit/mem_cache/test_hicache_pp_sync_drain.py

@@ -3,11 +3,17 @@
 import unittest
 from queue import Queue
 from types import SimpleNamespace
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import torch
 
-from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+from sglang.srt.mem_cache.base_prefix_cache import CacheRequestHandle
+from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
+from sglang.srt.mem_cache.unified_radix_cache import (
+    UnifiedRadixCache,
+    _OngoingPrefetch,
+    _StoragePrefetchWait,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
@@ -23,6 +29,10 @@ def _make_cache(self, pp_rank, write_ready, load_ready):
         cache.pp_rank = pp_rank
         cache.pp_size = 2
         cache.host_memory_mode = "cache"
+        # The per-tick deadline vote runs under storage; no live prefetch
+        # means it issues no collective, which the asserts below pin.
+        cache.prefetch_stop_policy = "timeout"
+        cache.ongoing_prefetch = {}
         cache.enable_storage_metrics = False
         cache.storage_metrics_collector = None
         cache.buffer_pipeline = None
@@ -48,6 +58,79 @@ def _make_cache(self, pp_rank, write_ready, load_ready):
         )
         return cache
 
+    def _arm_live_prefetch(self, cache):
+        """Cache-mode storage with one live prefetch on this PP stage."""
+        cache.tree_core.enable_storage = True
+        cache._drain_storage_control_queues_impl = MagicMock()
+        cache.storage_prefetch_retries = StoragePrefetchRetries()
+        cache._storage_prefetch_wait = {}
+        cc = cache.cache_controller
+        for name in (
+            "prefetch_hit_queue",
+            "ack_prefetch_queue",
+            "ack_backup_queue",
+            "host_mem_release_queue",
+        ):
+            setattr(cc, name, Queue())
+        cc.extra_host_mem_release_queues = {}
+        cc.terminate_prefetch = MagicMock()
+        handle = CacheRequestHandle("r", 0)
+        operation = SimpleNamespace(
+            handle=handle,
+            deadline_expired=False,
+            host_indices=None,
+            hash_value=[],
+            completed_tokens=0,
+        )
+        cache.ongoing_prefetch[handle] = _OngoingPrefetch(
+            0, None, None, operation, None, {}
+        )
+        return handle, operation
+
+    def test_pp_relays_the_deadline_vote_after_the_counts(self):
+        """PP0 evaluates the clocks and votes; a follower contributes zeros,
+        receives the verdict through the same tag after the counts, and marks
+        and terminates the operation without reading its own clock."""
+        leader = self._make_cache(0, [True], [True])
+        handle, operation = self._arm_live_prefetch(leader)
+        wait = _StoragePrefetchWait(timeout_s=1.0)
+        with patch("time.perf_counter", return_value=100.0):
+            wait.timer.observe(ready=False)
+        leader._storage_prefetch_wait[handle] = wait
+        with patch("time.perf_counter", return_value=103.0):
+            leader.check_hicache_events()
+
+        counts, vote = leader._all_reduce.call_args_list
+        self.assertIs(counts.args[1], torch.distributed.ReduceOp.MIN)
+        self.assertEqual(vote.args[0].tolist(), [1])
+        self.assertIs(vote.args[1], torch.distributed.ReduceOp.MAX)
+        self.assertTrue(wait.expired)
+        self.assertTrue(operation.deadline_expired)
+        leader.cache_controller.terminate_prefetch.assert_called_once_with(operation)
+
+        follower = self._make_cache(1, [True], [True])
+        handle, operation = self._arm_live_prefetch(follower)
+        follower._prefetch_deadline_expired = MagicMock()
+        relayed = []
+
+        def relay_from_pp0(data, _):
+            relayed.append(data.tolist())
+            if len(relayed) == 1:
+                data[0], data[1] = 1, 1  # counts; digest slots stay local
+            else:
+                data.fill_(1)  # PP0's verdict
+
+        follower._all_reduce.side_effect = relay_from_pp0
+        follower.check_hicache_events()
+
+        self.assertEqual(relayed[1], [0])
+        self.assertEqual(follower._all_reduce.call_count, 2)
+        follower._prefetch_deadline_expired.assert_not_called()
+        self.assertTrue(follower._storage_prefetch_wait[handle].expired)
+        self.assertTrue(operation.deadline_expired)
+        follower.cache_controller.terminate_prefetch.assert_called_once_with(operation)
+        follower.writing_check.assert_called_once_with(finish_count=1)
+
     def test_pp_batches_write_and_load_counts_once(self):
         leader = self._make_cache(0, [True, False], [True, True])
         leader.check_hicache_events()