chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Auto Sync] Update schedule_batch.py, schedule_policy.py, b... (20251122) (#13763)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
Co-authored-by: Hanming Lu <hanming@x.ai>
This commit is contained in:
Lianmin Zheng
2025-11-24 14:33:31 -08:00
committed by GitHub
co-authored by github-actions[bot] <github-actions[bot]@users.noreply.github.com> Hanming Lu Hanming Lu
parent 9dc15d8569
commit e83bd1fadc
7 changed files with 51 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_policy.py
+17 -9
View File
@@ -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