fix(disagg): account KV transfer metrics by actually-sent bytes

Port upstream sgl-project/sglang #24416 (staging-free). Our transfer metrics
were computed from the full prompt length, so total_mb / speed_gb_s were
systematically over-reported on every prefix-cache hit and every CP shared-KV
per-rank page filter. Now each sender accumulates the actually-sent KV/state
indices and reports bytes = sent_pages * per-page item bytes via a new
KVTransferMetric returned by get_transfer_metric(); prefill consumes that
instead of estimating from len(origin_input_ids), and skips fake-bootstrap and
dummy-CP-rank senders (which transfer nothing).

Also route convert_to_duration() phase deltas through duration_between(), which
returns 0 when a phase timestamp is uninitialized, fixing nonsensical negative
durations in the time-stats log.

Unit-tested (test_kv_transfer_metrics.py, 10 cases) in the CUDA-13 container.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

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

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

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:

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

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

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