From f2fe3d0b15a162d2f2e3dbc9c4412bb38067cff9 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Mon, 28 Sep 2026 14:34:41 -0700 Subject: [PATCH 01/11] [Scheduler] Keep waiting requests' cached prefixes resident under cache-agnostic policies Under FCFS and the other cache-agnostic policies a waiting request's prefix is never matched, and so never refreshed in the LRU, until the admission loop reaches it. A request queued behind prompts that exhaust the budget therefore ages into the eviction victim, loses its prefix while it waits, and is admitted cold, which lengthens the queue for the requests behind it. Run the existing waiting-queue pre-match for cache-agnostic policies too, bounded by the queue length at which LPM already stops matching, gated by SGLANG_ENABLE_WAITING_PREFIX_REFRESH (default on). Eviction order changes; budget and locks do not. Signed-off-by: Pei Li --- python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/schedule_policy.py | 19 ++- test/manual/test_schedule_policy.py | 5 + .../test_fcfs_waiting_prefix_refresh.py | 115 ++++++++++++++++++ 4 files changed, 136 insertions(+), 4 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 9af31cae8138..6cfb3d5ad5b4 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d5f1b80ae750..3fb51ce1b5c1 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -123,6 +123,10 @@ def _use_exact_chunk_fill() -> bool: ) PREFILL_TILE_BUDGET_MODE = "compact" +# Same queue length at which LPM stops matching the waiting queue; bounds the +# per-round cost of the cache-agnostic prefix refresh. +WAITING_PREFIX_REFRESH_MAX_QUEUE = 128 + def _ceil_div(value: int, divisor: int) -> int: return -(-value // divisor) @@ -266,12 +270,19 @@ def calc_priority( 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). + # set it in _compute_prefix_matches; do the same full match for cache-agnostic + # policies. The match also refreshes last_access_time on the waiting request's + # prefix; without it the LRU evicts that prefix while the request waits behind + # prompts that exhaust the budget. Skip on decode (never prefills). if ( not isinstance(policy, CacheAwarePolicy) - and self.tree_cache.supports_fast_match_prefix() + and ( + self.tree_cache.supports_fast_match_prefix() + or ( + envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get() + and len(waiting_queue) <= WAITING_PREFIX_REFRESH_MAX_QUEUE + ) + ) and get_disagg().disaggregation_mode != "decode" ): for r in waiting_queue: diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index 7258e31f4017..6afc173f7370 100644 --- a/test/manual/test_schedule_policy.py +++ b/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): diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py new file mode 100644 index 000000000000..81f7949f0392 --- /dev/null +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -0,0 +1,115 @@ +import time +import unittest +from array import array + +import torch + +from sglang.srt.environ import envs +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=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 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 = RadixCache( + 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 = SchedulePolicy( + policy="fcfs", + tree_cache=self.cache, + enable_hierarchical_cache=False, + enable_priority_scheduling=False, + schedule_low_priority_values_first=False, + ) + 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 = Req(1, "", array("q", OLDER_PREFIX + [9]), 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 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.assertEqual(len(self.waiting.prefix_indices), len(OLDER_PREFIX)) + 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_disabled_restores_plain_lru(self): + with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(False): + self.policy.calc_priority([self.waiting]) + self.assertEqual(len(self.waiting.prefix_indices), 0) + 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)) + + +if __name__ == "__main__": + unittest.main() From 3a79e2de106fe7078c442e7f8e620f681460160e Mon Sep 17 00:00:00 2001 From: Pei Li Date: Mon, 28 Sep 2026 23:50:45 -0700 Subject: [PATCH 02/11] Document the 128 waiting-queue cap on the prefix refresh Signed-off-by: Pei Li --- python/sglang/srt/managers/schedule_policy.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 3fb51ce1b5c1..53312179c911 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -123,8 +123,10 @@ def _use_exact_chunk_fill() -> bool: ) PREFILL_TILE_BUDGET_MODE = "compact" -# Same queue length at which LPM stops matching the waiting queue; bounds the -# per-round cost of the cache-agnostic prefix refresh. +# Bounds the per-round cost of the cache-agnostic prefix refresh to what LPM +# already pays: _determine_active_policy drops LPM/HRRN to FCFS above this same +# queue length (#1896). Inherited precedent, not a measured optimum; above it a +# waiting request is again unprotected from eviction. WAITING_PREFIX_REFRESH_MAX_QUEUE = 128 From a2f236794e75eaaa8e3b50c054d3f7976cbeffc8 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 30 Sep 2026 19:18:08 -0700 Subject: [PATCH 03/11] Refresh waiting prefixes on the device tree only and skip chunk caches The per-round refresh called match_prefix, whose external-cache implementations have side effects per call (FlexKV enqueues a lookup task and can allocate or load KV; hierarchical tiers consult the host). Add BasePrefixCache.refresh_device_prefix, a device-tree-only timestamp touch (no-op by default; RadixCache walks its tree and returns the matched length), and route the cache-agnostic refresh through it. Guard the refresh with is_chunk_cache so --disable-radix-cache deployments do no per-request work. The fast-match path for caches that support it is unchanged. Tests: the refresh never calls match_prefix; a chunk cache triggers neither match_prefix nor refresh_device_prefix. Signed-off-by: Pei Li --- python/sglang/srt/managers/schedule_policy.py | 56 ++++++++++++++----- .../sglang/srt/mem_cache/base_prefix_cache.py | 11 ++++ python/sglang/srt/mem_cache/radix_cache.py | 13 +++++ .../test_fcfs_waiting_prefix_refresh.py | 54 +++++++++++++++++- 4 files changed, 118 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 53312179c911..1028be297dd6 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -160,6 +160,29 @@ def estimate_prefill_extend_tile_metrics( } +def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None: + """Touch ``req``'s device-resident prefix and record its length on the request. + + Device-tree-only counterpart of :func:`match_prefix_for_req` for the per-round + waiting-queue refresh: it never calls ``match_prefix``, whose external-cache + implementations can allocate, load or enqueue lookups per call. + """ + 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 + matched = tree_cache.refresh_device_prefix( + RadixKey( + token_ids=token_ids, + extra_key=req.extra_key, + limit=key_limit, + cache_salt=req.cache_salt, + ) + ) + req.num_matched_prefix_tokens = min( + matched, req._compute_max_prefix_len(len(token_ids)) + ) + + def match_prefix_for_req( tree_cache: BasePrefixCache, req: Req, @@ -271,24 +294,29 @@ 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. The match also refreshes last_access_time on the waiting request's - # prefix; without it the LRU evicts that prefix while the request waits behind - # prompts that exhaust the budget. Skip on decode (never prefills). + # Skip on decode (never prefills). if ( not isinstance(policy, CacheAwarePolicy) - and ( - self.tree_cache.supports_fast_match_prefix() - or ( - envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get() - and len(waiting_queue) <= WAITING_PREFIX_REFRESH_MAX_QUEUE - ) - ) 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 the waiting requests' cached prefixes resident: without a + # touch the LRU evicts a prefix while its request waits behind prompts that + # exhaust the budget. Device tree only, so no host-tier or storage lookup + # runs as a side effect; nothing to touch under a chunk cache. + elif ( + envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.get() + and not self.tree_cache.is_chunk_cache() + and len(waiting_queue) <= WAITING_PREFIX_REFRESH_MAX_QUEUE + ): + for r in waiting_queue: + refresh_waiting_prefix(self.tree_cache, r) if self.policy == CacheAgnosticPolicy.FCFS: if self.enable_priority_scheduling: diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 41bfbe815dbe..f7f4d4734f44 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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. Must never reach a host tier, + storage backend or external connector: the scheduler calls this for every + waiting request on every round, and ``match_prefix`` on those caches can + allocate, load or enqueue lookups as a side effect. Caches without a + device tree (or whose device match has side effects) 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) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index fa5f9df2e61a..fce4482c5dc1 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -411,6 +411,19 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: best_match_node=last_node, ) + def refresh_device_prefix(self, key: RadixKey) -> int: + # The device-tree walk only: refreshes last_access_time along the matched + # path (and may split a node at the boundary, as match_prefix does) without + # the host-tier or storage lookups subclasses add to match_prefix. + 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) diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 81f7949f0392..155da368adaf 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -1,5 +1,6 @@ import time import unittest +import unittest.mock from array import array import torch @@ -14,6 +15,7 @@ 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.runtime_context import publish, reset_context @@ -97,7 +99,7 @@ def _matched_len(self, tokens): 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.assertEqual(len(self.waiting.prefix_indices), len(OLDER_PREFIX)) + self.assertEqual(self.waiting.num_matched_prefix_tokens, len(OLDER_PREFIX)) 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) @@ -105,11 +107,59 @@ def test_refresh_keeps_the_waiting_request_prefix_resident(self): def test_refresh_disabled_restores_plain_lru(self): with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(False): self.policy.calc_priority([self.waiting]) - self.assertEqual(len(self.waiting.prefix_indices), 0) + self.assertEqual(self.waiting.num_matched_prefix_tokens, 0) 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_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. + # The per-round refresh must stay on the device tree. + 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.assertEqual(self.waiting.num_matched_prefix_tokens, len(OLDER_PREFIX)) + + 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 = SchedulePolicy( + policy="fcfs", + tree_cache=chunk_cache, + enable_hierarchical_cache=False, + enable_priority_scheduling=False, + schedule_low_priority_values_first=False, + ) + 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() + if __name__ == "__main__": unittest.main() From 570867145f56a21db38c84a6e484ec367f7d5c44 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Wed, 30 Sep 2026 19:24:03 -0700 Subject: [PATCH 04/11] Give UnifiedRadixCache the device-only prefix refresh UnifiedRadixCache is the default prefix cache, so the refresh must reach it: walk the tree core (which refreshes last_access_time along the path) and apply the walk's bookkeeping actions, skipping the streaming-session shortcut, the component finalizers and the external linker that match_prefix adds. The refresh tests now run against both RadixCache and UnifiedRadixCache. Signed-off-by: Pei Li --- .../sglang/srt/mem_cache/unified_radix_cache.py | 12 ++++++++++++ .../mem_cache/test_fcfs_waiting_prefix_refresh.py | 15 ++++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 421a21572284..413e71e53c35 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -570,6 +570,18 @@ 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 tree walk alone: it refreshes last_access_time along the matched path + # and may split a node at the boundary, whose bookkeeping action is applied + # here. Skips what match_prefix adds on top -- the streaming-session + # shortcut, the component finalizers and the external linker -- so no + # request slot is written and no storage lookup is issued. + 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 diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 155da368adaf..29066919ec75 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -2,6 +2,7 @@ import unittest import unittest.mock from array import array +from dataclasses import replace import torch @@ -18,6 +19,8 @@ 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 @@ -34,6 +37,9 @@ 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) @@ -59,7 +65,7 @@ def setUp(self): req_to_token_pool = ReqToTokenPool( size=8, max_context_len=64, device="cpu", enable_memory_saver=False ) - self.cache = RadixCache( + self.cache = self._make_cache( CacheInitParams( disable=False, req_to_token_pool=req_to_token_pool, @@ -161,5 +167,12 @@ def test_chunk_cache_skips_the_refresh(self): 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() From 998f9e969ce9950c1c844c95b0785661c3c02c97 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Thu, 1 Oct 2026 14:58:24 -0700 Subject: [PATCH 05/11] Refresh waiting prefixes head-last at any depth, also through StreamingSession Reviewer follow-ups on the waiting-prefix refresh: - Iterate the refreshed slice tail to head so the next request FCFS admits carries the newest timestamp instead of the oldest. - Slice the first WAITING_PREFIX_REFRESH_MAX_QUEUE requests instead of switching the refresh off above that depth; a deep queue is when prefixes age out, so the head stays protected at any depth. - Forward refresh_device_prefix through StreamingSession, which otherwise inherits the base no-op (--enable-streaming-session wraps RadixCache subclasses in it). - Drop the num_matched_prefix_tokens write in refresh_waiting_prefix: nothing reads it on the cache-agnostic path, since waiting_queue_prefix_matched() is false there. - Keep the device-only contract once, in the BasePrefixCache docstring; cut the other copies to one or two lines. Tests: three new regression cases (head refreshed last, 129-deep queue still refreshes the head, StreamingSession wrapper reaches the inner tree), each failing on the previous code; existing cases assert eviction results only. Signed-off-by: Pei Li --- python/sglang/srt/managers/schedule_policy.py | 26 ++----- .../sglang/srt/mem_cache/base_prefix_cache.py | 10 +-- python/sglang/srt/mem_cache/radix_cache.py | 3 - .../srt/mem_cache/unified_radix_cache.py | 6 +- .../sglang/srt/session/streaming_session.py | 4 + .../test_fcfs_waiting_prefix_refresh.py | 73 +++++++++++++------ 6 files changed, 67 insertions(+), 55 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 1028be297dd6..57069b66e5da 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -123,10 +123,8 @@ def _use_exact_chunk_fill() -> bool: ) PREFILL_TILE_BUDGET_MODE = "compact" -# Bounds the per-round cost of the cache-agnostic prefix refresh to what LPM -# already pays: _determine_active_policy drops LPM/HRRN to FCFS above this same -# queue length (#1896). Inherited precedent, not a measured optimum; above it a -# waiting request is again unprotected from eviction. +# Inherited from the LPM/HRRN fallback bound in _determine_active_policy; +# not a measured optimum. WAITING_PREFIX_REFRESH_MAX_QUEUE = 128 @@ -161,16 +159,10 @@ def estimate_prefill_extend_tile_metrics( def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None: - """Touch ``req``'s device-resident prefix and record its length on the request. - - Device-tree-only counterpart of :func:`match_prefix_for_req` for the per-round - waiting-queue refresh: it never calls ``match_prefix``, whose external-cache - implementations can allocate, load or enqueue lookups per call. - """ 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 - matched = tree_cache.refresh_device_prefix( + tree_cache.refresh_device_prefix( RadixKey( token_ids=token_ids, extra_key=req.extra_key, @@ -178,9 +170,6 @@ def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None: cache_salt=req.cache_salt, ) ) - req.num_matched_prefix_tokens = min( - matched, req._compute_max_prefix_len(len(token_ids)) - ) def match_prefix_for_req( @@ -306,16 +295,13 @@ def calc_priority( 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 the waiting requests' cached prefixes resident: without a - # touch the LRU evicts a prefix while its request waits behind prompts that - # exhaust the budget. Device tree only, so no host-tier or storage lookup - # runs as a side effect; nothing to touch under a chunk cache. + # 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() - and len(waiting_queue) <= WAITING_PREFIX_REFRESH_MAX_QUEUE ): - for r in waiting_queue: + for r in reversed(waiting_queue[:WAITING_PREFIX_REFRESH_MAX_QUEUE]): refresh_waiting_prefix(self.tree_cache, r) if self.policy == CacheAgnosticPolicy.FCFS: diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index f7f4d4734f44..c1d390b8afcf 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -427,11 +427,11 @@ def supports_fast_match_prefix(self) -> bool: 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. Must never reach a host tier, - storage backend or external connector: the scheduler calls this for every - waiting request on every round, and ``match_prefix`` on those caches can - allocate, load or enqueue lookups as a side effect. Caches without a - device tree (or whose device match has side effects) leave this a no-op. + 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 diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index fce4482c5dc1..90e9444e1e43 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -412,9 +412,6 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: ) def refresh_device_prefix(self, key: RadixKey) -> int: - # The device-tree walk only: refreshes last_access_time along the matched - # path (and may split a node at the boundary, as match_prefix does) without - # the host-tier or storage lookups subclasses add to match_prefix. key, _ = key.maybe_to_bigram_view(self.is_eagle) if self.disable or len(key) == 0: return 0 diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 413e71e53c35..a7dd38b956df 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -571,11 +571,7 @@ def supports_fast_match_prefix(self) -> bool: return self.tree_core.supports_fast_match_prefix() def refresh_device_prefix(self, key: RadixKey) -> int: - # The tree walk alone: it refreshes last_access_time along the matched path - # and may split a node at the boundary, whose bookkeeping action is applied - # here. Skips what match_prefix adds on top -- the streaming-session - # shortcut, the component finalizers and the external linker -- so no - # request slot is written and no storage lookup is issued. + # 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)) diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index a90e61235519..b4517ddaec2e 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 29066919ec75..41580be70096 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -8,7 +8,10 @@ from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import Req -from sglang.srt.managers.schedule_policy import SchedulePolicy +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, @@ -24,6 +27,7 @@ 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 @@ -75,18 +79,24 @@ def setUp(self): enable_kv_cache_events=False, ) ) - self.policy = SchedulePolicy( + 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=self.cache, + tree_cache=cache, enable_hierarchical_cache=False, enable_priority_scheduling=False, schedule_low_priority_values_first=False, ) - 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 = Req(1, "", array("q", OLDER_PREFIX + [9]), SamplingParams()) + + def _make_req(self, rid, tokens): + return Req(rid, "", array("q", tokens), SamplingParams()) def _insert(self, tokens, slots): self.cache.insert( @@ -102,26 +112,43 @@ def _matched_len(self, tokens): ) return len(result.device_indices) - 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.assertEqual(self.waiting.num_matched_prefix_tokens, len(OLDER_PREFIX)) + 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.assertEqual(self.waiting.num_matched_prefix_tokens, 0) 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. - # The per-round refresh must stay on the device tree. calls = [] original = type(self.cache).match_prefix @@ -135,7 +162,15 @@ def spy(cache, params): ): self.policy.calc_priority([self.waiting]) self.assertEqual(calls, []) - self.assertEqual(self.waiting.num_matched_prefix_tokens, len(OLDER_PREFIX)) + 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 @@ -150,13 +185,7 @@ def test_chunk_cache_skips_the_refresh(self): enable_kv_cache_events=False, ) ) - policy = SchedulePolicy( - policy="fcfs", - tree_cache=chunk_cache, - enable_hierarchical_cache=False, - enable_priority_scheduling=False, - schedule_low_priority_values_first=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, From 80f7932cf971594fd55cc10ac08d4e7792c7bae9 Mon Sep 17 00:00:00 2001 From: Pei Li Date: Fri, 2 Oct 2026 12:29:22 -0700 Subject: [PATCH 06/11] Build the request radix key in one place for the matcher and the refresh refresh_waiting_prefix copied the key construction from match_prefix_for_req (token ids, SWA re-prefill limit, extra_key, cache_salt). A later edit to one copy would silently send the refresh down a different tree path. Both now call _req_radix_key. Signed-off-by: Pei Li --- python/sglang/srt/managers/schedule_policy.py | 39 +++++++++---------- 1 file changed, 18 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 57069b66e5da..b98984c5b376 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -158,17 +158,26 @@ 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 +def _req_radix_key( + tree_cache: BasePrefixCache, req: Req, token_ids: array[int] +) -> RadixKey: + # unified_kv SWA lives in a per-request ring that's not content-stable and is + # never stored in the radix tree, so a reused prefix carries stale SWA. Cap + # the match by the trailing sliding window so it gets re-prefilled, rewriting + # this request's SWA ring. No-op for other layouts. reprefill_tail = tree_cache.swa_reprefill_tail_tokens() key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None + return RadixKey( + token_ids=token_ids, + extra_key=req.extra_key, + limit=key_limit, + cache_salt=req.cache_salt, + ) + + +def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None: tree_cache.refresh_device_prefix( - RadixKey( - token_ids=token_ids, - extra_key=req.extra_key, - limit=key_limit, - cache_salt=req.cache_salt, - ) + _req_radix_key(tree_cache, req, req.origin_input_ids + req.output_ids) ) @@ -183,21 +192,9 @@ def match_prefix_for_req( if token_ids is None: token_ids = req.origin_input_ids + req.output_ids - # unified_kv SWA lives in a per-request ring that's not content-stable and is - # never stored in the radix tree, so a reused prefix carries stale SWA. Cap - # the match by the trailing sliding window so it gets re-prefilled, rewriting - # this request's SWA ring. No-op for other layouts. - reprefill_tail = tree_cache.swa_reprefill_tail_tokens() - key_limit = max(0, len(token_ids) - reprefill_tail) if reprefill_tail else None - 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), cow_mamba=cow_mamba, req=req if include_req else None, ) From 60bdb8c8d2ef9461400005fac722e17b31338bcb Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 8 Oct 2026 15:59:59 -0700 Subject: [PATCH 07/11] [Scheduler] Refresh waiting prefixes after sorting, throttled and LRU-only - Refresh after the queue is sorted, head last, so eviction follows the admission order of LOF, random, routing-key and priority scheduling too. - Throttle to once per WAITING_PREFIX_REFRESH_INTERVAL_S (0.5 s); a refresh only has to keep waiting prefixes newer than idle ones. - Skip unless --radix-eviction-policy is lru: other strategies rank by other keys, and MRU would evict the refreshed prefixes first. - refresh_device_prefix returns nothing; its docstring now states the call rate and that host-backed nodes are touched in HiCache mode. - Tests: cover the skips, the throttle and the sorted order; drop the cases the others already cover. --- python/sglang/srt/managers/schedule_policy.py | 45 ++++++--- .../sglang/srt/mem_cache/base_prefix_cache.py | 17 ++-- python/sglang/srt/mem_cache/radix_cache.py | 9 +- .../srt/mem_cache/unified_radix_cache.py | 5 +- .../test_fcfs_waiting_prefix_refresh.py | 97 +++++++++++++++---- 5 files changed, 124 insertions(+), 49 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index d76c76ced21c..f666169ee911 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 @@ -131,6 +133,9 @@ def _use_exact_chunk_fill() -> bool: # Inherited from the LPM/HRRN fallback bound in _determine_active_policy; # not a measured optimum. WAITING_PREFIX_REFRESH_MAX_QUEUE = 128 +# A refresh only has to keep waiting prefixes newer than idle ones, so it need not +# run every round; bounds the per-second walk cost under a long queue. +WAITING_PREFIX_REFRESH_INTERVAL_S = 0.5 def _ceil_div(value: int, divisor: int) -> int: @@ -199,6 +204,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() @@ -221,23 +227,13 @@ def calc_priority( if self.tree_cache.supports_fast_match_prefix(): for r in waiting_queue: match_kv_cache(r, self.tree_cache) - # 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 self.tree_cache.supports_prefix_sharing() - ): - 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: 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 ) @@ -281,6 +277,33 @@ def calc_priority( else: raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}") + self._refresh_waiting_prefixes(policy, waiting_queue) + + def _refresh_waiting_prefixes( + self, policy: Policy, waiting_queue: List[Req] + ) -> None: + # Under LRU, a waiting request's prefix ages while it waits and is evicted + # before idle prefixes. Refresh in admission order, head last, so eviction + # takes the prefix needed latest. Cache-aware policies already refresh as + # they match; other eviction strategies rank by other keys, and MRU would + # evict the refreshed prefixes first. + 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_PREFIX_REFRESH_MAX_QUEUE]): + refresh_waiting_prefix(self.tree_cache, r) + def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy: if ( self.policy diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index c0b6d6ae89d3..e8d7874a2549 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -438,16 +438,15 @@ 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. + def refresh_device_prefix(self, key: RadixKey) -> None: + """Touch the cached prefix of ``key`` so LRU eviction sees it as recent. + + The scheduler calls this for waiting requests, at most every + ``WAITING_PREFIX_REFRESH_INTERVAL_S``. It walks the tree only: with a host + tier it also touches host-backed nodes, but it never allocates, loads, or + enqueues lookups on a storage backend or external connector, as + ``match_prefix`` may. Caches without such a walk 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.""" diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6d5031b48bbd..6d4cd4cfd86b 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -412,15 +412,14 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: best_match_node=last_node, ) - def refresh_device_prefix(self, key: RadixKey) -> int: + def refresh_device_prefix(self, key: RadixKey) -> None: key, _ = key.maybe_to_bigram_view(self.is_eagle) if self.disable or len(key) == 0: - return 0 + return 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) + return + self._match_prefix_helper(self.root_node, key) def insert(self, params: InsertParams) -> InsertResult: if self.disable: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 0865cea6db67..64edc3637e08 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -599,13 +599,12 @@ def _match_tree(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: + def refresh_device_prefix(self, key: RadixKey) -> None: # The walk and its actions only; no session shortcut, finalizers or linker. if self.disable: - return 0 + return result = self.tree_core.match_prefix(MatchPrefixParams(key=key)) self._apply_cache_actions(result.cache_actions) - return len(result.device_indices) def supports_prefix_sharing(self) -> bool: return not self.disable diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index b050e36d0b20..1003774df5cb 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -7,8 +7,10 @@ import torch from sglang.srt.environ import envs +from sglang.srt.managers import schedule_policy from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_policy import ( + WAITING_PREFIX_REFRESH_INTERVAL_S, WAITING_PREFIX_REFRESH_MAX_QUEUE, SchedulePolicy, ) @@ -36,10 +38,7 @@ 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.""" - +class _WaitingPrefixFixture(CustomTestCase): def _make_cache(self, params): return RadixCache(params) @@ -85,17 +84,19 @@ def setUp(self): time.sleep(0.005) self.waiting = self._make_req(1, OLDER_PREFIX + [9]) - def _make_policy(self, cache): + def _make_policy(self, cache, policy="fcfs"): return SchedulePolicy( - policy="fcfs", + policy=policy, 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 _make_req(self, rid, tokens, max_new_tokens=128): + return Req( + rid, "", array("q", tokens), SamplingParams(max_new_tokens=max_new_tokens) + ) def _insert(self, tokens, slots): self.cache.insert( @@ -116,14 +117,66 @@ def _assert_older_prefix_survives_one_eviction(self): 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): +class TestFcfsWaitingPrefixRefresh(_WaitingPrefixFixture): + """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 _calc_priority_with_refresh_spy(self, policy, queue): + with unittest.mock.patch.object( + type(self.cache), + "refresh_device_prefix", + autospec=True, + side_effect=type(self.cache).refresh_device_prefix, + ) as refresh: + policy.calc_priority(queue) + return refresh.call_count + + def test_refresh_skipped_when_disabled_or_not_lru(self): + # Every configuration outside FCFS-like policies on LRU must keep main's + # eviction order; MRU would evict the refreshed prefixes first. with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(False): - self.policy.calc_priority([self.waiting]) + self.assertEqual( + self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 0 + ) + for eviction_policy in ("mru", "fifo", "lfu"): + with self.subTest(eviction_policy=eviction_policy): + reset_context() + publish( + ServerArgs( + model_path="dummy", radix_eviction_policy=eviction_policy + ), + role="test", + ) + policy = self._make_policy(self.cache) + self.assertEqual( + self._calc_priority_with_refresh_spy(policy, [self.waiting]), 0 + ) + + def test_refresh_is_throttled(self): + clock = [1000.0] + with unittest.mock.patch.object( + schedule_policy.time, "monotonic", lambda: clock[0] + ): + self.assertEqual( + self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 1 + ) + clock[0] += WAITING_PREFIX_REFRESH_INTERVAL_S / 2 + self.assertEqual( + self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 0 + ) + clock[0] += WAITING_PREFIX_REFRESH_INTERVAL_S + self.assertEqual( + self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 1 + ) + + def test_refresh_follows_the_sorted_order(self): + # LOF admits the longest output first, so after sorting it is the head and + # its prefix must be evicted last, whatever the arrival order. + short_output = self._make_req(2, OLDER_PREFIX + [9], max_new_tokens=8) + long_output = self._make_req(3, NEWER_PREFIX + [9], max_new_tokens=512) + policy = self._make_policy(self.cache, policy="lof") + policy.calc_priority([short_output, long_output]) 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)) @@ -163,6 +216,15 @@ def spy(cache, params): self.assertEqual(calls, []) self._assert_older_prefix_survives_one_eviction() + +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,))) + + +class TestChunkCacheSkipsWaitingPrefixRefresh(_WaitingPrefixFixture): 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. @@ -187,12 +249,5 @@ def test_chunk_cache_skips_the_refresh(self): 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() From 7a52058aca5229d9b729f34948d1cd944bd17c6a Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 8 Oct 2026 16:10:51 -0700 Subject: [PATCH 08/11] [Scheduler] Address review: rename refresh_prefix, trim tests and comments - Restore the fast-match block in calc_priority to main's form. - Rename refresh_device_prefix to refresh_prefix: with a host tier it also touches host-backed nodes. refresh_waiting_prefix takes (req, tree_cache) like match_kv_cache. - Condense the refresh comment; the base docstring no longer names a scheduler constant. - Keep two tests: under LRU the waiting prefixes are evicted in reverse admission order, and under MRU the refresh is skipped. --- python/sglang/srt/managers/schedule_policy.py | 21 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 12 +- python/sglang/srt/mem_cache/common.py | 4 +- python/sglang/srt/mem_cache/radix_cache.py | 2 +- .../srt/mem_cache/unified_radix_cache.py | 5 +- .../test_fcfs_waiting_prefix_refresh.py | 230 ++++-------------- 6 files changed, 64 insertions(+), 210 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index f666169ee911..f29d7fc687fb 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -217,16 +217,15 @@ def calc_priority( ) -> None: policy = self._determine_active_policy(waiting_queue) - # Skip on decode (never prefills). + # Cache-agnostic policies sort without matching; match anyway so the load + # snapshot sees num_matched_prefix_tokens. Decode never prefills. if ( not isinstance(policy, CacheAwarePolicy) + and self.tree_cache.supports_fast_match_prefix() and get_disagg().disaggregation_mode != "decode" ): - # Cache-agnostic policies sort without matching; match anyway so the - # load snapshot sees num_matched_prefix_tokens. - if self.tree_cache.supports_fast_match_prefix(): - for r in waiting_queue: - match_kv_cache(r, self.tree_cache) + for r in waiting_queue: + match_kv_cache(r, self.tree_cache) if self.policy == CacheAgnosticPolicy.FCFS: if self.enable_priority_scheduling: @@ -282,11 +281,9 @@ def calc_priority( def _refresh_waiting_prefixes( self, policy: Policy, waiting_queue: List[Req] ) -> None: - # Under LRU, a waiting request's prefix ages while it waits and is evicted - # before idle prefixes. Refresh in admission order, head last, so eviction - # takes the prefix needed latest. Cache-aware policies already refresh as - # they match; other eviction strategies rank by other keys, and MRU would - # evict the refreshed prefixes first. + # Under LRU a waiting prefix ages and is evicted before idle ones; refresh in + # admission order, head last. Cache-aware policies refresh as they match, and + # other eviction strategies do not rank by recency (MRU would invert it). if ( isinstance(policy, CacheAwarePolicy) or not waiting_queue @@ -302,7 +299,7 @@ def _refresh_waiting_prefixes( return self._last_waiting_prefix_refresh = now for r in reversed(waiting_queue[:WAITING_PREFIX_REFRESH_MAX_QUEUE]): - refresh_waiting_prefix(self.tree_cache, r) + refresh_waiting_prefix(r, self.tree_cache) def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy: if ( diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index e8d7874a2549..bc61652e26f4 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -438,14 +438,14 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def supports_fast_match_prefix(self) -> bool: return False - def refresh_device_prefix(self, key: RadixKey) -> None: + def refresh_prefix(self, key: RadixKey) -> None: """Touch the cached prefix of ``key`` so LRU eviction sees it as recent. - The scheduler calls this for waiting requests, at most every - ``WAITING_PREFIX_REFRESH_INTERVAL_S``. It walks the tree only: with a host - tier it also touches host-backed nodes, but it never allocates, loads, or - enqueues lookups on a storage backend or external connector, as - ``match_prefix`` may. Caches without such a walk leave this a no-op. + The scheduler calls this periodically for waiting requests. It walks the + tree only: with a host tier it also touches host-backed nodes, but it never + allocates, loads, or enqueues lookups on a storage backend or external + connector, as ``match_prefix`` may. Caches without such a walk leave this a + no-op. """ def dfs_weight_order(self, node_handles: Sequence[Any]) -> list[int]: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 8a8e89a9e5b1..1e3af4dc3d37 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -194,8 +194,8 @@ def _req_radix_key( ) -def refresh_waiting_prefix(tree_cache: BasePrefixCache, req: Req) -> None: - tree_cache.refresh_device_prefix( +def refresh_waiting_prefix(req: Req, tree_cache: BasePrefixCache) -> None: + tree_cache.refresh_prefix( _req_radix_key(tree_cache, req, req.origin_input_ids + req.output_ids) ) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6d4cd4cfd86b..5c554c2b54c5 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -412,7 +412,7 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: best_match_node=last_node, ) - def refresh_device_prefix(self, key: RadixKey) -> None: + def refresh_prefix(self, key: RadixKey) -> None: key, _ = key.maybe_to_bigram_view(self.is_eagle) if self.disable or len(key) == 0: return diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 64edc3637e08..14e758a569dd 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -599,8 +599,9 @@ def _match_tree(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) -> None: - # The walk and its actions only; no session shortcut, finalizers or linker. + def refresh_prefix(self, key: RadixKey) -> None: + # The walk and its actions only; no session shortcut, cache-level finalizers + # or linker. if self.disable: return result = self.tree_core.match_prefix(MatchPrefixParams(key=key)) diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 1003774df5cb..9e2bfe5698d3 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -1,19 +1,11 @@ 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 import schedule_policy from sglang.srt.managers.schedule_batch import Req -from sglang.srt.managers.schedule_policy import ( - WAITING_PREFIX_REFRESH_INTERVAL_S, - WAITING_PREFIX_REFRESH_MAX_QUEUE, - SchedulePolicy, -) +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, @@ -21,32 +13,30 @@ 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.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") +register_cpu_ci(est_time=5, suite="base-a-test-cpu") OLDER_PREFIX = [1, 2, 3, 4] NEWER_PREFIX = [5, 6, 7, 8] -class _WaitingPrefixFixture(CustomTestCase): - def _make_cache(self, params): - return RadixCache(params) +class TestFcfsWaitingPrefixRefresh(CustomTestCase): + """Under FCFS with LRU eviction, the prefixes of waiting requests are evicted + in reverse admission order, after prefixes no waiting request needs.""" def setUp(self): reset_context() self.addCleanup(reset_context) - publish(ServerArgs(model_path="dummy"), role="test") torch.set_default_device(None) + + def _run_and_evict_one_prefix(self, eviction_policy): kv_cache = MHATokenToKVPool( size=16, page_size=1, @@ -64,189 +54,55 @@ def setUp(self): 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( + cache = RadixCache( CacheInitParams( disable=False, - req_to_token_pool=req_to_token_pool, + 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="lru", + eviction_policy=eviction_policy, 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, policy="fcfs"): - return SchedulePolicy( - policy=policy, - tree_cache=cache, - enable_hierarchical_cache=False, - enable_priority_scheduling=False, - schedule_low_priority_values_first=False, - ) - - def _make_req(self, rid, tokens, max_new_tokens=128): - return Req( - rid, "", array("q", tokens), SamplingParams(max_new_tokens=max_new_tokens) - ) - - 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) - - -class TestFcfsWaitingPrefixRefresh(_WaitingPrefixFixture): - """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 _calc_priority_with_refresh_spy(self, policy, queue): - with unittest.mock.patch.object( - type(self.cache), - "refresh_device_prefix", - autospec=True, - side_effect=type(self.cache).refresh_device_prefix, - ) as refresh: - policy.calc_priority(queue) - return refresh.call_count - - def test_refresh_skipped_when_disabled_or_not_lru(self): - # Every configuration outside FCFS-like policies on LRU must keep main's - # eviction order; MRU would evict the refreshed prefixes first. - with envs.SGLANG_ENABLE_WAITING_PREFIX_REFRESH.override(False): - self.assertEqual( - self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 0 - ) - for eviction_policy in ("mru", "fifo", "lfu"): - with self.subTest(eviction_policy=eviction_policy): - reset_context() - publish( - ServerArgs( - model_path="dummy", radix_eviction_policy=eviction_policy - ), - role="test", - ) - policy = self._make_policy(self.cache) - self.assertEqual( - self._calc_priority_with_refresh_spy(policy, [self.waiting]), 0 - ) - - def test_refresh_is_throttled(self): - clock = [1000.0] - with unittest.mock.patch.object( - schedule_policy.time, "monotonic", lambda: clock[0] + for tokens, slots in ( + (OLDER_PREFIX, [10, 11, 12, 13]), + (NEWER_PREFIX, [20, 21, 22, 23]), ): - self.assertEqual( - self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 1 - ) - clock[0] += WAITING_PREFIX_REFRESH_INTERVAL_S / 2 - self.assertEqual( - self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 0 + cache.insert( + InsertParams( + key=RadixKey(array("q", tokens)), + value=torch.tensor(slots, dtype=torch.int64), + ) ) - clock[0] += WAITING_PREFIX_REFRESH_INTERVAL_S - self.assertEqual( - self._calc_priority_with_refresh_spy(self.policy, [self.waiting]), 1 + time.sleep(0.005) + # The head waits on the older prefix, the second request on the newer one. + 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 [ + len( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).device_indices ) - - def test_refresh_follows_the_sorted_order(self): - # LOF admits the longest output first, so after sorting it is the head and - # its prefix must be evicted last, whatever the arrival order. - short_output = self._make_req(2, OLDER_PREFIX + [9], max_new_tokens=8) - long_output = self._make_req(3, NEWER_PREFIX + [9], max_new_tokens=512) - policy = self._make_policy(self.cache, policy="lof") - policy.calc_priority([short_output, long_output]) - 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) + for tokens in (OLDER_PREFIX, NEWER_PREFIX) ] - 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() - - -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,))) + def test_lru_evicts_the_prefix_furthest_from_admission(self): + self.assertEqual(self._run_and_evict_one_prefix("lru"), [4, 0]) - -class TestChunkCacheSkipsWaitingPrefixRefresh(_WaitingPrefixFixture): - 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() + 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__": From 8c5f848b38fa4350dca8c018b09f6f7f53c00b8c Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 8 Oct 2026 16:28:27 -0700 Subject: [PATCH 09/11] [Scheduler] Rename the prefix touch API and share the queue bound - refresh_prefix -> touch_prefix (and touch_waiting_prefix, _touch_waiting_prefixes): it only bumps LRU recency, like _touch_node, and does not recompute prefix indices. - The LPM/HRRN fallback and the waiting-prefix touch read one bound, WAITING_QUEUE_PREFIX_MATCH_MAX, instead of two literals of 128. --- python/sglang/srt/managers/schedule_policy.py | 20 +++++++++---------- .../sglang/srt/mem_cache/base_prefix_cache.py | 2 +- python/sglang/srt/mem_cache/common.py | 4 ++-- python/sglang/srt/mem_cache/radix_cache.py | 2 +- .../srt/mem_cache/unified_radix_cache.py | 2 +- 5 files changed, 14 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index f29d7fc687fb..fcad3d433080 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -68,7 +68,7 @@ MatchPrefixParams, zero_match_result, ) -from sglang.srt.mem_cache.common import match_kv_cache, refresh_waiting_prefix +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, @@ -130,9 +130,9 @@ 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 +# Per-round prefix work on the waiting queue stops at this depth: LPM/HRRN fall +# back to FCFS above it, and the waiting-prefix refresh covers only this many. +WAITING_QUEUE_PREFIX_MATCH_MAX = 128 # A refresh only has to keep waiting prefixes newer than idle ones, so it need not # run every round; bounds the per-second walk cost under a long queue. WAITING_PREFIX_REFRESH_INTERVAL_S = 0.5 @@ -276,11 +276,9 @@ def calc_priority( else: raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}") - self._refresh_waiting_prefixes(policy, waiting_queue) + self._touch_waiting_prefixes(policy, waiting_queue) - def _refresh_waiting_prefixes( - self, policy: Policy, waiting_queue: List[Req] - ) -> None: + def _touch_waiting_prefixes(self, policy: Policy, waiting_queue: List[Req]) -> None: # Under LRU a waiting prefix ages and is evicted before idle ones; refresh in # admission order, head last. Cache-aware policies refresh as they match, and # other eviction strategies do not rank by recency (MRU would invert it). @@ -298,8 +296,8 @@ def _refresh_waiting_prefixes( 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_PREFIX_REFRESH_MAX_QUEUE]): - refresh_waiting_prefix(r, self.tree_cache) + 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 ( @@ -308,7 +306,7 @@ def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy: 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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index bc61652e26f4..97f74229f222 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -438,7 +438,7 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: def supports_fast_match_prefix(self) -> bool: return False - def refresh_prefix(self, key: RadixKey) -> None: + def touch_prefix(self, key: RadixKey) -> None: """Touch the cached prefix of ``key`` so LRU eviction sees it as recent. The scheduler calls this periodically for waiting requests. It walks the diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 1e3af4dc3d37..3fee07934ffc 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -194,8 +194,8 @@ def _req_radix_key( ) -def refresh_waiting_prefix(req: Req, tree_cache: BasePrefixCache) -> None: - tree_cache.refresh_prefix( +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) ) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 5c554c2b54c5..9f87ed4a9cd3 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -412,7 +412,7 @@ def match_prefix(self, params: MatchPrefixParams) -> MatchResult: best_match_node=last_node, ) - def refresh_prefix(self, key: RadixKey) -> None: + def touch_prefix(self, key: RadixKey) -> None: key, _ = key.maybe_to_bigram_view(self.is_eagle) if self.disable or len(key) == 0: return diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 14e758a569dd..c522575f4429 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -599,7 +599,7 @@ def _match_tree(self, params: MatchPrefixParams) -> MatchResult: def supports_fast_match_prefix(self) -> bool: return self.tree_core.supports_fast_match_prefix() - def refresh_prefix(self, key: RadixKey) -> None: + def touch_prefix(self, key: RadixKey) -> None: # The walk and its actions only; no session shortcut, cache-level finalizers # or linker. if self.disable: From 893db7d442ab732053076b95e7dbfd2998a6d16c Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 8 Oct 2026 16:36:15 -0700 Subject: [PATCH 10/11] [Scheduler] Trim the comments on the waiting-prefix touch State where the two bounds come from, keep the touch_prefix contract to two lines, and drop test comments that restate the code. --- python/sglang/srt/managers/schedule_policy.py | 11 ++++------- python/sglang/srt/mem_cache/base_prefix_cache.py | 10 ++-------- python/sglang/srt/mem_cache/unified_radix_cache.py | 3 +-- .../mem_cache/test_fcfs_waiting_prefix_refresh.py | 4 ---- 4 files changed, 7 insertions(+), 21 deletions(-) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index fcad3d433080..f06bd6dd393e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -130,11 +130,9 @@ def _use_exact_chunk_fill() -> bool: ) PREFILL_TILE_BUDGET_MODE = "compact" -# Per-round prefix work on the waiting queue stops at this depth: LPM/HRRN fall -# back to FCFS above it, and the waiting-prefix refresh covers only this many. +# Arbitrary; bounds per-round prefix work on the waiting queue. WAITING_QUEUE_PREFIX_MATCH_MAX = 128 -# A refresh only has to keep waiting prefixes newer than idle ones, so it need not -# run every round; bounds the per-second walk cost under a long queue. +# Arbitrary; bounds the walk cost of touching waiting prefixes. WAITING_PREFIX_REFRESH_INTERVAL_S = 0.5 @@ -279,9 +277,8 @@ def calc_priority( self._touch_waiting_prefixes(policy, waiting_queue) def _touch_waiting_prefixes(self, policy: Policy, waiting_queue: List[Req]) -> None: - # Under LRU a waiting prefix ages and is evicted before idle ones; refresh in - # admission order, head last. Cache-aware policies refresh as they match, and - # other eviction strategies do not rank by recency (MRU would invert it). + # 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 diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 97f74229f222..0f44eb3f0a3e 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -439,14 +439,8 @@ def supports_fast_match_prefix(self) -> bool: return False def touch_prefix(self, key: RadixKey) -> None: - """Touch the cached prefix of ``key`` so LRU eviction sees it as recent. - - The scheduler calls this periodically for waiting requests. It walks the - tree only: with a host tier it also touches host-backed nodes, but it never - allocates, loads, or enqueues lookups on a storage backend or external - connector, as ``match_prefix`` may. Caches without such a walk leave this a - no-op. - """ + """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.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index c522575f4429..06fbc60c6be1 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -600,8 +600,7 @@ def supports_fast_match_prefix(self) -> bool: return self.tree_core.supports_fast_match_prefix() def touch_prefix(self, key: RadixKey) -> None: - # The walk and its actions only; no session shortcut, cache-level finalizers - # or linker. + # 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)) diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 9e2bfe5698d3..2034d84137d0 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -28,9 +28,6 @@ class TestFcfsWaitingPrefixRefresh(CustomTestCase): - """Under FCFS with LRU eviction, the prefixes of waiting requests are evicted - in reverse admission order, after prefixes no waiting request needs.""" - def setUp(self): reset_context() self.addCleanup(reset_context) @@ -77,7 +74,6 @@ def _run_and_evict_one_prefix(self, eviction_policy): ) ) time.sleep(0.005) - # The head waits on the older prefix, the second request on the newer one. queue = [ Req(rid, "", array("q", tokens + [9]), SamplingParams()) for rid, tokens in ((1, OLDER_PREFIX), (2, NEWER_PREFIX)) From c42da2b1f675800039802688b30322ab7fe4e8f4 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 8 Oct 2026 16:50:30 -0700 Subject: [PATCH 11/11] [Scheduler] Read device_prefix_len in the waiting-prefix touch test main renamed MatchResult.device_indices to device_prefix_len. --- .../unit/mem_cache/test_fcfs_waiting_prefix_refresh.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py index 2034d84137d0..12a6e47c9bfa 100644 --- a/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py +++ b/test/registered/unit/mem_cache/test_fcfs_waiting_prefix_refresh.py @@ -85,11 +85,9 @@ def _run_and_evict_one_prefix(self, eviction_policy): SchedulePolicy("fcfs", cache, False, False, False).calc_priority(queue) cache.evict(EvictParams(num_tokens=len(OLDER_PREFIX))) return [ - len( - cache.match_prefix( - MatchPrefixParams(key=RadixKey(array("q", tokens))) - ).device_indices - ) + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).device_prefix_len for tokens in (OLDER_PREFIX, NEWER_PREFIX) ]