[metrics] Report logical prefill token counts (#38566)

python/sglang/srt/managers/schedule_policy.py

@@ -583,7 +583,6 @@ class PrefillAdder:
         self.log_device_hit_tokens = 0
         self.log_host_hit_tokens = 0
         self.log_storage_hit_tokens = 0
-        # TODO(lsyin): report the real input tokens excluding page alignment
         self.log_input_tokens = 0
         self.reprocessed_log_input_tokens = 0
 
@@ -906,6 +905,7 @@ class PrefillAdder:
         is_chunked_continuation: bool = False,
     ):
         # TODO(lsyin): check this workaround logic, which only ensures the prefill will not out of memory, and may be too conservative
+        raw_extend_input_len = extend_input_len
         extend_input_len = self.ceil_paged_tokens(extend_input_len)
 
         # alloc_extend reserves an extra page_size per request to make sure the budget doesn't over-commit
@@ -942,10 +942,10 @@ class PrefillAdder:
         # reprocessed_log_* is a subset of log_*; metrics_reporter subtracts it
         # when computing the first-attempt prefix cache hit rate.
         self.log_hit_tokens += prefix_len
-        self.log_input_tokens += extend_input_len
+        self.log_input_tokens += raw_extend_input_len
         if retracted_stain:
             self.reprocessed_log_hit_tokens += prefix_len
-            self.reprocessed_log_input_tokens += extend_input_len
+            self.reprocessed_log_input_tokens += raw_extend_input_len
 
     def _account_prefill_cache_admission(self, req: Req, prefix_len: int) -> None:
         if req.retracted_stain:
@@ -1439,7 +1439,7 @@ class PrefillAdder:
                 self._req_inc_lock_ref(req)
                 self._update_prefill_budget(
                     prefix_len,
-                    input_tokens,
+                    req.extend_range.length,
                     min(
                         req.sampling_params.max_new_tokens,
                         CLIP_MAX_NEW_TOKENS,

test/registered/unit/managers/test_prefill_adder.py

@@ -443,6 +443,13 @@ class TestPrefillAdder(CustomTestCase):
         req1.full_untruncated_fill_ids = list(range(56))
         req1.last_node = MagicMock()
         req1.sampling_params.ignore_eos = False
+        # add_one_req reads req.extend_range.length after set_extend_range;
+        # emulate the real Req writer (a spec=Req mock lacks the attribute).
+        req1.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req1, "extend_range", Range(start, end)
+            )
+        )
 
         result1 = adder.add_one_req(
             req1, has_chunked_req=False, truncation_align_size=None
@@ -476,6 +483,11 @@ class TestPrefillAdder(CustomTestCase):
         req2.full_untruncated_fill_ids = list(range(56))
         req2.last_node = MagicMock()
         req2.sampling_params.ignore_eos = False
+        req2.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req2, "extend_range", Range(start, end)
+            )
+        )
 
         result2 = adder2.add_one_req(
             req2, has_chunked_req=False, truncation_align_size=None
@@ -492,6 +504,11 @@ class TestPrefillAdder(CustomTestCase):
         req3.full_untruncated_fill_ids = list(range(3))
         req3.last_node = MagicMock()
         req3.sampling_params.ignore_eos = False
+        req3.set_extend_range = MagicMock(
+            side_effect=lambda start, end: setattr(
+                req3, "extend_range", Range(start, end)
+            )
+        )
 
         result3 = adder2.add_one_req(
             req3, has_chunked_req=False, truncation_align_size=None