minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Scheduler] Add shortest-prefill-first scheduling (#40024)

Browse Source
This commit is contained in:
Yuwei An
2026-09-17 21:18:58 -07:00
committed by GitHub
parent c055dc6ff6
commit 65ef55e2a8
6 changed files with 281 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_prefill_adder.py
+94
View File
@@ -9,6 +9,7 @@ from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_policy import (
AddReqResult,
PrefillAdder,
SchedulePolicy,
estimate_prefill_extend_tile_metrics,
)
from sglang.srt.mem_cache.base_prefix_cache import (
@@ -21,6 +22,7 @@ from sglang.srt.mem_cache.prefill_budget import (
SWAPrefillBudget,
estimate_swa_kv_tokens,
)
from sglang.srt.mem_cache.radix_cache import RadixCache
from sglang.srt.mem_cache.unified_memory_pool import init_unified_swa_pools
from sglang.srt.runtime_context import get_context
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
@@ -186,6 +188,98 @@ class TestPrefillAdder(CustomTestCase):
)
return req
def create_shortest_prefill_adder(self, *, chunk_tokens=4096):
override = get_context().override_server_args(
schedule_policy="shortest-prefill-first"
)
override.install()
self.addCleanup(override.restore)
self.mock_tree_cache.supports_mamba.return_value = False
self.mock_tree_cache.is_tree_cache.return_value = False
self.mock_token_allocator.available_size.return_value = 32768
return self.create_adder(
self.create_running_batch(), page_size=256, rem_chunk_tokens=chunk_tokens
)
def test_shortest_prefill_reserves_space_for_complete_waiting_requests(self):
adder = self.create_shortest_prefill_adder()
policy = SchedulePolicy(
policy="shortest-prefill-first",
tree_cache=RadixCache.create_simulated(),
enable_hierarchical_cache=True,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
continuation = self.create_shared_req("continuation")
continuation.full_untruncated_fill_ids = list(range(16384))
waiting = [self.create_shared_req("a"), self.create_shared_req("b")]
for req, length in zip(waiting, [512, 1024]):
req.origin_input_ids = list(range(length))
req.full_untruncated_fill_ids = list(range(length))
req.num_matched_prefix_tokens = 0
adder.chunked_req_limit = policy.shortest_prefill_chunk_limit(
continuation, waiting, adder.rem_chunk_tokens, adder.page_size
)
self.assertIs(adder.add_chunked_req(continuation), continuation)
self.assertEqual(continuation.extend_range.length, 2560)
for req in waiting:
adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None)
self.assertEqual(adder.can_run_list, [continuation, *waiting])
self.assertIsNone(adder.new_chunked_req)
self.assertEqual(adder.rem_chunk_tokens, 0)
self.assertGreaterEqual(adder.rem_total_tokens, 0)
def test_shortest_prefill_rejects_second_unfinished_chunk(self):
adder = self.create_shortest_prefill_adder(chunk_tokens=512)
req = self.create_shared_req("second-chunk")
req.full_untruncated_fill_ids = list(range(1024))
self.assertEqual(
adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
AddReqResult.OTHER,
)
self.assertEqual(adder.can_run_list, [])
self.assertIsNone(adder.new_chunked_req)
req.set_extend_range.assert_not_called()
self.mock_tree_cache.init_load_back.assert_not_called()
def test_shortest_prefill_rechecks_chunk_limit_after_host_miss(self):
adder = self.create_shortest_prefill_adder(chunk_tokens=512)
req = self.create_shared_req("host-miss")
req.full_untruncated_fill_ids = list(range(1024))
req.prefix_indices = torch.empty(0, dtype=torch.int64)
req.host_hit_length = 768
req.best_match_node = req.last_node
req.needs_host_load_back.return_value = True
self.mock_tree_cache.init_load_back.return_value = (
torch.empty(0, dtype=torch.int64),
req.last_node,
)
self.assertEqual(
adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
AddReqResult.OTHER,
)
self.mock_tree_cache.init_load_back.assert_called_once()
self.assertEqual(adder.can_run_list, [])
req.set_extend_range.assert_not_called()
def test_shortest_prefill_preserves_memory_admission(self):
adder = self.create_shortest_prefill_adder()
self.mock_token_allocator.available_size.return_value = 256
req = self.create_shared_req("no-memory")
req.full_untruncated_fill_ids = list(range(512))
self.assertEqual(
adder.add_one_req(req, has_chunked_req=True, truncation_align_size=None),
AddReqResult.NO_TOKEN,
)
self.assertEqual(adder.can_run_list, [])
def test_continuation_without_limit_keeps_normal_chunk_size(self):
adder = self.create_shortest_prefill_adder()
req = self.create_shared_req("continuation")
req.full_untruncated_fill_ids = list(range(8192))
self.assertIs(adder.add_chunked_req(req), req)
self.assertEqual(req.extend_range.length, 4096)
def test_shared_admission_reserves_all_pending_requests(self):
adder = self.create_shared_adder()
first, second = (
test/registered/unit/managers/test_schedule_policy.py
+104 -1
View File
@@ -1,8 +1,9 @@
import unittest
from array import array
from unittest.mock import patch
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_policy import SchedulePolicy
from sglang.srt.managers.schedule_policy import CacheAwarePolicy, SchedulePolicy
from sglang.srt.mem_cache.radix_cache import RadixCache
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.test.ci.ci_register import register_cpu_ci
@@ -124,5 +125,107 @@ class TestSchedulePolicyHRRN(CustomTestCase):
self.assertEqual(waiting_queue[2].rid, "c")
class TestShortestPrefillFirst(CustomTestCase):
def setUp(self):
self.policy = SchedulePolicy(
policy="shortest-prefill-first",
tree_cache=RadixCache.create_simulated(),
enable_hierarchical_cache=True,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
def make_req(self, rid, uncached, *, cached=0, arrived=0):
req = _make_req(rid, "", list(range(uncached + cached)))
req.full_untruncated_fill_ids = req.origin_input_ids[:]
req.num_matched_prefix_tokens = cached
req.prefix_indices = list(range(cached))
req.time_stats.wait_queue_entry_time = arrived
return req
def test_calc_priority_uses_uncached_work(self):
cached = self.make_req("cached", 16, cached=4096)
short = self.make_req("short", 32)
long = self.make_req("long", 1024)
queue = [long, short, cached]
with patch.object(self.policy, "_compute_prefix_matches", return_value=set()):
self.policy.calc_priority(queue)
self.assertEqual([req.rid for req in queue], ["cached", "short", "long"])
def test_equal_work_uses_arrival_time(self):
older = self.make_req("z", 32, arrived=1)
newer = self.make_req("a", 32, arrived=2)
queue = [newer, older]
self.policy._sort_by_shortest_prefill(queue, set())
self.assertEqual(queue, [older, newer])
def test_duplicate_prefix_is_deprioritized(self):
duplicate = self.make_req("duplicate", 1)
other = self.make_req("other", 1024)
queue = [duplicate, other]
with patch.object(
self.policy, "_compute_prefix_matches", return_value={duplicate.rid}
):
self.policy.calc_priority(queue)
self.assertEqual(queue, [other, duplicate])
def test_retracted_output_is_part_of_uncached_work(self):
replay = self.make_req("replay", 16, cached=1024)
replay.output_ids.extend([0] * 64)
short = self.make_req("short", 32)
queue = [replay, short]
self.policy._sort_by_shortest_prefill(queue, set())
self.assertEqual(queue, [short, replay])
def test_chunk_limit_reserves_complete_short_prefills(self):
continuation = self.make_req("continuation", 16384)
waiting = [self.make_req("a", 512), self.make_req("b", 1024)]
self.assertEqual(
self.policy.shortest_prefill_chunk_limit(continuation, waiting, 4096, 256),
2560,
)
def test_reservation_rounds_to_pages_and_keeps_continuation_progress(self):
continuation = self.make_req("continuation", 16384)
self.assertEqual(
self.policy.shortest_prefill_chunk_limit(
continuation, [self.make_req("short", 257)], 4096, 256
),
3584,
)
self.assertEqual(
self.policy.shortest_prefill_chunk_limit(
continuation, [self.make_req("short", 3840)], 4096, 256
),
256,
)
def test_no_reservation_when_request_cannot_fit_or_is_not_shorter(self):
continuation = self.make_req("continuation", 8192)
for waiting, budget in [
([], 4096),
([self.make_req("same", 8192)], 4096),
([self.make_req("too-large", 4096)], 4096),
([self.make_req("short", 1)], 256),
]:
with self.subTest(budget=budget, waiting=[req.rid for req in waiting]):
self.assertIsNone(
self.policy.shortest_prefill_chunk_limit(
continuation, waiting, budget, 256
)
)
def test_other_policy_keeps_normal_chunk_limit(self):
self.policy.policy = CacheAwarePolicy.HRRN
self.assertIsNone(
self.policy.shortest_prefill_chunk_limit(
self.make_req("continuation", 8192),
[self.make_req("short", 512)],
4096,
256,
)
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/test_server_args_cli_metadata.py
+1
View File
@@ -61,6 +61,7 @@ class TestServerArgsMigratedCliMetadata(CustomTestCase):
"priority",
"routing-key",
"hrrn",
"shortest-prefill-first",
],
)
self.assertEqual(