From 8eeffbe9aa0fcab23f2567ad4d3aadfaba7ce29c Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 9 Jan 2026 23:07:18 +0800 Subject: [PATCH] Support token low usage watermark in prefill delayer (#16814) --- python/sglang/srt/environ.py | 1 + python/sglang/srt/managers/prefill_delayer.py | 61 +++++- python/sglang/srt/managers/scheduler.py | 6 +- python/sglang/srt/metrics/collector.py | 4 +- test/srt/test_prefill_delayer.py | 191 +++++++++++++++++- 5 files changed, 241 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 445035643..84ca37b83 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -224,6 +224,7 @@ class Envs: SGLANG_SCHEDULER_SKIP_ALL_GATHER = EnvBool(False) SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE = EnvBool(False) SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES = EnvInt(30) + SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK = EnvFloat(None) SGLANG_DATA_PARALLEL_BUDGET_INTERVAL = EnvInt(1) # Test: pd-disaggregation diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index 967784f9d..8df34fe8e 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -31,6 +31,7 @@ class _NegotiateOutput(NamedTuple): output_allow: bool output_reason: str num_prefillable: int + num_token_watermark_force_allow: int class PrefillDelayer: @@ -41,15 +42,19 @@ class PrefillDelayer: cpu_group, server_args, max_delay_passes: int, + token_usage_low_watermark: Optional[float], metrics_collector: Optional["SchedulerMetricsCollector"] = None, ): self._max_delay_passes = max_delay_passes + self._token_usage_low_watermark = token_usage_low_watermark logger.info( - f"PrefillDelayer initialized with max_delay_passes={self._max_delay_passes}" + f"PrefillDelayer initialized with " + f"max_delay_passes={self._max_delay_passes} " + f"token_usage_low_watermark={self._token_usage_low_watermark}" ) self._global_info_buffer = torch.empty( - (dp_size, attn_tp_size, 1), + (dp_size, attn_tp_size, 2), dtype=torch.int64, device="cpu", ) @@ -70,33 +75,53 @@ class PrefillDelayer: ), "To use PrefillDelayer, disable_overlap_schedule must be False." def _negotiate_should_allow_prefill( - self, local_prefillable: bool + self, local_prefillable: bool, token_usage: float ) -> _NegotiateOutput: out = self._negotiate_should_allow_prefill_pure( prev_state=self._curr_state, local_prefillable=local_prefillable, + token_usage=token_usage, ) self._curr_state = out.next_state return out + # (Almost) pure function, do not modify self state def _negotiate_should_allow_prefill_pure( self, prev_state: Optional[_State], local_prefillable: bool, + token_usage: float, ) -> _NegotiateOutput: - global_prefillable = self._gather_info(local_prefillable=local_prefillable) + # Compute local states + local_token_watermark_force_allow = ( + local_prefillable + and ((x := self._token_usage_low_watermark) is not None) + and (token_usage < x) + ) + # Gather global states + global_prefillable, global_token_watermark_force_allow = self._gather_info( + local_prefillable=local_prefillable, + local_token_watermark_force_allow=local_token_watermark_force_allow, + ) + + # Compute derived global states if global_prefillable.min().item() > 0: prefillable_status = "all" elif global_prefillable.max().item() == 0: prefillable_status = "none" else: prefillable_status = "mixed" + global_exists_token_watermark_force_allow = ( + global_token_watermark_force_allow.max().item() > 0 + ) debug_info = dict( input_estimation=prefillable_status, num_prefillable=global_prefillable.sum().item(), + num_token_watermark_force_allow=global_token_watermark_force_allow.sum().item(), ) + # Compute outputs if prefillable_status == "all": exist_previous_wait = prev_state is not None return _NegotiateOutput( @@ -108,11 +133,20 @@ class PrefillDelayer: elif prefillable_status == "none": return _NegotiateOutput( next_state=None, + # It does not matter whether we allow or not, thus we allow for simplicity output_allow=True, output_reason="", **debug_info, ) elif prefillable_status == "mixed": + if global_exists_token_watermark_force_allow: + return _NegotiateOutput( + next_state=None, + output_allow=True, + output_reason="token_watermark", + **debug_info, + ) + prev_delayed_count = prev_state.delayed_count if prev_state else 0 if prev_delayed_count < self._max_delay_passes - 1: next_state = prev_state or _State() @@ -133,9 +167,11 @@ class PrefillDelayer: else: raise NotImplementedError - def _gather_info(self, local_prefillable: bool): + def _gather_info( + self, local_prefillable: bool, local_token_watermark_force_allow: bool + ): local_info = torch.tensor( - [int(local_prefillable)], + [int(local_prefillable), int(local_token_watermark_force_allow)], device="cpu", dtype=torch.int64, ) @@ -145,12 +181,13 @@ class PrefillDelayer: group=self._cpu_group, ) tp0_info = self._global_info_buffer[:, 0, :] - return tp0_info[:, 0] + return tp0_info[:, 0], tp0_info[:, 1] class PrefillDelayerSinglePassExecutor: - def __init__(self, prefill_delayer: PrefillDelayer): + def __init__(self, prefill_delayer: PrefillDelayer, token_usage: float): self._prefill_delayer = prefill_delayer + self._token_usage = token_usage self._result: Optional[_NegotiateOutput] = None @property @@ -171,6 +208,7 @@ class PrefillDelayerSinglePassExecutor: if not self._called: self._result = self._prefill_delayer._negotiate_should_allow_prefill( local_prefillable=local_prefillable, + token_usage=self._token_usage, ) return self._result.output_allow @@ -187,6 +225,13 @@ def _record_single_pass_result( f"(num_prefillable={output.num_prefillable}, " f"actual_execution={actual_execution})" ) + elif output.output_allow and (output.output_reason == "token_watermark"): + logger.info( + f"PrefillDelayer force allow prefill due to low watermark. " + f"(num_prefillable={output.num_prefillable}, " + f"num_token_watermark_force_allow={output.num_token_watermark_force_allow}, " + f"actual_execution={actual_execution})" + ) else: assert output.output_reason in { "", diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index fd720c5c8..77263a738 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -775,6 +775,9 @@ class Scheduler( self.metrics_collector if self.enable_metrics else None ), max_delay_passes=envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.get(), + token_usage_low_watermark=( + envs.SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK.get() + ), ) # Enable preemption for priority scheduling. self.try_preemption = self.enable_priority_scheduling @@ -1846,8 +1849,9 @@ class Scheduler( def get_new_batch_prefill(self) -> Optional[ScheduleBatch]: prefill_delayer_single_pass = None if self.prefill_delayer: + _, token_usage, _, _ = self._get_token_info() prefill_delayer_single_pass = PrefillDelayerSinglePassExecutor( - self.prefill_delayer + self.prefill_delayer, token_usage=token_usage ) ret = self._get_new_batch_prefill_raw( diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py index 9b9687f65..47d3df45c 100644 --- a/python/sglang/srt/metrics/collector.py +++ b/python/sglang/srt/metrics/collector.py @@ -826,9 +826,9 @@ class SchedulerMetricsCollector: self.prefill_delayer_outcomes_total.labels( **self.labels, input_estimation=input_estimation, - output_allow=str(int(output_allow)), + output_allow=str(output_allow).lower(), output_reason=output_reason, - actual_execution=str(int(actual_execution)), + actual_execution=str(actual_execution).lower(), ).inc(1) def increment_retracted_reqs( diff --git a/test/srt/test_prefill_delayer.py b/test/srt/test_prefill_delayer.py index 7e1062a95..a15854401 100644 --- a/test/srt/test_prefill_delayer.py +++ b/test/srt/test_prefill_delayer.py @@ -1,11 +1,14 @@ +import asyncio import os import re +import time import unittest from collections import defaultdict from dataclasses import dataclass from types import SimpleNamespace -from typing import List +from typing import List, Optional +import openai import requests import torch import torch.multiprocessing as mp @@ -32,12 +35,14 @@ WORLD_SIZE = os.environ.get("SGLANG_TEST_WORLD_SIZE", "8") @dataclass class NegotiateCall: prefillable: List[bool] + token_usage: List[float] @dataclass class NegotiateTestCase: name: str max_delay_passes: int + token_usage_low_watermark: Optional[float] calls: List[NegotiateCall] expected_allow: bool expected_reason: str @@ -63,11 +68,13 @@ def _run_negotiate_test(rank, world_size, test_cases, results_queue, port): disable_overlap_schedule=False, ), max_delay_passes=case.max_delay_passes, + token_usage_low_watermark=case.token_usage_low_watermark, ) for call in case.calls: result = delayer._negotiate_should_allow_prefill( local_prefillable=call.prefillable[rank], + token_usage=call.token_usage[rank], ) results_queue.put((rank, case.name, result.output_allow, result.output_reason)) @@ -79,8 +86,12 @@ _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( name="all_prefillable", max_delay_passes=100, + token_usage_low_watermark=0.8, calls=[ - NegotiateCall(prefillable=[True, True, True, True]), + NegotiateCall( + prefillable=[True, True, True, True], + token_usage=[0.9, 0.9, 0.9, 0.9], + ) ], expected_allow=True, expected_reason="no_wait", @@ -88,9 +99,16 @@ _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( name="all_prefillable_with_previous_wait", max_delay_passes=100, + token_usage_low_watermark=0.8, calls=[ - NegotiateCall(prefillable=[True, False, True, False]), - NegotiateCall(prefillable=[True, True, True, True]), + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ), + NegotiateCall( + prefillable=[True, True, True, True], + token_usage=[0.9, 0.9, 0.9, 0.9], + ), ], expected_allow=True, expected_reason="wait_success", @@ -98,8 +116,12 @@ _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( name="none_prefillable", max_delay_passes=100, + token_usage_low_watermark=0.8, calls=[ - NegotiateCall(prefillable=[False, False, False, False]), + NegotiateCall( + prefillable=[False, False, False, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ) ], expected_allow=True, expected_reason="", @@ -107,8 +129,51 @@ _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( name="mixed_delay", max_delay_passes=100, + token_usage_low_watermark=0.8, calls=[ - NegotiateCall(prefillable=[True, False, True, False]), + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ) + ], + expected_allow=False, + expected_reason="delay", + ), + NegotiateTestCase( + name="mixed_watermark_force_allow", + max_delay_passes=100, + token_usage_low_watermark=0.8, + calls=[ + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.5, 0.9, 0.9, 0.9], + ) + ], + expected_allow=True, + expected_reason="token_watermark", + ), + NegotiateTestCase( + name="mixed_watermark_disabled", + max_delay_passes=100, + token_usage_low_watermark=None, + calls=[ + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.5, 0.9, 0.9, 0.9], + ) + ], + expected_allow=False, + expected_reason="delay", + ), + NegotiateTestCase( + name="mixed_watermark_not_prefillable", + max_delay_passes=100, + token_usage_low_watermark=0.8, + calls=[ + NegotiateCall( + prefillable=[False, False, True, False], + token_usage=[0.5, 0.9, 0.9, 0.9], + ) ], expected_allow=False, expected_reason="delay", @@ -116,10 +181,20 @@ _NEGOTIATE_TEST_CASES = [ NegotiateTestCase( name="mixed_timeout", max_delay_passes=3, + token_usage_low_watermark=0.8, calls=[ - NegotiateCall(prefillable=[True, False, True, False]), - NegotiateCall(prefillable=[True, False, True, False]), - NegotiateCall(prefillable=[True, False, True, False]), + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ), + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ), + NegotiateCall( + prefillable=[True, False, True, False], + token_usage=[0.9, 0.9, 0.9, 0.9], + ), ], expected_allow=True, expected_reason="wait_timeout", @@ -172,6 +247,7 @@ class TestPrefillDelayerThroughputOnlineServing(CustomTestCase): self, test_name="online_serving", other_launch_args=[ + # Not really needed, only to test support non-FCFS algorithms "--schedule-policy", "lpm", ], @@ -196,6 +272,7 @@ class TestPrefillDelayerThroughputOfflineGen(CustomTestCase): random_input_len=30000, random_output_len=500, ), + token_usage_low_watermark=0.8, min_improvement_pct=20, ) @@ -206,11 +283,13 @@ def _run_throughput_comparison( other_launch_args, other_benchmark_args, min_improvement_pct: float, + token_usage_low_watermark: float = None, ): common_kwargs = dict( debug_name=test_name, other_launch_args=other_launch_args, other_benchmark_args=other_benchmark_args, + token_usage_low_watermark=token_usage_low_watermark, ) res_enabled = _run_throughput_test(prefill_delayer=True, **common_kwargs) res_disabled = _run_throughput_test(prefill_delayer=False, **common_kwargs) @@ -229,15 +308,17 @@ def _run_throughput_test( prefill_delayer: bool, other_launch_args, other_benchmark_args, + token_usage_low_watermark: float = None, ): model = "Qwen/Qwen3-0.6B" base_url = DEFAULT_URL_FOR_TEST process = _launch_server( + prefill_delayer=prefill_delayer, model=model, base_url=base_url, other_args=other_launch_args, - prefill_delayer=prefill_delayer, + token_usage_low_watermark=token_usage_low_watermark, ) try: @@ -290,6 +371,87 @@ def _assert_throughput_improvement( ) +class TestPrefillDelayerTokenUsageLowWatermark(CustomTestCase): + def test_1_with_low_watermark(self): + # The kv cache size here is deliberately small, thus we use smaller token usage + self._run(token_usage_low_watermark=0.5) + + def test_2_without_low_watermark(self): + self._run(token_usage_low_watermark=None) + + def _run(self, token_usage_low_watermark): + model = "Qwen/Qwen3-0.6B" + base_url = DEFAULT_URL_FOR_TEST + world_size = int(WORLD_SIZE) + + process = _launch_server( + model=model, + base_url=base_url, + prefill_delayer=True, + other_args=["--max-total-tokens", "50000"], + # e.g. gen throughput is 370 tok/s on H200. + # Will need a different threshold on B200 + max_delay_passes=3000, + token_usage_low_watermark=token_usage_low_watermark, + ) + + async def run_test(): + client = openai.AsyncClient(base_url=f"{base_url}/v1", api_key="EMPTY") + long_prompt = "Hello " * 5000 + + async def send_blocking_request(): + return await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": long_prompt}], + max_tokens=10000, + extra_body={"data_parallel_rank": 0}, + ) + + async def send_normal_request(dp_rank, req_idx): + start = time.time() + await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "Say hi"}], + max_tokens=10, + extra_body={"data_parallel_rank": dp_rank}, + ) + elapsed = time.time() - start + return dp_rank, req_idx, elapsed + + asyncio.create_task(send_blocking_request()) + await asyncio.sleep(3) + + num_reqs_per_rank = 10 + results = await asyncio.gather( + *[ + send_normal_request(dp_rank, req_idx) + for dp_rank in range(1, world_size) + for req_idx in range(num_reqs_per_rank) + ] + ) + + enabled = token_usage_low_watermark is not None + thresh = 5 + for dp_rank, req_idx, elapsed in results: + print(f"DP rank {dp_rank} req {req_idx} completed in {elapsed:.2f}s") + self.assertTrue( + (elapsed < thresh) if enabled else (elapsed > thresh), + f"DP rank {dp_rank} req {req_idx}: elapsed={elapsed:.2f}s, thresh={thresh}, enabled={enabled}. " + f"Maybe you need a different `max_delay_passes` when using hardware other than H200.", + ) + + try: + asyncio.run(run_test()) + + metrics_text = _print_prefill_delayer_metrics(base_url, expect_metrics=True) + if token_usage_low_watermark is not None: + total = _sum_prometheus_metric_values(metrics_text, "token_watermark") + self.assertGreater(total, 0, "Expected token_watermark > 0") + print(f"total token_watermark: {total}") + finally: + kill_process_tree(process.pid) + + class TestPrefillDelayerAccuracy(CustomTestCase): def test_1_mgsm_en_has_prefill_delayer(self): self._run_accuracy_test(prefill_delayer=True) @@ -305,8 +467,10 @@ class TestPrefillDelayerAccuracy(CustomTestCase): model=model, base_url=base_url, other_args=[ + # Not really needed, only to test support non-FCFS algorithms "--schedule-policy", "lpm", + # Use this to ensure prefill delayer will be run "--max-total-tokens", "4096", ], @@ -334,12 +498,17 @@ def _launch_server( prefill_delayer: bool, other_args, max_delay_passes: int = 100, + token_usage_low_watermark: float = None, ): os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1" with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override( prefill_delayer - ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(max_delay_passes): + ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override( + max_delay_passes + ), envs.SGLANG_PREFILL_DELAYER_TOKEN_USAGE_LOW_WATERMARK.override( + token_usage_low_watermark + ): return popen_launch_server( model, base_url,