From e83bd1fadcb6fb860a5fcac61a754f4e946ada17 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 24 Nov 2025 14:33:31 -0800 Subject: [PATCH] [Auto Sync] Update schedule_batch.py, schedule_policy.py, b... (20251122) (#13763) Co-authored-by: github-actions[bot] Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Co-authored-by: Hanming Lu --- python/sglang/srt/managers/schedule_batch.py | 18 ++++++++----- python/sglang/srt/managers/schedule_policy.py | 26 ++++++++++++------- .../sglang/srt/mem_cache/base_prefix_cache.py | 4 +++ .../sglang/srt/mem_cache/mamba_radix_cache.py | 6 ++++- python/sglang/srt/mem_cache/radix_cache.py | 6 ++++- .../storage/lmcache/lmc_radix_cache.py | 3 ++- .../sglang/srt/mem_cache/swa_radix_cache.py | 7 ++++- 7 files changed, 51 insertions(+), 19 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 6a7f44a31..45451c92e 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -762,12 +762,7 @@ class Req: token_ids = self.fill_ids[:max_prefix_len] if tree_cache is not None: - ( - self.prefix_indices, - self.last_node, - self.last_host_node, - self.host_hit_length, - ) = tree_cache.match_prefix( + match_result = tree_cache.match_prefix( key=RadixKey(token_ids=token_ids, extra_key=self.extra_key), **( {"req": self, "cow_mamba": True} @@ -775,6 +770,17 @@ class Req: else {} ), ) + ( + self.prefix_indices, + self.last_node, + self.last_host_node, + self.host_hit_length, + ) = ( + match_result.device_indices, + match_result.last_device_node, + match_result.last_host_node, + match_result.host_hit_length, + ) self.cache_protected_len = len(self.prefix_indices) self.extend_input_len = len(self.fill_ids) - len(self.prefix_indices) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index aee0d1fc2..3c95e72e9 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -175,10 +175,19 @@ class SchedulePolicy: extra_key = r.extra_key # NOTE: the prefix_indices must always be aligned with last_node - r.prefix_indices, r.last_node, r.last_host_node, r.host_hit_length = ( - self.tree_cache.match_prefix( - rid=r.rid, key=RadixKey(token_ids=prefix_ids, extra_key=extra_key) - ) + match_result = self.tree_cache.match_prefix( + rid=r.rid, key=RadixKey(token_ids=prefix_ids, extra_key=extra_key) + ) + ( + r.prefix_indices, + r.last_node, + r.last_host_node, + r.host_hit_length, + ) = ( + match_result.device_indices, + match_result.last_device_node, + match_result.last_host_node, + match_result.host_hit_length, ) # NOTE(sang): This logic is for in-batch prefix caching; @@ -189,12 +198,11 @@ class SchedulePolicy: # threshold means we cannot use in-batch prefix caching for short prefixes. # It is kind of common when the engine is long running (e.g., imagine the prefix "the"). if len(r.prefix_indices) <= IN_BATCH_PREFIX_CACHING_CHECK_THRESHOLD: - in_batch_matching_prefixes, _, _, _ = ( - self.waiting_queue_radix_tree.match_prefix( - rid=r.rid, - key=RadixKey(token_ids=prefix_ids, extra_key=extra_key), - ) + match_result = self.waiting_queue_radix_tree.match_prefix( + rid=r.rid, + key=RadixKey(token_ids=prefix_ids, extra_key=extra_key), ) + in_batch_matching_prefixes = match_result.device_indices if ( len(in_batch_matching_prefixes) >= IN_BATCH_PREFIX_CACHING_DEPRIORITIZE_THRESHOLD diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 9c16b639b..826bfa6c5 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -41,12 +41,16 @@ class MatchResult(NamedTuple): this **must** be the same as `last_device_node`. host_hit_length : Length of the KV cache hit on the host, if applicable. 0 if HiCache is not enabled. + mamba_branching_seqlen: The mamba radix cache branching point, which is the longest + page-aligned position that could've been cache hit if there + exists a mamba state. """ device_indices: torch.Tensor last_device_node: Any last_host_node: Any host_hit_length: int = 0 + mamba_branching_seqlen: Optional[int] = None class BasePrefixCache(ABC, PrefixCacheTrait): diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 31a37f86d..239f37bdc 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -515,9 +515,13 @@ class MambaRadixCache(BasePrefixCache): self.req_to_token_pool.mamba_pool.free(mamba_value_forked) # The prefix indices could be updated, reuse it - new_indices, new_last_node, _, _ = self.match_prefix( + match_result = self.match_prefix( RadixKey(page_aligned_token_ids, req.extra_key) ) + (new_indices, new_last_node) = ( + match_result.device_indices, + match_result.last_device_node, + ) if not mamba_exist: assert torch.equal(new_last_node.mamba_value, mamba_value_forked) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index fae975de8..0ac534165 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -436,7 +436,11 @@ class RadixCache(BasePrefixCache): ) # The prefix indices could be updated, reuse it - new_indices, new_last_node, _, _ = self.match_prefix(radix_key) + match_result = self.match_prefix(radix_key) + (new_indices, new_last_node) = ( + match_result.device_indices, + match_result.last_device_node, + ) assert len(new_indices) == len(keys), f"{len(new_indices)=}, {len(keys)=}" self.req_to_token_pool.write( diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index c697ebc79..be23e7818 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -217,7 +217,8 @@ class LMCRadixCache(RadixCache): req.req_pool_idx, :kv_committed_len ] - _, new_last_node, _, _ = self.match_prefix(RadixKey(token_ids, req.extra_key)) + match_result = self.match_prefix(RadixKey(token_ids, req.extra_key)) + new_last_node = match_result.last_device_node assert new_last_node is not None self.inc_lock_ref(new_last_node) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 33f55c1bb..037ea4ec3 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -541,9 +541,14 @@ class SWARadixCache(BasePrefixCache): ) # The prefix indices could be updated, reuse it - new_indices, new_last_node, _, _ = self.match_prefix( + match_result = self.match_prefix( RadixKey(page_aligned_token_ids, req.extra_key) ) + (new_indices, new_last_node) = ( + match_result.device_indices, + match_result.last_device_node, + ) + assert old_prefix_len <= len( new_indices ), f"{req.prefix_indices=}, {new_indices=}"