From 2d868656d0290bea39c4c36535e1017393ab8909 Mon Sep 17 00:00:00 2001
From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com>
Date: Tue, 19 May 2026 09:19:42 +0800
Subject: [PATCH] Move the retract-decode ratio estimation onto the
 new-token-ratio tracker (#25717)

---
 python/sglang/srt/managers/schedule_batch.py  | 14 +++++---------
 .../new_token_ratio_tracker.py                | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 9 deletions(-)

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