[PD] Count transferred state bytes per component

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

@@ -170,7 +170,7 @@ def __init__(
         self.kv_cache_dtype_str = args.kv_cache_dtype_str
         self.dsv41_spec_layout = get_dsv41_spec_layout(args)
         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.state_item_lens_sums = [sum(lens) for lens in args.state_item_lens]
         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).
@@ -697,6 +697,16 @@ def register_deferred_ack_target(
         worker ack, then a new chunk writes pages the decode already released."""
         self._deferred_ack_targets[room] = (decode_ip, decode_port)
 
+    def get_state_transfer_bytes(self, component: int, indices) -> int:
+        row_bytes = self.state_item_lens_sums[component]
+        if self.kv_args.state_types[component] == StateType.DSA_TAIL:
+            # indices = [req_pool_idx, start, first_n, 0, second_n, tail_size];
+            # only the live slots of the one ring row move.
+            if len(indices) == 0:
+                return 0
+            return (indices[2] + indices[4]) * row_bytes // indices[5]
+        return len(indices) * row_bytes
+
     def get_kv_replica_factor(self) -> int:
         if self._kv_replica_factor is None:
             logger.warning_once(
@@ -1494,7 +1504,7 @@ def __init__(
         self.conclude_state: Optional[KVPoll] = None
         self._transfer_metric = KVTransferMetric()
         self._transfer_num_kv_indices = 0
-        self._transfer_num_state_indices = 0
+        self._transfer_state_bytes = 0
         # inner state
         self.curr_idx = 0
         self.init_time: Optional[float] = None
@@ -1570,9 +1580,7 @@ def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
 
     def get_transfer_metric(self) -> KVTransferMetric:
         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
-        )
+        total_bytes += self._transfer_state_bytes
         # 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
@@ -1585,9 +1593,12 @@ def _record_transfer_indices(
     ):
         self._transfer_num_kv_indices += len(kv_indices)
         if state_indices:
-            for component_indices in state_indices:
+            for component, component_indices in enumerate(state_indices):
                 if component_indices is not None:
-                    self._transfer_num_state_indices += len(component_indices)
+                    # Slot sizes differ per component (a Mamba slot vs an SWA page).
+                    self._transfer_state_bytes += self.kv_mgr.get_state_transfer_bytes(
+                        component, component_indices
+                    )
 
     def _prepare_send_indices(
         self,

test/registered/unit/disaggregation/test_dcp_pack.py

@@ -202,6 +202,7 @@ def test_cached_prefix_fits_pack_capacity_and_preserves_pages_and_state(self):
             gpu_id=0,
             state_types=[StateType.MAMBA],
         )
+        mgr.state_item_lens_sums = [0]
         mgr._dcp_pack_buffers = None
         mgr._dcp_pack_max_tokens = None
         mgr.transfer_queues = [None]

test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py

@@ -12,8 +12,12 @@
 
 import numpy as np
 
-from sglang.srt.disaggregation.base.conn import KVTransferMetric
+from sglang.srt.disaggregation.base.conn import KVTransferMetric, StateType
 from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVSender
+from sglang.srt.disaggregation.utils import (
+    build_dsa_tail_transfer_blocks,
+    get_dsa_tail_state_indices,
+)
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -40,11 +44,18 @@ def _make_kv_mgr(is_mla_backend):
     mgr = CommonKVManager.__new__(CommonKVManager)
     mgr.is_mla_backend = is_mla_backend
     mgr.kv_item_lens_sum = KV_ITEM_LENS_SUM
-    mgr.state_item_lens_sum = STATE_ITEM_LENS_SUM
+    _set_state_components(mgr, [StateType.MAMBA], [[STATE_ITEM_LENS_SUM]])
     mgr._kv_replica_factor = None if is_mla_backend else 1
     return mgr
 
 
+def _set_state_components(mgr, state_types, state_item_lens):
+    mgr.kv_args = SimpleNamespace(
+        state_types=state_types, state_item_lens=state_item_lens
+    )
+    mgr.state_item_lens_sums = [sum(lens) for lens in state_item_lens]
+
+
 class _ConcreteKVSender(CommonKVSender):
     """CommonKVSender is abstract: `poll` and `failure_exception` are backend
     duties, so it cannot be instantiated (not even via __new__, which enforces
@@ -64,7 +75,7 @@ def _make_sender(kv_mgr):
     sender = _ConcreteKVSender.__new__(_ConcreteKVSender)
     sender._transfer_metric = KVTransferMetric()
     sender._transfer_num_kv_indices = 0
-    sender._transfer_num_state_indices = 0
+    sender._transfer_state_bytes = 0
     sender.kv_mgr = kv_mgr
     return sender
 
@@ -84,6 +95,46 @@ def test_mla_scales_kv_and_state_bytes_by_fan_out(self):
         expected = (8 * KV_ITEM_LENS_SUM + 5 * STATE_ITEM_LENS_SUM) * 4
         self.assertEqual(sender.get_transfer_metric().transfer_total_bytes, expected)
 
+    def test_components_use_their_own_slot_sizes(self):
+        mgr = _make_kv_mgr(is_mla_backend=True)
+        _set_state_components(mgr, [StateType.MAMBA, StateType.SWA], [[3, 4], [10, 20]])
+        sender = _make_sender(mgr)
+        mgr.resolve_kv_replica_factor(_room(4))
+        sender._record_transfer_indices(
+            np.arange(8, dtype=np.int32),
+            [np.arange(5, dtype=np.int32), np.arange(2, dtype=np.int32)],
+        )
+        expected = (8 * KV_ITEM_LENS_SUM + 5 * 7 + 2 * 30) * 4
+        self.assertEqual(sender.get_transfer_metric().transfer_total_bytes, expected)
+
+    def test_dsa_tail_counts_live_slots_not_descriptor_length(self):
+        """A DSA tail index is a 6-int descriptor of the live slots in one ring
+        row; counting it as 6 whole rows overstated the bytes on the wire."""
+        tail_size, pool_size = 6, 4
+        row_lens = [tail_size * 128, tail_size * 128, tail_size * 4, tail_size * 4]
+        mgr = _make_kv_mgr(is_mla_backend=False)
+        _set_state_components(mgr, [StateType.DSA_TAIL], [row_lens])
+        pool = SimpleNamespace(
+            kpool_use_compress=True,
+            index_kpool=pool_size,
+            tail_extra_slots=tail_size - pool_size,
+        )
+        # 7: 3 live slots split across the ring end; 13: 1 slot; 16: none live.
+        for seq_len in (7, 13, 16):
+            sender = _make_sender(mgr)
+            indices = get_dsa_tail_state_indices(pool, 3, seq_len)
+            src_ptrs = [(i + 1) << 20 for i in range(len(row_lens))]
+            wire_bytes = sum(
+                n
+                for _, _, n in build_dsa_tail_transfer_blocks(
+                    src_ptrs, row_lens, src_ptrs, indices, list(indices)
+                )
+            )
+            sender._record_transfer_indices(np.arange(0, dtype=np.int32), [indices])
+            self.assertEqual(
+                sender.get_transfer_metric().transfer_total_bytes, wire_bytes
+            )
+
     def test_non_mla_factor_is_one_regardless_of_destinations(self):
         # Non-MLA head slices sum to one logical copy: factor stays 1.
         mgr = _make_kv_mgr(is_mla_backend=False)