chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fix disagg PD bootstrap and KV transfer metrics (#19009)

Browse Source 
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
2026-03-04 09:08:10 -08:00
committed by GitHub
parent 34c19a32c1
commit c03deb8175
5 changed files with 179 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/disaggregation/decode.py
View File

@@ -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:

29
python/sglang/srt/disaggregation/prefill.py
View File

@@ -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,

57
python/sglang/srt/observability/metrics_collector.py
View File

@@ -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)

121
python/sglang/srt/observability/req_time_stats.py
View File

@@ -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)}, "

6
python/sglang/srt/observability/scheduler_metrics_mixin.py
View File

@@ -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