[PD] Unify dsv4 dispatch with swa (#24888)

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

@@ -31,7 +31,7 @@ class KVArgs:
     state_data_ptrs: List[int]
     state_data_lens: List[int]
     state_item_lens: List[int]
-    state_type: str  # "none", "mamba", "swa", "nsa", "dsv4"
+    state_type: str  # "none", "mamba", "swa", "nsa"
     # for mamba state different tp slice transfer
     state_dim_per_tensor: List[int]  # dimension to slice for each state tensor
     ib_device: str

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

@@ -995,13 +995,8 @@ class MooncakeKVManager(CommonKVManager):
                     prefill_state_indices,
                     dst_state_data_ptrs,
                 )
-        elif state_type in ["swa", "nsa", "dsv4"]:
-            # SWA / NSA / DSv4 hybrid models do not support different TP sizes
-            # yet. (DSv4 carries a flat heterogeneous state pool of
-            # SWA + compress + indexer buffers; reusing this branch routes it
-            # through the same ``_send_kvcache_generic`` path that
-            # ``get_mla_kv_ptrs_with_pp`` already handles for compressed-MLA
-            # PP/MTP layouts.)
+        elif state_type in ["swa", "nsa"]:
+            # Non-MLA SWA / NSA hybrid models do not support different TP sizes yet.
             if (
                 target_rank_registration_info is not None
                 and not self.is_mla_backend

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

@@ -1010,62 +1010,6 @@ class NixlKVManager(CommonKVManager):
             raise Exception("Failed to post Mamba state slice transfer")
         return xfer_handle
 
-    def _send_state_pages_flat(
-        self,
-        peer_name: str,
-        prefill_state_indices: List[int],
-        dst_state_data_ptrs: list[int],
-        dst_state_indices: List[int],
-        dst_state_item_lens: list[int],
-        dst_gpu_id: int,
-        notif: str,
-    ):
-        """Per-page WRITE transfer of a flat (heterogeneous) state pool.
-
-        Used by V4 whose state pool is a flat list of buffers (SWA + compress
-        + indexer pools) that does not match the per-layer K/V layout assumed
-        by ``_send_kvcache_generic``. Both sides must have identical
-        ``state_item_lens`` (no TP-slicing path).
-        """
-        src_state_ptrs = self.kv_args.state_data_ptrs
-        src_state_item_lens = self.kv_args.state_item_lens
-        assert len(src_state_ptrs) == len(dst_state_data_ptrs)
-        assert len(src_state_item_lens) == len(dst_state_item_lens)
-        assert len(prefill_state_indices) == len(dst_state_indices), (
-            f"State index length mismatch: prefill={len(prefill_state_indices)}, "
-            f"dst={len(dst_state_indices)}"
-        )
-        for i in range(len(src_state_item_lens)):
-            assert src_state_item_lens[i] == dst_state_item_lens[i], (
-                f"V4 state item length mismatch at index {i}: "
-                f"{src_state_item_lens[i]} != {dst_state_item_lens[i]}"
-            )
-
-        src_addrs = []
-        dst_addrs = []
-        for i in range(len(src_state_ptrs)):
-            item_len = src_state_item_lens[i]
-            for src_idx, dst_idx in zip(prefill_state_indices, dst_state_indices):
-                src_addr = src_state_ptrs[i] + int(src_idx) * item_len
-                dst_addr = dst_state_data_ptrs[i] + int(dst_idx) * item_len
-                src_addrs.append((src_addr, item_len, self.kv_args.gpu_id))
-                dst_addrs.append((dst_addr, item_len, dst_gpu_id))
-
-        if not src_addrs:
-            return None
-
-        src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
-        dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
-        xfer_handle = self.agent.initialize_xfer(
-            "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii")
-        )
-        if not xfer_handle:
-            raise Exception("KVSender failed to create state transfer")
-        state = self.agent.transfer(xfer_handle)
-        if state == "ERR":
-            raise Exception("KVSender failed to post state transfer")
-        return xfer_handle
-
     def maybe_send_extra(
         self,
         peer_name: str,
@@ -1124,16 +1068,6 @@ class NixlKVManager(CommonKVManager):
                 dst_gpu_id=dst_gpu_id,
                 notif=notif,
             )
-        elif state_type == "dsv4":
-            return self._send_state_pages_flat(
-                peer_name,
-                prefill_state_indices,
-                dst_state_data_ptrs,
-                dst_state_indices,
-                dst_state_item_lens or [],
-                dst_gpu_id,
-                notif,
-            )
         else:
             if state_type != "none":
                 raise RuntimeError(

python/sglang/srt/disaggregation/utils.py

@@ -543,7 +543,6 @@ def setup_state_kv_args(
     lives in one place.
     """
     from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
-    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
     from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool
 
     if not hasattr(token_to_kv_pool, "get_state_buf_infos"):
@@ -560,12 +559,9 @@ def setup_state_kv_args(
     kv_args.state_data_lens = state_data_lens
     kv_args.state_item_lens = state_item_lens
 
-    # V4 must be checked before BaseSWAKVPool: V4's state pool is a flat
-    # heterogeneous list (SWA + compress + indexer), so the per-layer K/V
-    # transfer path used for "swa"/"nsa" does not apply.
-    if isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool):
-        kv_args.state_type = "dsv4"
-    elif isinstance(token_to_kv_pool, BaseSWAKVPool):
+    # DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous
+    # state list is described per-entry via get_state_buf_infos.
+    if isinstance(token_to_kv_pool, BaseSWAKVPool):
         kv_args.state_type = "swa"
     elif isinstance(token_to_kv_pool, HybridLinearKVPool):
         kv_args.state_type = "mamba"