[Bug fix] Account for KV replication fan-out in transfer-byte metrics (#30351)

Co-authored-by: Hun-ho Kim <hunho.kim@samsung.com>

python/sglang/srt/disaggregation/common/conn.py

@@ -127,6 +127,11 @@ class CommonKVManager(BaseKVManager):
         self.kv_item_lens_sum = sum(args.kv_item_lens)
         self.state_item_lens_sum = sum(x for comp in args.state_item_lens for x in comp)
         self.is_mla_backend = is_mla_backend
+        # Per-sender fan-out of a KV copy onto N decode destinations
+        # (MLA under Prefill-CP + Decode-TP, or decode_tp > prefill_tp).
+        # MLA is resolved lazily at bootstrap (see resolve_kv_replica_factor);
+        # MHA never replicates, so it stays pinned at 1.
+        self._kv_replica_factor: Optional[int] = None if is_mla_backend else 1
         self.is_hybrid_mla_backend = getattr(args, "is_hybrid_mla_backend", False)
         self.disaggregation_mode = disaggregation_mode
         self.server_args = server_args
@@ -263,6 +268,30 @@ class CommonKVManager(BaseKVManager):
         with self.failure_lock:
             self.failure_records[bootstrap_room] = failure_reason
 
+    def get_kv_replica_factor(self) -> int:
+        if self._kv_replica_factor is None:
+            logger.warning_once(
+                "get_kv_replica_factor called before resolve_kv_replica_factor; "
+                "assuming 1, but the metrics may be inaccurate."
+            )
+            return 1
+        return self._kv_replica_factor
+
+    def resolve_kv_replica_factor(self, transfer_infos: Dict) -> None:
+        if not self.is_mla_backend:
+            # Only MLA replicates its KV across decode ranks; non-MLA head slices are
+            # disjoint and stay pinned at the factor of 1 set in __init__.
+            return
+
+        info = next(iter(transfer_infos.values()), None)
+        if info is None or info.required_dst_info_num is None:
+            logger.warning_once(
+                "resolve_kv_replica_factor: no decode destinations available; "
+                "KV transfer metrics may be inaccurate."
+            )
+            return
+        self._kv_replica_factor = info.required_dst_info_num
+
     def _ensure_prefill_recompute_executor(
         self,
     ) -> concurrent.futures.ThreadPoolExecutor:
@@ -1033,6 +1062,8 @@ class CommonKVSender(BaseKVSender):
         total_bytes += (
             self._transfer_num_state_indices * self.kv_mgr.state_item_lens_sum
         )
+        # Pinned to 1 for MHA (disjoint slices); only MLA replication makes it > 1.
+        total_bytes *= self.kv_mgr.get_kv_replica_factor()
         self._transfer_metric.transfer_total_bytes = total_bytes
         return self._transfer_metric
 

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -1560,6 +1560,7 @@ class MooncakeKVManager(CommonKVManager):
                     )
                     # NOTE: after bootstrapping we can mark the req as waiting for input
                     if len(self.transfer_infos[room]) == required_dst_info_num:
+                        self.resolve_kv_replica_factor(self.transfer_infos[room])
                         self.req_to_decode_prefix_len[room] = next(
                             (
                                 info.decode_prefix_len

python/sglang/srt/disaggregation/mori/conn.py

@@ -506,6 +506,7 @@ class MoriKVManager(CommonKVManager):
                 infos[transfer_info.engine_key] = transfer_info
 
                 if len(infos) >= transfer_info.required_dst_info_num:
+                    self.resolve_kv_replica_factor(infos)
                     # All decode peers reported their dst metadata; pick a
                     # non-None decode_prefix_len if any peer set it (they
                     # should all agree, but be defensive). 0 means "no

python/sglang/srt/disaggregation/nixl/conn.py

@@ -2412,6 +2412,7 @@ class NixlKVManager(CommonKVManager):
                 ].required_dst_info_num
                 logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}")
                 if len(self.transfer_infos[room]) == required_dst_info_num:
+                    self.resolve_kv_replica_factor(self.transfer_infos[room])
                     self.req_to_decode_prefix_len[room] = next(
                         (
                             info.decode_prefix_len