From e797f0c570ade57cc70849010feeb1d9e7ffe7da Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 4 Jan 2026 11:05:10 +0800 Subject: [PATCH] Support offline generation scenario for prefill delayer (#16363) --- python/sglang/srt/managers/prefill_delayer.py | 56 ++++++++--- python/sglang/srt/managers/schedule_policy.py | 11 +++ python/sglang/srt/managers/scheduler.py | 32 ++++--- test/srt/test_prefill_delayer.py | 94 +++++++++++++------ 4 files changed, 138 insertions(+), 55 deletions(-) diff --git a/python/sglang/srt/managers/prefill_delayer.py b/python/sglang/srt/managers/prefill_delayer.py index e33f68518..0112a82cd 100644 --- a/python/sglang/srt/managers/prefill_delayer.py +++ b/python/sglang/srt/managers/prefill_delayer.py @@ -1,4 +1,5 @@ import logging +from typing import Optional import torch @@ -35,21 +36,7 @@ class PrefillDelayer: not server_args.disable_overlap_schedule ), "To use PrefillDelayer, disable_overlap_schedule must be False." - def _gather_info(self, local_prefillable: int): - local_info = torch.tensor( - [local_prefillable], - device="cpu", - dtype=torch.int64, - ) - torch.distributed.all_gather_into_tensor( - self.global_info.flatten(), - local_info, - group=self.cpu_group, - ) - tp0_info = self.global_info[:, 0, :] - return tp0_info - - def should_allow_prefill(self, local_prefillable: int) -> bool: + def _negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool: tp0_info = self._gather_info(local_prefillable=local_prefillable) global_prefillable = tp0_info[:, 0] global_exists_not_prefillable = global_prefillable.min().item() == 0 @@ -70,3 +57,42 @@ class PrefillDelayer: self.curr_delayed_count = 0 return True + + def _gather_info(self, local_prefillable: bool): + local_info = torch.tensor( + [int(local_prefillable)], + device="cpu", + dtype=torch.int64, + ) + torch.distributed.all_gather_into_tensor( + self.global_info.flatten(), + local_info, + group=self.cpu_group, + ) + tp0_info = self.global_info[:, 0, :] + return tp0_info + + +class PrefillDelayerSinglePassExecutor: + def __init__(self, prefill_delayer: PrefillDelayer): + self._prefill_delayer = prefill_delayer + self._result: Optional[bool] = None + + @property + def _called(self) -> bool: + return self._result is not None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if not self._called: + self.negotiate_should_allow_prefill(local_prefillable=False) + return False + + def negotiate_should_allow_prefill(self, local_prefillable: bool) -> bool: + if not self._called: + self._result = self._prefill_delayer._negotiate_should_allow_prefill( + local_prefillable=local_prefillable + ) + return self._result diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 806650822..e03fd8a27 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -1,5 +1,7 @@ from __future__ import annotations +from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor + # Copyright 2023-2024 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -327,6 +329,7 @@ class PrefillAdder: mixed_with_decode_tokens: int = 0, priority_scheduling_preemption_threshold: int = 0, prefill_max_requests: Optional[int] = None, + prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] = None, ): self.page_size = page_size self.tree_cache = tree_cache @@ -366,6 +369,7 @@ class PrefillAdder: ) self.nsa_prefill_cp_in_seq_split = is_nsa_prefill_cp_in_seq_split() self.prefill_max_requests = prefill_max_requests + self.prefill_delayer_single_pass = prefill_delayer_single_pass def _get_running_request_total_token_offset(self, req: Req) -> int: return ( @@ -614,6 +618,13 @@ class PrefillAdder: if input_tokens >= self.rem_input_tokens and len(self.can_run_list) != 0: return AddReqResult.OTHER + if (self.prefill_delayer_single_pass is not None) and ( + not self.prefill_delayer_single_pass.negotiate_should_allow_prefill( + local_prefillable=True + ) + ): + return AddReqResult.OTHER + if self.rem_chunk_tokens is None or input_tokens <= self.rem_chunk_tokens: # Non-chunked prefill self.can_run_list.append(req) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 84c8d2bd4..dedff1f69 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -22,6 +22,7 @@ import sys import time from collections import deque from concurrent import futures +from contextlib import nullcontext from dataclasses import dataclass from http import HTTPStatus from typing import Any, Deque, Dict, List, Optional, Tuple, Union @@ -120,7 +121,10 @@ from sglang.srt.managers.io_struct import ( ) from sglang.srt.managers.mm_utils import init_mm_embedding_cache from sglang.srt.managers.overlap_utils import FutureMap -from sglang.srt.managers.prefill_delayer import PrefillDelayer +from sglang.srt.managers.prefill_delayer import ( + PrefillDelayer, + PrefillDelayerSinglePassExecutor, +) from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, ModelWorkerBatch, @@ -1911,6 +1915,18 @@ class Scheduler( return res def get_new_batch_prefill(self) -> Optional[ScheduleBatch]: + with ( + PrefillDelayerSinglePassExecutor(self.prefill_delayer) + if self.prefill_delayer + else nullcontext() + ) as prefill_delayer_single_pass: + return self._get_new_batch_prefill_raw( + prefill_delayer_single_pass=prefill_delayer_single_pass + ) + + def _get_new_batch_prefill_raw( + self, prefill_delayer_single_pass: Optional[PrefillDelayerSinglePassExecutor] + ) -> Optional[ScheduleBatch]: # Check if the grammar is ready in the grammar queue if self.grammar_queue: self.move_ready_grammar_requests() @@ -1919,19 +1935,8 @@ class Scheduler( # Reset batch_is_full to try preemption with a prefill adder. self.running_batch.batch_is_full = False - # The `should_allow_prefill` needs to be called on all ranks since contains communication - delayer_allow_prefill = ( - self.prefill_delayer.should_allow_prefill( - # TODO: consider offline generation cases when there are a lot of waiting requests - local_prefillable=len(self.waiting_queue) - > 0, - ) - if self.prefill_delayer - else True - ) if ( - (not delayer_allow_prefill) - or ((self.running_batch.batch_is_full or len(self.waiting_queue) == 0)) + self.running_batch.batch_is_full or len(self.waiting_queue) == 0 ) and self.chunked_req is None: return None @@ -1981,6 +1986,7 @@ class Scheduler( running_bs if self.is_mixed_chunk else 0, self.priority_scheduling_preemption_threshold, prefill_max_requests=self.server_args.prefill_max_requests, + prefill_delayer_single_pass=prefill_delayer_single_pass, ) if self.chunked_req is not None: diff --git a/test/srt/test_prefill_delayer.py b/test/srt/test_prefill_delayer.py index fd70de0fa..adbb7c583 100644 --- a/test/srt/test_prefill_delayer.py +++ b/test/srt/test_prefill_delayer.py @@ -14,59 +14,99 @@ from sglang.test.test_utils import ( class TestPrefillDelayerThroughput(CustomTestCase): - def _run_throughput_test(self, with_prefill_delayer: bool): + def test_1_online_serving_has_prefill_delayer(self): + self._run_throughput_test_online_serving(prefill_delayer=True) + + def test_2_online_serving_no_prefill_delayer(self): + self._run_throughput_test_online_serving(prefill_delayer=False) + + def test_3_offline_gen_has_prefill_delayer(self): + self._run_throughput_test_offline_gen(prefill_delayer=True) + + def test_4_offline_gen_no_prefill_delayer(self): + self._run_throughput_test_offline_gen(prefill_delayer=False) + + def _run_throughput_test_online_serving(self, prefill_delayer: bool): + self._run_throughput_test( + debug_name=f"online_serving ({prefill_delayer=})", + prefill_delayer=prefill_delayer, + other_launch_args=[ + "--mem-fraction-static", + "0.6", + ], + other_benchmark_args=dict( + num_prompts=500, + # trigger chunked prefill + random_input_len=30000, + random_output_len=256, + request_rate=32, + ), + ) + + def _run_throughput_test_offline_gen(self, prefill_delayer: bool): + self._run_throughput_test( + debug_name=f"offline_gen ({prefill_delayer=})", + prefill_delayer=prefill_delayer, + other_benchmark_args=dict( + num_prompts=800, + random_input_len=30000, + random_output_len=500, + ), + other_launch_args=[ + "--max-total-tokens", + "200000", + "--mem-fraction-static", + "0.6", + ], + ) + + def _run_throughput_test( + self, + debug_name: str, + prefill_delayer: bool, + other_launch_args, + other_benchmark_args, + ): os.environ["SGLANG_PREFILL_DELAYER_DEBUG_LOG"] = "1" model = "Qwen/Qwen3-0.6B" base_url = DEFAULT_URL_FOR_TEST - other_args = [ - "--trust-remote-code", - "--tp", - "8", - "--enable-dp-attention", - "--dp", - "8", - "--chunked-prefill-size", - "131072", - "--mem-fraction-static", - "0.6", - ] with envs.SGLANG_SCHEDULER_DECREASE_PREFILL_IDLE.override( - with_prefill_delayer + prefill_delayer ), envs.SGLANG_PREFILL_DELAYER_MAX_DELAY_PASSES.override(100): process = popen_launch_server( model, base_url, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, - other_args=other_args, + other_args=[ + "--trust-remote-code", + "--tp", + "8", + "--enable-dp-attention", + "--dp", + "8", + "--chunked-prefill-size", + "131072", + *other_launch_args, + ], ) try: args = get_benchmark_args( base_url=base_url, dataset_name="random", - num_prompts=500, - # trigger chunked prefill - random_input_len=30000, - random_output_len=256, - request_rate=32, tokenizer=model, + **other_benchmark_args, ) res = run_benchmark(args) finally: kill_process_tree(process.pid) - print(f"=== {with_prefill_delayer=} ===") + print(f"=== {debug_name} ===") print(f"Input throughput: {res['input_throughput']:.2f} token/s") print(f"Output throughput: {res['output_throughput']:.2f} token/s") - def test_1_dp_attention_throughput_with_prefill_delayer(self): - self._run_throughput_test(with_prefill_delayer=True) - - def test_2_dp_attention_throughput_without_prefill_delayer(self): - self._run_throughput_test(with_prefill_delayer=False) - if __name__ == "__main__": unittest.main()