Add prefill and decode load counters to LoadSnapshot (#32245)

python/sglang/srt/managers/load_snapshot.py

@@ -173,20 +173,10 @@ class QueueMetrics(msgspec.Struct, array_like=True):
     retracted: int
 
 
-_CORE_KEYS = (
+# LoadSnapshot's nested sub-struct fields; every other struct field is a flat
-    "timestamp",
+# scalar returned under "core".
-    "dp_rank",
+_SECTION_FIELDS = frozenset(
-    "num_running_reqs",
+    {"memory", "speculative", "lora", "disaggregation", "queues"}
-    "num_waiting_reqs",
-    "num_waiting_uncached_tokens",
-    "num_used_tokens",
-    "num_total_tokens",
-    "max_total_num_tokens",
-    "max_running_requests",
-    "token_usage",
-    "gen_throughput",
-    "cache_hit_rate",
-    "utilization",
 )
 
 
@@ -200,12 +190,20 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True):
     num_waiting_uncached_tokens: int = 0
     num_used_tokens: int = 0
     num_total_tokens: int = 0
+    # num_total_tokens minus tokens still awaiting a KV transfer (equal to it
+    # outside disaggregated decode).
+    num_active_tokens: int = 0
     max_total_num_tokens: int = 0
     max_running_requests: int = 0
     token_usage: float = 0.0
     gen_throughput: float = 0.0
     cache_hit_rate: float = 0.0
     utilization: float = 0.0
+    # cumulative counters
+    total_prefill_uncached_tokens: int = 0
+    total_prefill_busy_us: int = 0
+    # Decode step-time moment sums
+    decode_moments: Optional[list[float]] = None
 
     memory: Optional[MemoryMetrics] = None
     speculative: Optional[SpeculativeMetrics] = None
@@ -246,6 +244,12 @@ class LoadSnapshot(msgspec.Struct, omit_defaults=True):
         return load
 
 
+# Flat scalar fields returned under "core": every struct field but the sections.
+_CORE_KEYS = tuple(
+    f for f in LoadSnapshot.__struct_fields__ if f not in _SECTION_FIELDS
+)
+
+
 def _enc_hook(obj):
     """Coerce numpy scalars to native Python; msgpack has no numpy types."""
     to_item = getattr(obj, "item", None)

python/sglang/srt/managers/schedule_batch.py

@@ -1877,6 +1877,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     dp_cooperation_info: Optional[DPCooperationInfo] = None
     prefill_stats: Optional[PrefillStats] = None
     forward_iter: Optional[int] = None
+    launch_ts: Optional[float] = None
 
     # === GPU tensors crossing to ForwardBatch (clone targets for stream isolation) ===
     # Batched arguments to model runner
@@ -3082,6 +3083,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             prefill_stats=self.prefill_stats,
             fpm_start_time=self.fpm_start_time,
             forward_iter=self.forward_iter,
+            launch_ts=self.launch_ts,
+            extend_num_tokens=self.extend_num_tokens,
         )
 
     def maybe_evict_swa(self):

python/sglang/srt/managers/scheduler.py

@@ -296,6 +296,29 @@ TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get()
 TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get()
 TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get()
 
+
+DECODE_STEP_MAX_US = 2_000_000
+
+
+def _accumulate_decode_moment(
+    totals: list[float],
+    batch_size: int,
+    step_us: int,
+    generated: int,
+) -> None:
+    if batch_size <= 0 or step_us <= 0:
+        return
+    b = float(batch_size)
+    t = float(step_us)
+    g = float(generated)
+    totals[0] += 1.0
+    totals[1] += b
+    totals[2] += t
+    totals[3] += b * b
+    totals[4] += b * t
+    totals[5] += g
+
+
 _is_npu = is_npu()
 _is_hip = is_hip()
 
@@ -1842,6 +1865,10 @@ class Scheduler(
         )
 
     def init_load_inquirer(self) -> None:
