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

python/sglang/srt/environ.py

@@ -602,6 +602,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

@@ -7,6 +7,7 @@
 from sglang.srt.runtime_context import (
     get_disagg,
     get_exec,
+    get_memory,
     get_schedule,
 )
 from sglang.srt.utils import get_bool_env_var, is_gfx95_supported, is_hip
@@ -31,6 +32,7 @@
 
 import os
 import random
+import time
 from collections import Counter
 from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
@@ -66,7 +68,7 @@
     MatchPrefixParams,
     zero_match_result,
 )
-from sglang.srt.mem_cache.common import match_kv_cache
+from sglang.srt.mem_cache.common import match_kv_cache, touch_waiting_prefix
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 from sglang.srt.mem_cache.unified_cache.components import (
     CacheTransferPhase,
@@ -128,6 +130,11 @@ def _use_exact_chunk_fill() -> bool:
     )
     PREFILL_TILE_BUDGET_MODE = "compact"
 
+# Arbitrary; bounds per-round prefix work on the waiting queue.
+WAITING_QUEUE_PREFIX_MATCH_MAX = 128
+# Arbitrary; bounds the walk cost of touching waiting prefixes.
+WAITING_PREFIX_REFRESH_INTERVAL_S = 0.5
+
 
 def _ceil_div(value: int, divisor: int) -> int:
     return -(-value // divisor)
@@ -195,6 +202,7 @@ def __init__(
         self.schedule_low_priority_values_first = schedule_low_priority_values_first
         self.priority_sign = 1 if schedule_low_priority_values_first else -1
         self._shortest_prefill_calls = 0
+        self._last_waiting_prefix_refresh = float("-inf")
 
         # It is used to find the matching prefix for in-batch prefix caching.
         self.waiting_queue_radix_tree = RadixCache.create_simulated()
@@ -222,9 +230,7 @@ def calc_priority(
                 SchedulePolicy._sort_by_priority_and_fcfs(
                     waiting_queue, self.priority_sign
                 )
-            return
-
-        if isinstance(policy, CacheAwarePolicy):
+        elif isinstance(policy, CacheAwarePolicy):
             temporary_deprioritized = self._compute_prefix_matches(
                 waiting_queue, policy
             )
@@ -268,14 +274,36 @@ def calc_priority(
             else:
                 raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}")
 
+        self._touch_waiting_prefixes(policy, waiting_queue)
+
+    def _touch_waiting_prefixes(self, policy: Policy, waiting_queue: List[Req]) -> None:
+        # Head last, so LRU evicts in reverse admission order; cache-aware policies
+        # already touch while matching, and other eviction strategies ignore recency.
+        if (
+            isinstance(policy, CacheAwarePolicy)
+            or not waiting_queue
+            or not envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get()
+            or get_disagg().disaggregation_mode == "decode"
+            or get_memory().radix_eviction_policy.lower() != "lru"
+            or self.tree_cache.supports_fast_match_prefix()
+            or not self.tree_cache.supports_prefix_sharing()
+        ):
+            return
+        now = time.monotonic()
+        if now - self._last_waiting_prefix_refresh < WAITING_PREFIX_REFRESH_INTERVAL_S:
+            return
+        self._last_waiting_prefix_refresh = now
+        for r in reversed(waiting_queue[:WAITING_QUEUE_PREFIX_MATCH_MAX]):
+            touch_waiting_prefix(r, self.tree_cache)
+
     def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy:
         if (
             self.policy
             in (
                 CacheAwarePolicy.LPM,
                 CacheAwarePolicy.HRRN,
             )
-            and len(waiting_queue) > 128
+            and len(waiting_queue) > WAITING_QUEUE_PREFIX_MATCH_MAX
         ):
             # Turn off the expensive prefix matching and sorting when the #queue is large.
             return CacheAgnosticPolicy.FCFS

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -440,6 +440,10 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
     def supports_fast_match_prefix(self) -> bool:
         return False
 
+    def touch_prefix(self, key: RadixKey) -> None:
+        """Bump LRU recency along ``key``'s cached path. Unlike ``match_prefix``,
+        never allocates, loads, or queries external tiers."""
+
     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/common.py

@@ -172,6 +172,34 @@ def free_kv_row_segments(
         allocator.free_segments(swa_alive)
 
 
+def _req_radix_key(
+    tree_cache: BasePrefixCache,
+    req: Req,
+    token_ids: array[int],
+    max_prefix_len: Optional[int] = None,
+) -> RadixKey:
+    # unified_kv SWA lives in a per-request ring the tree never stores, so a reused
+    # prefix carries stale SWA; cap the match so the window is re-prefilled.
+    reprefill_tail = tree_cache.swa_reprefill_tail_tokens()
+    key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None
+    if max_prefix_len is not None:
+        key_limit = (
+            max_prefix_len if key_limit is None else min(key_limit, max_prefix_len)
+        )
+    return RadixKey(
+        token_ids=token_ids,
+        extra_key=req.extra_key,
+        limit=key_limit,
+        cache_salt=req.cache_salt,
+    )
+
+
+def touch_waiting_prefix(req: Req, tree_cache: BasePrefixCache) -> None:
+    tree_cache.touch_prefix(
+        _req_radix_key(tree_cache, req, req.origin_input_ids + req.output_ids)
+    )
+
+
 def match_kv_cache(
     req: Req,
     tree_cache: BasePrefixCache,
@@ -184,23 +212,9 @@ def match_kv_cache(
     if token_ids is None:
         token_ids = req.origin_input_ids + req.output_ids
 
-    # unified_kv SWA lives in a per-request ring the tree never stores, so a reused
-    # prefix carries stale SWA; cap the match so the window is re-prefilled.
-    reprefill_tail = tree_cache.swa_reprefill_tail_tokens()
-    key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None
-    if max_prefix_len is not None:
-        key_limit = (
-            max_prefix_len if key_limit is None else min(key_limit, max_prefix_len)
-        )
-
     match_result = tree_cache.match_prefix(
         MatchPrefixParams(
-            key=RadixKey(
-                token_ids=token_ids,
-                extra_key=req.extra_key,
-                limit=key_limit,
-                cache_salt=req.cache_salt,
-            ),
+            key=_req_radix_key(tree_cache, req, token_ids, max_prefix_len),
             cow_mamba=cow_mamba,
             req=req,
         )

python/sglang/srt/mem_cache/radix_cache.py

@@ -406,6 +406,15 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult:
             best_match_node=last_node,
         )
 
+    def touch_prefix(self, key: RadixKey) -> None:
+        key, _ = key.maybe_to_bigram_view(self.is_eagle)
+        if self.disable or len(key) == 0:
+            return
+        key = key.page_aligned(self.page_size)
+        if len(key) == 0:
+            return
+        self._match_prefix_helper(self.root_node, key)
+
     def insert(self, params: InsertParams) -> InsertResult:
         if self.disable:
             return InsertResult(prefix_len=0)

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -599,6 +599,13 @@ def _match_tree(self, params: MatchPrefixParams) -> MatchResult:
     def supports_fast_match_prefix(self) -> bool:
         return self.tree_core.supports_fast_match_prefix()
 
+    def touch_prefix(self, key: RadixKey) -> None:
+        # Tree walk only; skips match_prefix's session shortcut, finalizers and linker.
+        if self.disable:
+            return
+        result = self.tree_core.match_prefix(MatchPrefixParams(key=key))
+        self._apply_cache_actions(result.cache_actions)
+
     def supports_prefix_sharing(self) -> bool:
         return not self.disable
 

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,103 @@
+import time
+import unittest
+from array import array
+
+import torch
+
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.schedule_policy import 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.memory_pool import MHATokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
+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.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+OLDER_PREFIX = [1, 2, 3, 4]
+NEWER_PREFIX = [5, 6, 7, 8]
+
+
+class TestFcfsWaitingPrefixRefresh(CustomTestCase):
+    def setUp(self):
+        reset_context()
+        self.addCleanup(reset_context)
+        torch.set_default_device(None)
+
+    def _run_and_evict_one_prefix(self, eviction_policy):
+        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,
+        )
+        cache = RadixCache(
+            CacheInitParams(
+                disable=False,
+                req_to_token_pool=ReqToTokenPool(
+                    size=8, max_context_len=64, device="cpu", enable_memory_saver=False
+                ),
+                token_to_kv_pool_allocator=allocator,
+                page_size=1,
+                eviction_policy=eviction_policy,
+                enable_kv_cache_events=False,
+            )
+        )
+        for tokens, slots in (
+            (OLDER_PREFIX, [10, 11, 12, 13]),
+            (NEWER_PREFIX, [20, 21, 22, 23]),
+        ):
+            cache.insert(
+                InsertParams(
+                    key=RadixKey(array("q", tokens)),
+                    value=torch.tensor(slots, dtype=torch.int64),
+                )
+            )
+            time.sleep(0.005)
+        queue = [
+            Req(rid, "", array("q", tokens + [9]), SamplingParams())
+            for rid, tokens in ((1, OLDER_PREFIX), (2, NEWER_PREFIX))
+        ]
+        publish(
+            ServerArgs(model_path="dummy", radix_eviction_policy=eviction_policy),
+            role="test",
+        )
+        SchedulePolicy("fcfs", cache, False, False, False).calc_priority(queue)
+        cache.evict(EvictParams(num_tokens=len(OLDER_PREFIX)))
+        return [
+            cache.match_prefix(
+                MatchPrefixParams(key=RadixKey(array("q", tokens)))
+            ).device_prefix_len
+            for tokens in (OLDER_PREFIX, NEWER_PREFIX)
+        ]
+
+    def test_lru_evicts_the_prefix_furthest_from_admission(self):
+        self.assertEqual(self._run_and_evict_one_prefix("lru"), [4, 0])
+
+    def test_mru_keeps_plain_order(self):
+        # MRU evicts the most recent first, so a refresh would invert its intent.
+        self.assertEqual(self._run_and_evict_one_prefix("mru"), [4, 0])
+
+
+if __name__ == "__main__":
+    unittest.main()