diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 2309e8a83..c87b7ce41 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dataclasses from abc import ABC, abstractmethod from typing import TYPE_CHECKING, List, Optional @@ -12,6 +13,13 @@ if TYPE_CHECKING: from sglang.srt.disaggregation.utils import DisaggregationMode +@dataclasses.dataclass +class KVTransferMetric: + # Backends that cannot isolate transfer latency can leave this as None. + transfer_latency_s: Optional[float] = None + transfer_total_bytes: Optional[int] = None + + class KVArgs: engine_rank: int kv_data_ptrs: List[int] @@ -95,6 +103,11 @@ class BaseKVSender(ABC): """ ... + @abstractmethod + def get_transfer_metric(self) -> KVTransferMetric: + """Return backend-specific transfer metrics for this sender.""" + ... + @abstractmethod def poll(self) -> KVPoll: """ diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 796bf92b6..824cce206 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -25,6 +25,7 @@ from sglang.srt.disaggregation.base.conn import ( BaseKVSender, KVArgs, KVPoll, + KVTransferMetric, ) from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.distributed import get_pp_group @@ -178,6 +179,10 @@ class CommonKVManager(BaseKVManager): is_mla_backend: Optional[bool] = False, ): self.kv_args = args + # Per-page byte totals used for accurate transfer-metric accounting + # (ported from upstream sgl-project/sglang #24416). + self.kv_item_lens_sum = sum(getattr(args, "kv_item_lens", None) or []) + self.state_item_lens_sum = sum(getattr(args, "state_item_lens", None) or []) self.is_mla_backend = is_mla_backend self.disaggregation_mode = disaggregation_mode self.server_args = server_args @@ -533,6 +538,12 @@ class CommonKVSender(BaseKVSender): self.bootstrap_room = bootstrap_room self.aux_index = None self.bootstrap_server_url = bootstrap_addr + # Transfer-metric accounting: count the KV/state indices actually sent + # (ported from upstream sgl-project/sglang #24416). Set before the dummy + # CP-rank early return so every sender has the fields. + self._transfer_metric = KVTransferMetric() + self._transfer_num_kv_indices = 0 + self._transfer_num_state_indices = 0 # inner state self.curr_idx = 0 if self.kv_mgr.is_dummy_cp_rank: @@ -563,6 +574,21 @@ class CommonKVSender(BaseKVSender): except Exception as e: logger.error(f"Failed to register prefill dp_rank: {e}") + def get_transfer_metric(self) -> KVTransferMetric: + # Bytes actually transferred = sent KV pages * per-page KV bytes (all + # layers) + sent state pages * per-page state bytes. This replaces the + # old full-prompt-length estimate, which over-reported on cache hits and + # CP shared-KV per-rank filtering. Ported from upstream #24416. + total_bytes = self._transfer_num_kv_indices * self.kv_mgr.kv_item_lens_sum + total_bytes += self._transfer_num_state_indices * self.kv_mgr.state_item_lens_sum + self._transfer_metric.transfer_total_bytes = total_bytes + return self._transfer_metric + + def _record_transfer_indices(self, kv_indices, state_indices) -> None: + self._transfer_num_kv_indices += len(kv_indices) + if state_indices is not None: + self._transfer_num_state_indices += len(state_indices) + def init(self, num_kv_indices: int, aux_index: Optional[int] = None): self.num_kv_indices = num_kv_indices self.aux_index = aux_index diff --git a/python/sglang/srt/disaggregation/fake/conn.py b/python/sglang/srt/disaggregation/fake/conn.py index 4a3841e68..d03b9d382 100644 --- a/python/sglang/srt/disaggregation/fake/conn.py +++ b/python/sglang/srt/disaggregation/fake/conn.py @@ -10,6 +10,7 @@ from sglang.srt.disaggregation.base.conn import ( BaseKVSender, KVArgs, KVPoll, + KVTransferMetric, ) from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.server_args import ServerArgs @@ -52,6 +53,9 @@ class FakeKVSender(BaseKVSender): logger.debug("FakeKVSender poll success") return KVPoll.Success + def get_transfer_metric(self) -> KVTransferMetric: + return KVTransferMetric() + def init( self, kv_indices: list[int], diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 472b92671..0e08c1cb3 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1546,6 +1546,9 @@ class MooncakeKVSender(CommonKVSender): logical_page_positions=logical_page_positions, state_logical_page_positions=state_logical_page_positions, ) + # Record the actually-sent (post CP shared-KV filter) indices for accurate + # transfer metrics. Ported from upstream sgl-project/sglang #24416. + self._record_transfer_indices(kv_indices, state_indices) def poll(self) -> KVPoll: if self.conclude_state is None: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 6ab291196..95e20199c 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -1008,19 +1008,17 @@ 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 + if req.bootstrap_host == FAKE_BOOTSTRAP_HOST: + continue + kv_mgr = getattr(req.disagg_kv_sender, "kv_mgr", None) + if kv_mgr and getattr(kv_mgr, "is_dummy_cp_rank", False): + # Dummy CP ranks transfer nothing; skip so they don't pollute the metric. + 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, + req.disagg_kv_sender.get_transfer_metric() ) if metrics: # Update last-value for REST API diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py index 8caf21c32..26e55e723 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -37,6 +37,7 @@ from sglang.srt.observability.trace import ( from sglang.srt.utils import get_bool_env_var if TYPE_CHECKING: + from sglang.srt.disaggregation.base.conn import KVTransferMetric from sglang.srt.managers.schedule_batch import ScheduleBatch SGLANG_TEST_REQUEST_TIME_STATS = get_bool_env_var("SGLANG_TEST_REQUEST_TIME_STATS") @@ -772,27 +773,30 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): def compute_and_observe_kv_transfer_metrics( self, - num_tokens: int, - page_size: int, - bytes_per_page_all_layers: int, + transfer_metric: KVTransferMetric, ) -> 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 = {} + if transfer_metric.transfer_total_bytes is None: + return result if result else None # Transfer latency, size, and speed - if self.prefill_transfer_queue_entry_time > 0 and self.completion_time > 0: + if transfer_metric.transfer_latency_s is not None: + transfer_latency_s = transfer_metric.transfer_latency_s + else: + if self.prefill_transfer_queue_entry_time <= 0 or self.completion_time <= 0: + return result if result else None transfer_latency_s = ( self.completion_time - self.prefill_transfer_queue_entry_time ) + + if transfer_latency_s > 0: 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_bytes = transfer_metric.transfer_total_bytes total_mb = total_bytes / (1024 * 1024) self.transfer_total_mb = total_mb @@ -918,8 +922,12 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): def convert_to_duration(self) -> str: if self.disagg_mode == DisaggregationMode.NULL: - queue_duration = self.forward_entry_time - self.wait_queue_entry_time - forward_duration = self.completion_time - self.forward_entry_time + queue_duration = self.duration_between( + self.wait_queue_entry_time, self.forward_entry_time + ) + forward_duration = self.duration_between( + self.forward_entry_time, self.completion_time + ) if SGLANG_TEST_REQUEST_TIME_STATS: assert ( @@ -928,11 +936,15 @@ 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_queue_duration = ( - self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time + bootstrap_queue_duration = self.duration_between( + self.prefill_bootstrap_queue_entry_time, self.wait_queue_entry_time + ) + queue_duration = self.duration_between( + self.wait_queue_entry_time, self.forward_entry_time + ) + forward_duration = self.duration_between( + self.forward_entry_time, self.completion_time ) - queue_duration = self.forward_entry_time - self.wait_queue_entry_time - forward_duration = self.completion_time - self.forward_entry_time if SGLANG_TEST_REQUEST_TIME_STATS: if self.wait_queue_entry_time > 0: @@ -944,11 +956,11 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): # 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 + bootstrap_duration = self.duration_between( + self.prefill_bootstrap_queue_entry_time, self.bootstrap_done_time ) - alloc_wait_duration = ( - self.wait_queue_entry_time - self.bootstrap_done_time + alloc_wait_duration = self.duration_between( + self.bootstrap_done_time, self.wait_queue_entry_time ) if SGLANG_TEST_REQUEST_TIME_STATS: assert ( @@ -972,15 +984,22 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): f"#retries={self.prefill_retry_count}" ) elif self.disagg_mode == DisaggregationMode.DECODE: - prealloc_duration = ( - self.decode_transfer_queue_entry_time - - self.decode_prealloc_queue_entry_time + prealloc_duration = self.duration_between( + self.decode_prealloc_queue_entry_time, + self.decode_transfer_queue_entry_time, ) - transfer_duration = ( - self.wait_queue_entry_time - self.decode_transfer_queue_entry_time + transfer_duration = self.duration_between( + self.decode_transfer_queue_entry_time, + self.wait_queue_entry_time, + ) + queue_duration = self.duration_between( + self.wait_queue_entry_time, + self.forward_entry_time, + ) + forward_duration = self.duration_between( + self.forward_entry_time, + self.completion_time, ) - queue_duration = self.forward_entry_time - self.wait_queue_entry_time - forward_duration = self.completion_time - self.forward_entry_time if SGLANG_TEST_REQUEST_TIME_STATS: if self.wait_queue_entry_time > 0: @@ -993,11 +1012,11 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): # 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 + bootstrap_duration = self.duration_between( + self.decode_prealloc_queue_entry_time, self.bootstrap_done_time ) - alloc_wait_duration = ( - self.decode_transfer_queue_entry_time - self.bootstrap_done_time + alloc_wait_duration = self.duration_between( + self.bootstrap_done_time, self.decode_transfer_queue_entry_time ) if SGLANG_TEST_REQUEST_TIME_STATS: assert ( @@ -1043,6 +1062,13 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): def format_duration(self, duration: float) -> str: return f"{duration * 1e3:.2f}ms" + def duration_between(self, start: float, end: float) -> float: + # Guard against uninitialized (<=0) phase timestamps producing nonsensical + # negative durations. Ported from upstream sgl-project/sglang #24416. + if start <= 0 or end <= 0: + return 0.0 + return end - start + def set_schedule_time_batch(batch: ScheduleBatch): # only for tracing diff --git a/test/registered/unit/disaggregation/test_kv_transfer_metrics.py b/test/registered/unit/disaggregation/test_kv_transfer_metrics.py new file mode 100644 index 000000000..28a97d50c --- /dev/null +++ b/test/registered/unit/disaggregation/test_kv_transfer_metrics.py @@ -0,0 +1,122 @@ +"""Unit tests for the KV-transfer metric accounting (port of upstream #24416). + +These exercise the pure metric logic without standing up a real PD transfer: + - bytes are counted from the *actually-sent* indices (not prompt length), + - duration_between guards uninitialized (<=0) timestamps, + - compute_and_observe_kv_transfer_metrics consumes a KVTransferMetric. +""" + +import unittest + +import numpy as np + +from sglang.srt.disaggregation.base.conn import KVTransferMetric +from sglang.srt.disaggregation.common.conn import CommonKVSender +from sglang.srt.observability.req_time_stats import SchedulerReqTimeStats + + +class _FakeMgr: + def __init__(self, kv_sum, state_sum): + self.kv_item_lens_sum = kv_sum + self.state_item_lens_sum = state_sum + + +def _make_sender(kv_sum, state_sum): + s = CommonKVSender.__new__(CommonKVSender) # bypass heavy __init__ + s.kv_mgr = _FakeMgr(kv_sum, state_sum) + s._transfer_metric = KVTransferMetric() + s._transfer_num_kv_indices = 0 + s._transfer_num_state_indices = 0 + return s + + +def _make_stats(): + st = SchedulerReqTimeStats.__new__(SchedulerReqTimeStats) + st.enable_metrics = False + st.transfer_total_mb = 0.0 + st.transfer_speed_gb_s = 0.0 + st.prefill_transfer_queue_entry_time = 0.0 + st.completion_time = 0.0 + st.prefill_bootstrap_queue_entry_time = 0.0 + st.bootstrap_done_time = 0.0 + st.wait_queue_entry_time = 0.0 + return st + + +class TestKVTransferMetricBytes(unittest.TestCase): + def test_records_actually_sent_indices(self): + s = _make_sender(kv_sum=100, state_sum=10) + s._record_transfer_indices(np.arange(5, dtype=np.int32), [1, 2, 3]) + s._record_transfer_indices(np.arange(2, dtype=np.int32), None) + m = s.get_transfer_metric() + # 7 kv pages * 100 + 3 state pages * 10 + self.assertEqual(m.transfer_total_bytes, 7 * 100 + 3 * 10) + + def test_no_state_pool(self): + s = _make_sender(kv_sum=64, state_sum=0) + s._record_transfer_indices(np.arange(4, dtype=np.int32), None) + self.assertEqual(s.get_transfer_metric().transfer_total_bytes, 4 * 64) + + def test_bytes_track_sent_pages_not_prompt_length(self): + # The whole point of #24416: under cache hits / CP shared-KV filtering only + # a subset of pages is sent, and the metric must reflect that subset. + s = _make_sender(kv_sum=100, state_sum=0) + s._record_transfer_indices(np.arange(3, dtype=np.int32), None) + self.assertEqual(s.get_transfer_metric().transfer_total_bytes, 3 * 100) + + +class TestDurationBetween(unittest.TestCase): + def test_guards_nonpositive(self): + st = _make_stats() + self.assertEqual(st.duration_between(0, 5), 0.0) + self.assertEqual(st.duration_between(5, 0), 0.0) + self.assertEqual(st.duration_between(-1.0, 5.0), 0.0) + + def test_positive(self): + st = _make_stats() + self.assertAlmostEqual(st.duration_between(2.0, 5.0), 3.0) + + +class TestComputeAndObserve(unittest.TestCase): + def test_none_bytes_returns_none(self): + st = _make_stats() + m = KVTransferMetric(transfer_latency_s=1.0, transfer_total_bytes=None) + self.assertIsNone(st.compute_and_observe_kv_transfer_metrics(m)) + + def test_explicit_latency_and_bytes(self): + st = _make_stats() + mb = 1024 * 1024 + m = KVTransferMetric(transfer_latency_s=2.0, transfer_total_bytes=20 * mb) + out = st.compute_and_observe_kv_transfer_metrics(m) + self.assertAlmostEqual(out["total_mb"], 20.0) + self.assertAlmostEqual(out["latency_ms"], 2000.0) + self.assertAlmostEqual(out["speed_gb_s"], (20.0 / 1024) / 2.0) + + def test_latency_derived_from_timestamps(self): + st = _make_stats() + st.prefill_transfer_queue_entry_time = 1.0 + st.completion_time = 1.5 + m = KVTransferMetric(transfer_latency_s=None, transfer_total_bytes=1024 * 1024) + out = st.compute_and_observe_kv_transfer_metrics(m) + self.assertAlmostEqual(out["latency_ms"], 500.0) + + def test_missing_timestamps_returns_none(self): + st = _make_stats() # both timestamps are 0 + m = KVTransferMetric(transfer_latency_s=None, transfer_total_bytes=1024) + self.assertIsNone(st.compute_and_observe_kv_transfer_metrics(m)) + + +class TestSendersImplementMetric(unittest.TestCase): + def test_senders_have_get_transfer_metric(self): + from sglang.srt.disaggregation.fake.conn import FakeKVSender + from sglang.srt.disaggregation.mooncake.conn import MooncakeKVSender + + for cls in (CommonKVSender, FakeKVSender, MooncakeKVSender): + self.assertTrue( + callable(getattr(cls, "get_transfer_metric", None)), + f"{cls.__name__} missing get_transfer_metric", + ) + + +if __name__ == "__main__": + unittest.main()