diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 109fa0e64..ffd77cc4d 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -56,6 +56,9 @@ from sglang.srt.dllm.mixin.req import ReqDllmMixin from sglang.srt.environ import envs from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.managers.embed_types import PositionalEmbeds +from sglang.srt.managers.scheduler_components.new_token_ratio_tracker import ( + NewTokenRatioTracker, +) from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( BasePrefixCache, @@ -2229,16 +2232,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.filter_batch(keep_indices=sorted_indices) # Reqs in batch are filtered - total_decoded_tokens = sum(len(r.output_ids) for r in self.reqs) - total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in self.reqs) - new_estimate_ratio = ( - total_decoded_tokens - + envs.SGLANG_RETRACT_DECODE_STEPS.get() * len(self.reqs) - ) / ( - total_max_new_tokens + 1 - ) # avoid zero division - new_estimate_ratio = min(1.0, new_estimate_ratio) + NewTokenRatioTracker.estimate_new_token_ratio_after_retract(self.reqs) + ) return retracted_reqs, new_estimate_ratio, reqs_to_abort diff --git a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py index 57512a067..81f7e7801 100644 --- a/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py +++ b/python/sglang/srt/managers/scheduler_components/new_token_ratio_tracker.py @@ -1,8 +1,14 @@ +from __future__ import annotations + from dataclasses import dataclass +from typing import TYPE_CHECKING, Sequence from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import Req + @dataclass(slots=True, kw_only=True) class NewTokenRatioTracker: @@ -30,3 +36,16 @@ class NewTokenRatioTracker: def reset(self) -> None: self.current = self.init + + @staticmethod + def estimate_new_token_ratio_after_retract(reqs: Sequence[Req]) -> float: + total_decoded_tokens = sum(len(r.output_ids) for r in reqs) + total_max_new_tokens = sum(r.sampling_params.max_new_tokens for r in reqs) + + new_estimate_ratio = ( + total_decoded_tokens + envs.SGLANG_RETRACT_DECODE_STEPS.get() * len(reqs) + ) / ( + total_max_new_tokens + 1 + ) # avoid zero division + new_estimate_ratio = min(1.0, new_estimate_ratio) + return new_estimate_ratio