[Scheduler] Refresh waiting requests' cached prefixes under FCFS-like policies with LRU eviction

python/sglang/srt/environ.py

@@ -592,6 +592,7 @@ class Envs:
     # ===================================================================
     # Scheduler token budgeting and admission
     # ===================================================================
+    SGLANG_ENABLE_WAITING_PREFIX_REFRESH = EnvBool(True)
     SGLANG_INIT_NEW_TOKEN_RATIO = EnvFloat(0.7)
     SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR = EnvFloat(0.14)
     SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS = EnvInt(600)

python/sglang/srt/managers/schedule_policy.py

@@ -123,6 +123,10 @@ def _use_exact_chunk_fill() -> bool:
     )
     PREFILL_TILE_BUDGET_MODE = "compact"
 
+# Inherited from the LPM/HRRN fallback bound in _determine_active_policy;
+# not a measured optimum.
+WAITING_PREFIX_REFRESH_MAX_QUEUE = 128
+
 
 def _ceil_div(value: int, divisor: int) -> int:
     return -(-value // divisor)
@@ -154,6 +158,20 @@ def estimate_prefill_extend_tile_metrics(
     }
 
 
+def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None:
+    token_ids = req.origin_input_ids + req.output_ids
+    reprefill_tail = tree_cache.swa_reprefill_tail_tokens()
+    key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None
+    tree_cache.refresh_device_prefix(
+        RadixKey(
+            token_ids=token_ids,
+            extra_key=req.extra_key,
+            limit=key_limit,
+            cache_salt=req.cache_salt,
+        )
+    )
+
+
 def match_prefix_for_req(
     tree_cache: BasePrefixCache,
     req: Req,
@@ -265,17 +283,26 @@ def calc_priority(
     ) -> None:
         policy = self._determine_active_policy(waiting_queue)
 
-        # Populate req.num_matched_prefix_tokens at schedule time. Cache-aware policies
-        # set it in _compute_prefix_matches; do the same full match for
-        # cache-agnostic policies when the radix supports it, so the load
-        # snapshot has it. Skip on decode (never prefills).
+        # Skip on decode (never prefills).
         if (
             not isinstance(policy, CacheAwarePolicy)
-            and self.tree_cache.supports_fast_match_prefix()
             and get_disagg().disaggregation_mode != "decode"
         ):
-            for r in waiting_queue:
-                match_prefix_for_req(self.tree_cache, r, include_req=True)
+            # Populate req.num_matched_prefix_tokens at schedule time. Cache-aware
+            # policies set it in _compute_prefix_matches; do the same full match for
+            # cache-agnostic policies when the radix supports it, so the load
+            # snapshot has it.
+            if self.tree_cache.supports_fast_match_prefix():
+                for r in waiting_queue:
+                    match_prefix_for_req(self.tree_cache, r, include_req=True)
+            # Otherwise keep waiting requests' prefixes resident, or the LRU evicts
+            # them while they wait. Head last, so it carries the newest timestamp.
+            elif (
+                envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get()
+                and not self.tree_cache.is_chunk_cache()
+            ):
+                for r in reversed(waiting_queue[:WAITING_PREFIX_REFRESH_MAX_QUEUE]):
+                    refresh_waiting_prefix(self.tree_cache, r)
 
         if self.policy == CacheAgnosticPolicy.FCFS:
             if self.enable_priority_scheduling:

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -424,6 +424,17 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
     def supports_fast_match_prefix(self) -> bool:
         return False
 
+    def refresh_device_prefix(self, key: RadixKey) -> int:
+        """Touch the device-resident prefix of ``key`` so LRU eviction sees it as recent.
+
+        Returns the number of device tokens matched. The scheduler calls this for
+        every waiting request on every round, so it must stay on the device tree:
+        no host tier, storage backend or external connector, which ``match_prefix``
+        may allocate, load or enqueue lookups on. Caches without a side-effect-free
+        device match leave this a no-op.
+        """
+        return 0
+
     def dfs_weight_order(self, node_handles: Sequence[Any]) -> list[int]:
         """Return request indices in depth-first, subtree-weight order."""
         return _dfs_weight_order(self.root_node, node_handles, self.resolve_node_handle)

python/sglang/srt/mem_cache/radix_cache.py

@@ -411,6 +411,16 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
             best_match_node=last_node,
         )
 
+    def refresh_device_prefix(self, key: RadixKey) -> int:
+        key, _ = key.maybe_to_bigram_view(self.is_eagle)
+        if self.disable or len(key) == 0:
+            return 0
+        key = key.page_aligned(self.page_size)
+        if len(key) == 0:
+            return 0
+        value, _ = self._match_prefix_helper(self.root_node, key)
+        return sum(len(v) for v in value)
+
     def insert(self, params: InsertParams) -> InsertResult:
         if self.disable:
             return InsertResult(prefix_len=0)

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -570,6 +570,14 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
     def supports_fast_match_prefix(self) -> bool:
         return self.tree_core.supports_fast_match_prefix()
 
+    def refresh_device_prefix(self, key: RadixKey) -> int:
+        # The walk and its actions only; no session shortcut, finalizers or linker.
+        if self.disable:
+            return 0
+        result = self.tree_core.match_prefix(MatchPrefixParams(key=key))
+        self._apply_cache_actions(result.cache_actions)
+        return len(result.device_indices)
+
     def is_chunk_cache(self) -> bool:
         return self.disable
 

python/sglang/srt/session/streaming_session.py

@@ -26,6 +26,7 @@
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
     from sglang.srt.mem_cache.buffer_mode.pipeline import BufferModePipeline
+    from sglang.srt.mem_cache.radix_cache import RadixKey
     from sglang.srt.mem_cache.storage_prefetch import StoragePrefetchRetries
 
 
@@ -626,6 +627,9 @@ def supports_streaming_session(self) -> bool:
     def is_chunk_cache(self):
         return self.inner.is_chunk_cache()
 
+    def refresh_device_prefix(self, key: RadixKey) -> int:
+        return self.inner.refresh_device_prefix(key)
+
     def is_tree_cache(self):
         return self.inner.is_tree_cache()
 

test/manual/test_schedule_policy.py

@@ -8,7 +8,9 @@
     SchedulePolicy,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache
+from sglang.srt.runtime_context import publish, reset_context
 from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.server_args import ServerArgs
 from sglang.test.test_utils import CustomTestCase
 
 
@@ -26,6 +28,9 @@ def _make_req(rid, origin_input_text, origin_input_ids, sampling_params=None, **
 
 class TestSchedulePolicy(CustomTestCase):
     def setUp(self):
+        reset_context()
+        self.addCleanup(reset_context)
+        publish(ServerArgs(model_path="dummy"), role="test")
         self.tree_cache = RadixCache.create_simulated()
 
     def test_init_with_cache_aware_policy(self):

test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py

@@ -0,0 +1,207 @@
+import time
+import unittest
+import unittest.mock
+from array import array
+from dataclasses import replace
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.schedule_policy import (
+    WAITING_PREFIX_REFRESH_MAX_QUEUE,
+    SchedulePolicy,
+)
+from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import (
+    EvictParams,
+    InsertParams,
+    MatchPrefixParams,
+)
+from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.chunk_cache import ChunkCache
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
+from sglang.srt.mem_cache.unified_cache.components import ComponentType
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
+from sglang.srt.runtime_context import publish, reset_context
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.session.streaming_session import StreamingSession
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=10, suite="base-a-test-cpu")
+
+OLDER_PREFIX = [1, 2, 3, 4]
+NEWER_PREFIX = [5, 6, 7, 8]
+
+
+class TestFcfsWaitingPrefixRefresh(CustomTestCase):
+    """Under FCFS an LRU eviction must not take the cached prefix of a request that
+    is waiting in the queue while a prefix no pending request needs is available."""
+
+    def _make_cache(self, params):
+        return RadixCache(params)
+
+    def setUp(self):
+        reset_context()
+        self.addCleanup(reset_context)
+        publish(ServerArgs(model_path="dummy"), role="test")
+        torch.set_default_device(None)
+        kv_cache = MHATokenToKVPool(
+            size=16,
+            page_size=1,
+            dtype=torch.float16,
+            head_num=1,
+            head_dim=8,
+            layer_num=1,
+            device="cpu",
+            enable_memory_saver=False,
+        )
+        allocator = TokenToKVPoolAllocator(
+            size=16,
+            dtype=torch.float16,
+            device="cpu",
+            kvcache=kv_cache,
+            need_sort=False,
+        )
+        req_to_token_pool = ReqToTokenPool(
+            size=8, max_context_len=64, device="cpu", enable_memory_saver=False
+        )
+        self.cache = self._make_cache(
+            CacheInitParams(
+                disable=False,
+                req_to_token_pool=req_to_token_pool,
+                token_to_kv_pool_allocator=allocator,
+                page_size=1,
+                eviction_policy="lru",
+                enable_kv_cache_events=False,
+            )
+        )
+        self.policy = self._make_policy(self.cache)
+        self._insert(OLDER_PREFIX, [10, 11, 12, 13])
+        time.sleep(0.005)
+        self._insert(NEWER_PREFIX, [20, 21, 22, 23])
+        time.sleep(0.005)
+        self.waiting = self._make_req(1, OLDER_PREFIX + [9])
+
+    def _make_policy(self, cache):
+        return SchedulePolicy(
+            policy="fcfs",
+            tree_cache=cache,
+            enable_hierarchical_cache=False,
+            enable_priority_scheduling=False,
+            schedule_low_priority_values_first=False,
+        )
+
+    def _make_req(self, rid, tokens):
+        return Req(rid, "", array("q", tokens), SamplingParams())
+
+    def _insert(self, tokens, slots):
+        self.cache.insert(
+            InsertParams(
+                key=RadixKey(array("q", tokens)),
+                value=torch.tensor(slots, dtype=torch.int64),
+            )
+        )
+
+    def _matched_len(self, tokens):
+        result = self.cache.match_prefix(
+            MatchPrefixParams(key=RadixKey(array("q", tokens)))
+        )
+        return len(result.device_indices)
+
+    def _assert_older_prefix_survives_one_eviction(self):
+        self.cache.evict(EvictParams(num_tokens=len(NEWER_PREFIX)))
+        self.assertEqual(self._matched_len(OLDER_PREFIX), len(OLDER_PREFIX))
+        self.assertEqual(self._matched_len(NEWER_PREFIX), 0)
+
+    def test_refresh_keeps_the_waiting_request_prefix_resident(self):
+        with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True):
+            self.policy.calc_priority([self.waiting])
+        self._assert_older_prefix_survives_one_eviction()
+
+    def test_refresh_disabled_restores_plain_lru(self):
+        with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(False):
+            self.policy.calc_priority([self.waiting])
+        self.cache.evict(EvictParams(num_tokens=len(OLDER_PREFIX)))
+        self.assertEqual(self._matched_len(OLDER_PREFIX), 0)
+        self.assertEqual(self._matched_len(NEWER_PREFIX), len(NEWER_PREFIX))
+
+    def test_head_of_queue_is_refreshed_last(self):
+        # FCFS admits the head next, so when eviction must take a waiting prefix
+        # it has to be the one furthest from admission.
+        second = self._make_req(2, NEWER_PREFIX + [9])
+        with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True):
+            self.policy.calc_priority([self.waiting, second])
+        self._assert_older_prefix_survives_one_eviction()
+
+    def test_deep_queue_still_refreshes_the_head(self):
+        fillers = [
+            self._make_req(2 + i, [1000 + i])
+            for i in range(WAITING_PREFIX_REFRESH_MAX_QUEUE)
+        ]
+        with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True):
+            self.policy.calc_priority([self.waiting] + fillers)
+        self._assert_older_prefix_survives_one_eviction()
+
+    def test_refresh_never_calls_the_general_matcher(self):
+        # External-cache implementations (FlexKV, LMCache, hierarchical tiers) give
+        # match_prefix side effects: lookups enqueued, KV allocated or loaded per call.
+        calls = []
+        original = type(self.cache).match_prefix
+
+        def spy(cache, params):
+            calls.append(params)
+            return original(cache, params)
+
+        with (
+            unittest.mock.patch.object(type(self.cache), "match_prefix", spy),
+            envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True),
+        ):
+            self.policy.calc_priority([self.waiting])
+        self.assertEqual(calls, [])
+        self._assert_older_prefix_survives_one_eviction()
+
+    def test_streaming_session_wrapper_forwards_the_refresh(self):
+        # --enable-streaming-session wraps caches without native session support;
+        # the wrapper inherits the base no-op unless it forwards explicitly.
+        policy = self._make_policy(StreamingSession(self.cache))
+        with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True):
+            policy.calc_priority([self.waiting])
+        self._assert_older_prefix_survives_one_eviction()
+
+    def test_chunk_cache_skips_the_refresh(self):
+        # --disable-radix-cache deployments have no tree or LRU state to refresh; the
+        # scheduler hot path must not build keys or match for every waiting request.
+        chunk_cache = ChunkCache(
+            CacheInitParams(
+                disable=True,
+                req_to_token_pool=self.cache.req_to_token_pool,
+                token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator,
+                page_size=1,
+                eviction_policy="lru",
+                enable_kv_cache_events=False,
+            )
+        )
+        policy = self._make_policy(chunk_cache)
+        with (
+            unittest.mock.patch.object(ChunkCache, "match_prefix") as match_prefix,
+            unittest.mock.patch.object(ChunkCache, "refresh_device_prefix") as refresh,
+            envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(True),
+        ):
+            policy.calc_priority([self.waiting])
+        match_prefix.assert_not_called()
+        refresh.assert_not_called()
+
+
+class TestFcfsWaitingPrefixRefreshUnified(TestFcfsWaitingPrefixRefresh):
+    """The same contract on UnifiedRadixCache, the default prefix cache."""
+
+    def _make_cache(self, params):
+        return UnifiedRadixCache(replace(params, tree_components=(ComponentType.FULL,)))
+
+
+if __name__ == "__main__":
+    unittest.main()