diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index c09d1b354..5964e1cbd 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -480,6 +480,7 @@ class DecodePreallocQueue: pass elif poll == KVPoll.WaitingForInput: decode_req.waiting_for_input = True + decode_req.req.time_stats.set_bootstrap_done_time() elif poll == KVPoll.Failed: error_message = f"Decode handshake failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" try: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 01f721e0d..fa9964026 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -42,7 +42,12 @@ from sglang.srt.disaggregation.utils import ( poll_and_all_reduce_attn_cp_tp_group, prepare_abort, ) -from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch +from sglang.srt.managers.schedule_batch import ( + FINISH_ABORT, + FINISH_LENGTH, + Req, + ScheduleBatch, +) from sglang.srt.mem_cache.common import release_kv_cache from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool @@ -299,6 +304,7 @@ class PrefillBootstrapQueue: continue # KV.WaitingForInput - init here + req.time_stats.set_bootstrap_done_time() num_kv_indices = len(req.origin_input_ids) if self.req_to_metadata_buffer_idx_allocator.available_size() == 0: break @@ -597,6 +603,27 @@ class SchedulerDisaggregationPrefillMixin: for req in done_reqs: req.time_stats.set_completion_time() + page_size = self.token_to_kv_pool_allocator.page_size + kv_item_lens = ( + self.disagg_prefill_bootstrap_queue.kv_manager.kv_args.kv_item_lens + ) + bytes_per_page_all_layers = sum(kv_item_lens) + + for req in done_reqs: + if isinstance(req.finished_reason, FINISH_ABORT): + continue + metrics = req.time_stats.compute_and_observe_kv_transfer_metrics( + num_tokens=len(req.origin_input_ids), + page_size=page_size, + bytes_per_page_all_layers=bytes_per_page_all_layers, + ) + if metrics: + # Update last-value for REST API + if "latency_ms" in metrics: + self.kv_transfer_latency_ms = metrics["latency_ms"] + if "speed_gb_s" in metrics: + self.kv_transfer_speed_gb_s = metrics["speed_gb_s"] + # Stream requests which have finished transfer self.stream_output( done_reqs, diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 07c24f950..f5855f1db 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -80,9 +80,6 @@ class SchedulerStats: num_decode_transfer_queue_reqs: int = 0 kv_transfer_speed_gb_s: float = 0.0 kv_transfer_latency_ms: float = 0.0 - kv_transfer_bootstrap_ms: float = 0.0 - kv_transfer_alloc_ms: float = 0.0 - kv_transfer_total_mb: float = 0.0 # Utilization utilization: float = 0.0 @@ -322,35 +319,35 @@ class SchedulerMetricsCollector: documentation="Total number of prefill retries.", labelnames=labels.keys(), ) - self.kv_transfer_speed_gb_s = Gauge( + self.kv_transfer_speed_gb_s = Histogram( name="sglang:kv_transfer_speed_gb_s", - documentation="The transfer speed of the KV cache in GB/s.", + documentation="Histogram of KV cache transfer speed in GB/s.", labelnames=labels.keys(), - multiprocess_mode="mostrecent", + buckets=(0.1, 0.5, 1, 5, 10, 25, 50, 100, 200, 400), ) - self.kv_transfer_latency_ms = Gauge( + self.kv_transfer_latency_ms = Histogram( name="sglang:kv_transfer_latency_ms", - documentation="The transfer latency of the KV cache in ms.", + documentation="Histogram of KV cache transfer latency in ms.", labelnames=labels.keys(), - multiprocess_mode="mostrecent", + buckets=(1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000), ) - self.kv_transfer_bootstrap_ms = Gauge( + self.kv_transfer_bootstrap_ms = Histogram( name="sglang:kv_transfer_bootstrap_ms", - documentation="The bootstrap time of the KV transfer in ms.", + documentation="Histogram of KV transfer bootstrap time in ms.", labelnames=labels.keys(), - multiprocess_mode="mostrecent", + buckets=(1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500), ) - self.kv_transfer_alloc_ms = Gauge( + self.kv_transfer_alloc_ms = Histogram( name="sglang:kv_transfer_alloc_ms", - documentation="The allocation waiting time of the KV transfer in ms.", + documentation="Histogram of KV transfer allocation waiting time in ms.", labelnames=labels.keys(), - multiprocess_mode="mostrecent", + buckets=(1, 2, 5, 10, 25, 50, 100, 250, 500, 1000, 2500), ) - self.kv_transfer_total_mb = Gauge( + self.kv_transfer_total_mb = Histogram( name="sglang:kv_transfer_total_mb", - documentation="The total number of tokens transferred in the KV cache.", + documentation="Histogram of KV cache transfer size in MB.", labelnames=labels.keys(), - multiprocess_mode="mostrecent", + buckets=(1, 5, 10, 50, 100, 500, 1000, 5000, 10000), ) # Utilization @@ -750,6 +747,24 @@ class SchedulerMetricsCollector: if count > 0: self.num_prefill_retries_total.labels(**self.labels).inc(count) + def observe_kv_transfer_metrics( + self, + latency_ms: float, + total_mb: float, + speed_gb_s: float, + ) -> None: + self._log_histogram(self.kv_transfer_latency_ms, latency_ms) + self._log_histogram(self.kv_transfer_total_mb, total_mb) + self._log_histogram(self.kv_transfer_speed_gb_s, speed_gb_s) + + def observe_kv_transfer_bootstrap( + self, + bootstrap_ms: float, + alloc_ms: float, + ) -> None: + self._log_histogram(self.kv_transfer_bootstrap_ms, bootstrap_ms) + self._log_histogram(self.kv_transfer_alloc_ms, alloc_ms) + def observe_per_stage_req_latency(self, stage: str, latency: float) -> None: labels_with_stage = {**self.labels, "stage": stage} self.per_stage_req_latency_seconds.labels(**labels_with_stage).observe(latency) @@ -880,12 +895,6 @@ class SchedulerMetricsCollector: self._log_gauge( self.num_decode_transfer_queue_reqs, stats.num_decode_transfer_queue_reqs ) - self._log_gauge(self.kv_transfer_speed_gb_s, stats.kv_transfer_speed_gb_s) - self._log_gauge(self.kv_transfer_latency_ms, stats.kv_transfer_latency_ms) - self._log_gauge(self.kv_transfer_bootstrap_ms, stats.kv_transfer_bootstrap_ms) - self._log_gauge(self.kv_transfer_alloc_ms, stats.kv_transfer_alloc_ms) - self._log_gauge(self.kv_transfer_total_mb, stats.kv_transfer_total_mb) - # Retract self._log_gauge(self.num_retracted_reqs, stats.num_retracted_reqs) self._log_gauge(self.num_paused_reqs, stats.num_paused_reqs) diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index bb859aefd..5dc93474b 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -547,6 +547,9 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): decode_transfer_queue_entry_time: float = 0.0 decode_prebuilt_finish_time: float = 0.0 + # bootstrap sub-phase tracking (PD disagg) + bootstrap_done_time: float = 0.0 + # only for request tracing scheduler_recv_time: float = 0.0 last_chunked_prefill_finish_time: float = 0.0 @@ -778,6 +781,70 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): self.trace_ctx.abort() + def compute_and_observe_kv_transfer_metrics( + self, + num_tokens: int, + page_size: int, + bytes_per_page_all_layers: int, + ) -> Optional[dict]: + """Compute KV transfer metrics and observe them via the metrics collector. + + Returns a dict with latency_ms, total_mb, speed_gb_s if computable, else None. + """ + from sglang.srt.disaggregation.utils import kv_to_page_num + + result = {} + + # Transfer latency, size, and speed + if self.prefill_transfer_queue_entry_time > 0 and self.completion_time > 0: + transfer_latency_s = ( + self.completion_time - self.prefill_transfer_queue_entry_time + ) + latency_ms = transfer_latency_s * 1000 + + num_pages = kv_to_page_num(num_tokens, page_size) + total_bytes = bytes_per_page_all_layers * num_pages + total_mb = total_bytes / (1024 * 1024) + self.transfer_total_mb = total_mb + + speed_gb_s = 0.0 + if transfer_latency_s > 0: + speed_gb_s = (total_mb / 1024) / transfer_latency_s + self.transfer_speed_gb_s = speed_gb_s + + result["latency_ms"] = latency_ms + result["total_mb"] = total_mb + result["speed_gb_s"] = speed_gb_s + + if self.enable_metrics: + self.metrics_collector.observe_kv_transfer_metrics( + latency_ms=latency_ms, + total_mb=total_mb, + speed_gb_s=speed_gb_s, + ) + + # Bootstrap and alloc durations + if ( + self.prefill_bootstrap_queue_entry_time > 0 + and self.bootstrap_done_time > 0 + and self.wait_queue_entry_time > 0 + ): + bootstrap_ms = ( + self.bootstrap_done_time - self.prefill_bootstrap_queue_entry_time + ) * 1000 + alloc_ms = (self.wait_queue_entry_time - self.bootstrap_done_time) * 1000 + + result["bootstrap_ms"] = bootstrap_ms + result["alloc_ms"] = alloc_ms + + if self.enable_metrics: + self.metrics_collector.observe_kv_transfer_bootstrap( + bootstrap_ms=bootstrap_ms, + alloc_ms=alloc_ms, + ) + + return result if result else None + def set_quick_finish_time(self, ts=None): if ts is None: ts = time.perf_counter() @@ -829,6 +896,12 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): ) self.trace_slice(stage, self.decode_prealloc_queue_entry_time, ts) + def set_bootstrap_done_time(self, ts=None): + if ts is None: + ts = time.perf_counter() + if self.bootstrap_done_time == 0.0: + self.bootstrap_done_time = ts + def set_decode_prebuilt_finish_time(self, ts=None): if ts is None: ts = time.perf_counter() @@ -866,7 +939,7 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): return f"queue_duration={self.format_duration(queue_duration)}, forward_duration={self.format_duration(forward_duration)}, start_time={self.wait_queue_entry_time:.3f}" elif self.disagg_mode == DisaggregationMode.PREFILL: - bootstrap_duration = ( + bootstrap_queue_duration = ( self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time ) queue_duration = self.forward_entry_time - self.wait_queue_entry_time @@ -875,13 +948,33 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): if SGLANG_TEST_REQUEST_TIME_STATS: if self.wait_queue_entry_time > 0: assert ( - bootstrap_duration >= 0 + bootstrap_queue_duration >= 0 and queue_duration >= 0 and forward_duration >= 0 - ), f"bootstrap_duration={bootstrap_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + ), f"bootstrap_queue_duration={bootstrap_queue_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0" + + # Break down bootstrap_queue_duration into sub-phases + if self.bootstrap_done_time > 0: + bootstrap_duration = ( + self.bootstrap_done_time - self.prefill_bootstrap_queue_entry_time + ) + alloc_wait_duration = ( + self.wait_queue_entry_time - self.bootstrap_done_time + ) + if SGLANG_TEST_REQUEST_TIME_STATS: + assert ( + bootstrap_duration >= 0 and alloc_wait_duration >= 0 + ), f"bootstrap_duration={bootstrap_duration} < 0 or alloc_wait_duration={alloc_wait_duration} < 0" + bootstrap_breakdown = ( + f"= bootstrap({self.format_duration(bootstrap_duration)}) " + f"+ alloc_wait({self.format_duration(alloc_wait_duration)}); " + ) + else: + bootstrap_breakdown = "" return ( - f"bootstrap_queue_duration({self.format_duration(bootstrap_duration)}) " + f"bootstrap_queue_duration({self.format_duration(bootstrap_queue_duration)}) " + f"{bootstrap_breakdown}" f"queue_duration={self.format_duration(queue_duration)}, " f"forward_duration={self.format_duration(forward_duration)}, " f"start={self.prefill_bootstrap_queue_entry_time:.3f}, " @@ -909,8 +1002,28 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): and forward_duration >= 0 ), f"prealloc_duration={prealloc_duration} < 0 or transfer_duration={transfer_duration} < 0 or queue_duration={queue_duration} < 0 or forward_duration={forward_duration} < 0. {self=}" + # Break down prealloc_duration into sub-phases + if self.bootstrap_done_time > 0: + bootstrap_duration = ( + self.bootstrap_done_time - self.decode_prealloc_queue_entry_time + ) + alloc_wait_duration = ( + self.decode_transfer_queue_entry_time - self.bootstrap_done_time + ) + if SGLANG_TEST_REQUEST_TIME_STATS: + assert ( + bootstrap_duration >= 0 and alloc_wait_duration >= 0 + ), f"bootstrap_duration={bootstrap_duration} < 0 or alloc_wait_duration={alloc_wait_duration} < 0" + prealloc_breakdown = ( + f"= bootstrap({self.format_duration(bootstrap_duration)}) " + f"+ alloc_wait({self.format_duration(alloc_wait_duration)}); " + ) + else: + prealloc_breakdown = "" + return ( f"prealloc_queue_duration({self.format_duration(prealloc_duration)}) " + f"{prealloc_breakdown}" f"transfer_duration={self.format_duration(transfer_duration)}; " f"queue_duration={self.format_duration(queue_duration)}, " f"forward_duration={self.format_duration(forward_duration)}, " diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py index 487482498..6ddcc7c4d 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -90,9 +90,6 @@ class SchedulerMetricsMixin: # For PD disaggregation self.kv_transfer_speed_gb_s: float = 0.0 self.kv_transfer_latency_ms: float = 0.0 - self.kv_transfer_bootstrap_ms: float = 0.0 - self.kv_transfer_alloc_ms: float = 0.0 - self.kv_transfer_total_mb: float = 0.0 self.stats = SchedulerStats() @@ -284,9 +281,6 @@ class SchedulerMetricsMixin: ) self.stats.kv_transfer_speed_gb_s = self.kv_transfer_speed_gb_s self.stats.kv_transfer_latency_ms = self.kv_transfer_latency_ms - self.stats.kv_transfer_bootstrap_ms = self.kv_transfer_bootstrap_ms - self.stats.kv_transfer_alloc_ms = self.kv_transfer_alloc_ms - self.stats.kv_transfer_total_mb = self.kv_transfer_total_mb elif self.disaggregation_mode == DisaggregationMode.DECODE: self.stats.num_decode_prealloc_queue_reqs = len( self.disagg_decode_prealloc_queue.queue