env: add SGLANG_RADIX_FORCE_MISS to force radix prefix-cache miss (#24726)
Co-authored-by: sihan-zzz <228612289+sihan-zzz@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
sihan-zzz
Claude Opus 4.7
parent
560829a171
commit
d1c5937428
@@ -254,6 +254,7 @@ class Envs:
|
|||||||
SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False)
|
SGLANG_DISABLE_CONSECUTIVE_PREFILL_OVERLAP = EnvBool(False)
|
||||||
SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1)
|
SGLANG_SCHEDULER_MAX_RECV_PER_POLL = EnvInt(-1)
|
||||||
SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
|
SGLANG_EXPERIMENTAL_CPP_RADIX_TREE = EnvBool(False)
|
||||||
|
SGLANG_RADIX_FORCE_MISS = EnvBool(False)
|
||||||
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
|
SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR = EnvFloat(0.75)
|
||||||
SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False)
|
SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False)
|
||||||
SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False)
|
SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False)
|
||||||
|
|||||||
@@ -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.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
|
||||||
from sglang.srt.managers.embed_types import PositionalEmbeds
|
from sglang.srt.managers.embed_types import PositionalEmbeds
|
||||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
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 (
|
from sglang.srt.mem_cache.common import (
|
||||||
alloc_for_decode,
|
alloc_for_decode,
|
||||||
alloc_for_extend,
|
alloc_for_extend,
|
||||||
@@ -1029,6 +1033,8 @@ class Req(ReqDllmMixin):
|
|||||||
cow_mamba=cow_mamba,
|
cow_mamba=cow_mamba,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
if envs.SGLANG_RADIX_FORCE_MISS.get():
|
||||||
|
match_result = zero_match_result(tree_cache, match_result)
|
||||||
(
|
(
|
||||||
self.prefix_indices,
|
self.prefix_indices,
|
||||||
self.last_node,
|
self.last_node,
|
||||||
|
|||||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
|||||||
|
|
||||||
import logging
|
import logging
|
||||||
|
|
||||||
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
|
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
|
||||||
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
|
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
|
||||||
from sglang.srt.utils import get_bool_env_var
|
from sglang.srt.utils import get_bool_env_var
|
||||||
@@ -42,6 +43,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
|
|||||||
InitLoadBackParams,
|
InitLoadBackParams,
|
||||||
InsertParams,
|
InsertParams,
|
||||||
MatchPrefixParams,
|
MatchPrefixParams,
|
||||||
|
zero_match_result,
|
||||||
)
|
)
|
||||||
from sglang.srt.mem_cache.hisparse_memory_pool import (
|
from sglang.srt.mem_cache.hisparse_memory_pool import (
|
||||||
DeepSeekV4HiSparseTokenToKVPoolAllocator,
|
DeepSeekV4HiSparseTokenToKVPoolAllocator,
|
||||||
@@ -98,6 +100,8 @@ def match_prefix_for_req(
|
|||||||
req=req if include_req else None,
|
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.prefix_indices,
|
||||||
req.last_node,
|
req.last_node,
|
||||||
@@ -249,6 +253,10 @@ class SchedulePolicy:
|
|||||||
key=RadixKey(token_ids=prefix_ids, extra_key=extra_key)
|
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
|
in_batch_matching_prefixes = match_result.device_indices
|
||||||
if (
|
if (
|
||||||
len(in_batch_matching_prefixes)
|
len(in_batch_matching_prefixes)
|
||||||
|
|||||||
@@ -151,6 +151,24 @@ class MatchResult(NamedTuple):
|
|||||||
cache_protected_len: Optional[int] = None
|
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):
|
class BasePrefixCache(ABC, PrefixCacheTrait):
|
||||||
"""Cache can be indexed by either rid or key."""
|
"""Cache can be indexed by either rid or key."""
|
||||||
|
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user