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