diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 5cf9d060890e..7f1c30c8e391 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/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, diff --git a/test/registered/unit/disaggregation/test_dcp_pack.py b/test/registered/unit/disaggregation/test_dcp_pack.py index dc5190d37224..c137b2726a33 100644 --- a/test/registered/unit/disaggregation/test_dcp_pack.py +++ b/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] diff --git a/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py b/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py index 75baa3f52abf..9710df60a9f1 100644 --- a/test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py +++ b/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)