+        self.total_prefill_uncached_tokens = 0
+        self.total_prefill_busy_us = 0
+        self.decode_moment_totals: list[float] = [0.0] * 6
+        self._prev_decode_launch_ts: Optional[float] = None
         self.load_inquirer = SchedulerLoadInquirer(
             disaggregation_mode=self.disaggregation_mode,
             ps=self.ps,
@@ -1862,6 +1889,9 @@ class Scheduler(
             get_disagg_decode_transfer_queue=lambda: self.disagg_decode_transfer_queue,
             get_spec_total_num_accept_tokens=lambda: self.metrics_reporter.spec_total_num_accept_tokens,
             get_spec_total_num_forward_ct=lambda: self.metrics_reporter.spec_total_num_forward_ct,
+            get_total_prefill_uncached_tokens=lambda: self.total_prefill_uncached_tokens,
+            get_total_prefill_busy_us=lambda: self.total_prefill_busy_us,
+            get_decode_moment_totals=lambda: self.decode_moment_totals,
         )
 
     def init_output_streamer(self) -> None:
@@ -3313,6 +3343,7 @@ class Scheduler(
         """Run a batch."""
         self.forward_ct += 1
         batch.forward_iter = self.forward_ct
+        batch.launch_ts = time.monotonic()
 
         if self.scripted_scheduler_hook is not None:
             self.scripted_scheduler_hook.on_run_batch(batch)
@@ -3602,6 +3633,8 @@ class Scheduler(
         elif batch.forward_mode.is_idle():
             self.batch_result_processor.process_batch_result_idle(batch, result)
 
+        self._record_step_counters(batch, result)
+
         self.metrics_reporter.log_batch_result_stats(batch, result)
 
         # Emit forward pass metrics (every iteration when enabled)
@@ -3612,6 +3645,33 @@ class Scheduler(
         self.maybe_send_health_check_signal()
         self.metrics_reporter.update_device_timer()
 
+    def _record_step_counters(
+        self, batch: ScheduleBatch, result: GenerationBatchResult
+    ) -> None:
+        mode = batch.forward_mode
+        is_prefill = mode.is_extend_without_speculative()
+        if not (is_prefill or mode.is_decode() or mode.is_target_verify()):
+            return
+        if all(is_health_check_generate_req(req) for req in batch.reqs):
+            return
+        if is_prefill:
+            # Busy span = run_batch entry -> result processed.
+            span_us = int((time.monotonic() - batch.launch_ts) * 1e6)
+            self.total_prefill_busy_us += span_us
+            self.total_prefill_uncached_tokens += batch.extend_num_tokens
+        else:
+            batch_size = len(batch.reqs)
+            if self._prev_decode_launch_ts is not None:
+                step_us = int((batch.launch_ts - self._prev_decode_launch_ts) * 1e6)
+                if 0 < step_us < DECODE_STEP_MAX_US:
+                    _accumulate_decode_moment(
+                        self.decode_moment_totals,
+                        batch_size,
+                        step_us,
+                        batch_size + result.num_correct_drafts,
+                    )
+            self._prev_decode_launch_ts = batch.launch_ts
+
     def maybe_send_health_check_signal(self):
         if self.return_health_check_ipcs:
             # Return some signal for the health check.

python/sglang/srt/managers/scheduler_components/load_inquirer.py

@@ -50,6 +50,9 @@ class SchedulerLoadInquirer:
     get_disagg_decode_transfer_queue: Callable
     get_spec_total_num_accept_tokens: Callable
     get_spec_total_num_forward_ct: Callable
+    get_total_prefill_uncached_tokens: Callable
+    get_total_prefill_busy_us: Callable
+    get_decode_moment_totals: Callable
 
     def _get_num_pending_tokens(self, chunk_deduct: int = 0) -> int:
         """Get the total number of tokens pending prefill.
@@ -91,6 +94,7 @@ class SchedulerLoadInquirer:
 
         waiting_queues = [self.get_waiting_queue()]
         pending_token_queues = [self.get_waiting_queue()]
+        awaiting_kv_tokens = 0
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
             prefill_bootstrap_queue = self.get_disagg_prefill_bootstrap_queue().queue
             waiting_queues.append(prefill_bootstrap_queue)
@@ -108,6 +112,12 @@ class SchedulerLoadInquirer:
             # waiting-queue requests have already pre-allocated decode-side KV
             # slots, so they are already included in num_used_tokens.
             pending_token_queues = [decode_prealloc_queue, decode_retracted_queue]
+            # KV not yet arrived from the prefill side.
+            awaiting_kv_tokens = sum(
+                req.seqlen
+                for queue in (decode_prealloc_queue, decode_transfer_queue)
+                for req in queue
+            )
 
         num_waiting_reqs = sum(len(queue) for queue in waiting_queues)
         num_used_tokens, kv_token_usage = (
@@ -116,6 +126,7 @@ class SchedulerLoadInquirer:
         num_total_tokens = num_used_tokens + sum(
             req.seqlen for queue in pending_token_queues for req in queue
         )
+        num_active_tokens = max(0, num_total_tokens - awaiting_kv_tokens)
 
         memory = None
         try:
@@ -183,6 +194,9 @@ class SchedulerLoadInquirer:
             retracted=stats.num_retracted_reqs,
         )
 
+        totals = self.get_decode_moment_totals()
+        decode_moments = list(totals) if totals[0] > 0 else None
+
         return LoadSnapshot(
             dp_rank=int(self.ps.dp_rank) if self.ps.dp_rank is not None else 0,
             timestamp=time.time(),
@@ -191,6 +205,7 @@ class SchedulerLoadInquirer:
             num_waiting_uncached_tokens=self.get_num_waiting_uncached_tokens(),
             num_used_tokens=num_used_tokens,
             num_total_tokens=num_total_tokens,
+            num_active_tokens=num_active_tokens,
             max_total_num_tokens=self.max_total_num_tokens,
             max_running_requests=self.max_running_requests,
             token_usage=round(kv_token_usage, 4),
@@ -202,4 +217,7 @@ class SchedulerLoadInquirer:
             lora=lora,
             disaggregation=disaggregation,
             queues=queues,
+            total_prefill_uncached_tokens=self.get_total_prefill_uncached_tokens(),
+            total_prefill_busy_us=self.get_total_prefill_busy_us(),
+            decode_moments=decode_moments,
         )

test/registered/unit/managers/test_load_snapshot_backends.py

@@ -64,11 +64,28 @@ class TestShmRoundTrip(CustomTestCase):
         writer = ShmLoadSnapshotWriter(path, dp_size=1, dp_rank=0)
         reader = ShmLoadSnapshotReader(path, dp_size=1)
         try:
-            writer.write(LoadSnapshot(dp_rank=0, num_running_reqs=5, timestamp=1.0))
+            writer.write(
+                LoadSnapshot(
+                    dp_rank=0,
+                    num_running_reqs=5,
+                    timestamp=1.0,
+                    num_active_tokens=4096,
+                    total_prefill_uncached_tokens=1000,
+                    total_prefill_busy_us=250_000,
+                    decode_moments=[2, 30, 3000, 500, 50_000, 60],
+                )
+            )
             load = reader.read(0)
             self.assertIsNotNone(load)
             self.assertEqual(load.num_running_reqs, 5)
             self.assertEqual(load.timestamp, 1.0)
+            self.assertEqual(load.num_active_tokens, 4096)
+            # The cumulative total_* counters round-trip like any other
+            # core scalar.
+            self.assertEqual(load.total_prefill_uncached_tokens, 1000)
+            self.assertEqual(load.total_prefill_busy_us, 250_000)
+            self.assertEqual(load.decode_moments[0], 2)
+            self.assertEqual(load.decode_moments[5], 60)
         finally:
             reader.close()
             writer.close()