diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c87321c35..97e8040e8 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -254,6 +254,7 @@ class Envs: SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False) SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1) SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False) + SGLANG_RADIX_FORCE_MISS = EnvBool(False) SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75) SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False) SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 07a1c867d..2450d49fb 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -61,7 +61,11 @@ from sglang.srt.environ import envs from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchPrefixParams +from sglang.srt.mem_cache.base_prefix_cache import ( + BasePrefixCache, + MatchPrefixParams, + zero_match_result, +) from sglang.srt.mem_cache.common import ( alloc_for_decode, alloc_for_extend, @@ -1029,6 +1033,8 @@ class Req(ReqDllmMixin): cow_mamba=cow_mamba, ) ) + if envs.SGLANG_RADIX_FORCE_MISS.get(): + match_result = zero_match_result(tree_cache, match_result) ( self.prefix_indices, self.last_node, diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 84ec14c36..257d30293 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -2,6 +2,7 @@ from __future__ import annotations import logging +from sglang.srt.environ import envs from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams from sglang.srt.utils import get_bool_env_var @@ -42,6 +43,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( InitLoadBackParams, InsertParams, MatchPrefixParams, + zero_match_result, ) from sglang.srt.mem_cache.hisparse_memory_pool import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, @@ -98,6 +100,8 @@ def match_prefix_for_req( req=req if include_req else None, ) ) + if envs.SGLANG_RADIX_FORCE_MISS.get(): + match_result = zero_match_result(tree_cache, match_result) ( req.prefix_indices, req.last_node, @@ -249,6 +253,10 @@ class SchedulePolicy: key=RadixKey(token_ids=prefix_ids, extra_key=extra_key) ) ) + if envs.SGLANG_RADIX_FORCE_MISS.get(): + match_result = zero_match_result( + self.waiting_queue_radix_tree, match_result + ) in_batch_matching_prefixes = match_result.device_indices if ( len(in_batch_matching_prefixes) diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 4316697ef..28670c5a4 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -151,6 +151,24 @@ class MatchResult(NamedTuple): cache_protected_len: Optional[int] = None +def zero_match_result(tree_cache, match_result: "MatchResult") -> "MatchResult": + root = getattr(tree_cache, "root_node", None) + if root is None: + raise RuntimeError( + f"SGLANG_RADIX_FORCE_MISS is not supported by {type(tree_cache).__name__} " + "(no `root_node` attribute). Disable the flag or use a cache backend " + "that exposes a tree root." + ) + return match_result._replace( + # [:0] keeps dtype and device of the original tensor (e.g. CUDA int64) + # without allocating a fresh empty tensor. + device_indices=match_result.device_indices[:0], + last_device_node=root, + last_host_node=root, + host_hit_length=0, + ) + + class BasePrefixCache(ABC, PrefixCacheTrait): """Cache can be indexed by either rid or key.""" diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py new file mode 100644 index 000000000..d9a72c89d --- /dev/null +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -0,0 +1,99 @@ +"""Unit tests for SGLANG_RADIX_FORCE_MISS. + +The flag is gated at the scheduler boundary, so we test the helper directly +plus an end-to-end check of `match_prefix_for_req` driving a populated +RadixCache. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="stage-a-test-cpu") + +import unittest +import unittest.mock + +import torch + +from sglang.srt.environ import envs +from sglang.srt.managers.schedule_policy import match_prefix_for_req +from sglang.srt.mem_cache.base_prefix_cache import ( + InsertParams, + MatchPrefixParams, + MatchResult, + zero_match_result, +) +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey + + +class _StubReq: + def __init__(self, token_ids): + self.origin_input_ids = list(token_ids) + self.output_ids = [] + self.extra_key = None + self.prefix_indices = None + self.last_node = None + self.last_host_node = None + self.host_hit_length = None + self.mamba_branching_seqlen = None + self.cache_protected_len = None + + +class TestZeroMatchResult(unittest.TestCase): + def test_zero_replaces_indices_and_nodes(self): + tree = RadixCache.create_simulated() + tree.insert(InsertParams(key=RadixKey(token_ids=[1, 2, 3, 4, 5]))) + match = tree.match_prefix( + MatchPrefixParams(key=RadixKey(token_ids=[1, 2, 3, 9])) + ) + self.assertGreater(len(match.device_indices), 0) + zeroed = zero_match_result(tree, match) + self.assertEqual(int(zeroed.device_indices.numel()), 0) + self.assertIs(zeroed.last_device_node, tree.root_node) + self.assertIs(zeroed.last_host_node, tree.root_node) + self.assertEqual(zeroed.host_hit_length, 0) + # dtype/device preserved (slice-not-allocate). + self.assertEqual(zeroed.device_indices.dtype, match.device_indices.dtype) + self.assertEqual(zeroed.device_indices.device, match.device_indices.device) + + def test_no_root_node_raises(self): + # tree_cache without a root_node: must raise loudly rather than silently + # leak cache hits past the gate. + class _NoRoot: + pass + + original = MatchResult( + device_indices=torch.tensor([7, 8, 9], dtype=torch.int64), + last_device_node="sentinel-device", + last_host_node="sentinel-host", + host_hit_length=4, + ) + with self.assertRaisesRegex(RuntimeError, "SGLANG_RADIX_FORCE_MISS"): + zero_match_result(_NoRoot(), original) + + +class TestMatchPrefixForReqForceMiss(unittest.TestCase): + def test_force_miss_zeros_req_prefix(self): + tree = RadixCache.create_simulated() + tree.insert( + InsertParams(key=RadixKey(token_ids=[10, 11, 12, 13, 14, 15, 16, 17])) + ) + + # Sanity: without the flag, the same lookup hits. + baseline_req = _StubReq([10, 11, 12, 13, 99, 100]) + with envs.SGLANG_RADIX_FORCE_MISS.override(False): + match_prefix_for_req(tree, baseline_req) + self.assertGreater(int(baseline_req.prefix_indices.numel()), 0) + self.assertIsNot(baseline_req.last_node, tree.root_node) + + # With the flag, the same lookup is forced to miss. + forced_req = _StubReq([10, 11, 12, 13, 99, 100]) + with envs.SGLANG_RADIX_FORCE_MISS.override(True): + match_prefix_for_req(tree, forced_req) + self.assertEqual(int(forced_req.prefix_indices.numel()), 0) + self.assertIs(forced_req.last_node, tree.root_node) + self.assertIs(forced_req.last_host_node, tree.root_node) + self.assertEqual(forced_req.host_hit_length, 0) + + +if __name__ == "__main__": + unittest.main()