diff --git a/python/sglang/srt/managers/load_snapshot.py b/python/sglang/srt/managers/load_snapshot.py index fd0f53246..950ef4f18 100644 --- a/python/sglang/srt/managers/load_snapshot.py +++ b/python/sglang/srt/managers/load_snapshot.py @@ -173,20 +173,10 @@ class QueueMetrics(msgspec.Struct, array_like=True): retracted: int -_CORE_KEYS = ( - "timestamp", - "dp_rank", - "num_running_reqs", - "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", +# LoadSnapshot's nested sub-struct fields; every other struct field is a flat +# scalar returned under "core". +_SECTION_FIELDS = frozenset( + {"memory", "speculative", "lora", "disaggregation", "queues"} ) @@ -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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index cc02f4a61..b1f8224f9 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 967e1120e..5c2a091d0 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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. diff --git a/python/sglang/srt/managers/scheduler_components/load_inquirer.py b/python/sglang/srt/managers/scheduler_components/load_inquirer.py index e8619a2f7..3edcf86a1 100644 --- a/python/sglang/srt/managers/scheduler_components/load_inquirer.py +++ b/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, ) diff --git a/test/registered/unit/managers/test_load_snapshot_backends.py b/test/registered/unit/managers/test_load_snapshot_backends.py index 12fb3386f..e689d4958 100644 --- a/test/registered/unit/managers/test_load_snapshot_backends.py +++ b/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()