[DSV4] Fix draft SWA transfer for disaggregated MTP (#30461)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

python/sglang/srt/disaggregation/utils.py

@@ -668,6 +668,7 @@ def setup_state_kv_args(
     from sglang.srt.disaggregation.base.conn import StateType
     from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
     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 (
         DSATokenToKVPool,
         HybridLinearKVPool,
@@ -757,6 +758,82 @@ def setup_state_kv_args(
                     kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
                 )
 
+    # DSV4 NextN shares the target allocator, so target and draft use the same
+    # local SWA indices. Keep draft buffers in a separate positional component
+    # to avoid mixing them into the target's heterogeneous state layout, while
+    # reusing the existing SWA transport dispatch. NPU has a different paged
+    # state layout and is intentionally left unchanged.
+    if (
+        not is_npu()
+        and isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+        and isinstance(draft_token_to_kv_pool, DeepSeekV4TokenToKVPool)
+    ):
+        if not draft_token_to_kv_pool.compression_ratios or not all(
+            ratio == 0 for ratio in draft_token_to_kv_pool.compression_ratios
+        ):
+            raise RuntimeError(
+                "DSV4 draft state transfer expects SWA-only NextN layers"
+            )
+        if token_to_kv_pool._unified_kv != draft_token_to_kv_pool._unified_kv:
+            raise RuntimeError(
+                "DSV4 target and draft pools must use the same unified-KV mode"
+            )
+
+        if token_to_kv_pool._unified_kv:
+            target_geometry = (
+                token_to_kv_pool.unified_swa_window,
+                token_to_kv_pool.unified_swa_ring_size,
+                token_to_kv_pool.unified_swa_pages,
+            )
+            draft_geometry = (
+                draft_token_to_kv_pool.unified_swa_window,
+                draft_token_to_kv_pool.unified_swa_ring_size,
+                draft_token_to_kv_pool.unified_swa_pages,
+            )
+            if target_geometry != draft_geometry:
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share SWA ring geometry: "
+                    f"target={target_geometry}, draft={draft_geometry}"
+                )
+            draft_ptrs, draft_lens, draft_item_lens = (
+                draft_token_to_kv_pool.get_unified_swa_ring_buf_infos()
+            )
+            draft_state_type = StateType.SWA_RING
+        else:
+            if (
+                token_to_kv_pool.full_to_swa_index_mapping
+                is not draft_token_to_kv_pool.full_to_swa_index_mapping
+            ):
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share the SWA index mapping"
+                )
+            target_geometry = (
+                token_to_kv_pool.page_size,
+                token_to_kv_pool.sliding_window,
+            )
+            draft_geometry = (
+                draft_token_to_kv_pool.page_size,
+                draft_token_to_kv_pool.sliding_window,
+            )
+            if target_geometry != draft_geometry:
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share paged SWA geometry: "
+                    f"target={target_geometry}, draft={draft_geometry}"
+                )
+            draft_ptrs, draft_lens, draft_item_lens = (
+                draft_token_to_kv_pool.get_state_buf_infos()
+            )
+            draft_state_type = StateType.SWA
+
+        if draft_ptrs:
+            append_state_component(
+                kv_args,
+                draft_state_type,
+                draft_ptrs,
+                draft_lens,
+                draft_item_lens,
+            )
+
     if (
         StateType.MAMBA not in kv_args.state_types
         and req_to_token_pool is not None

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -1,7 +1,10 @@
 import unittest
+from types import SimpleNamespace
 
 import numpy as np
+import torch
 
+from sglang.srt.disaggregation.base.conn import KVArgs, StateType
 from sglang.srt.disaggregation.common.utils import (
     group_concurrent_contiguous,
     pack_int_lists,
@@ -9,7 +12,11 @@ from sglang.srt.disaggregation.common.utils import (
     unpack_int_lists,
     unpack_list_of_buffers,
 )
-from sglang.srt.disaggregation.utils import get_dsv4_c128_state_indices
+from sglang.srt.disaggregation.utils import (
+    get_dsv4_c128_state_indices,
+    setup_state_kv_args,
+)
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=2, suite="base-a-test-cpu")
@@ -117,5 +124,87 @@ class TestDSV4C128StateIndices(unittest.TestCase):
         )
 
 
+def _buf_infos(*ptrs):
+    return list(ptrs), [ptr + 100 for ptr in ptrs], [ptr + 200 for ptr in ptrs]
+
+
+def _make_dsv4_target(*, unified, mapping=None):
+    pool = object.__new__(DeepSeekV4TokenToKVPool)
+    pool._unified_kv = unified
+    pool.page_size = 256
+    pool.sliding_window = 128
+    pool.full_to_swa_index_mapping = mapping
+    pool.unified_swa_window = 128
+    pool.unified_swa_ring_size = 131
+    pool.unified_swa_pages = 524
+    pool.get_state_buf_infos = lambda: _buf_infos(11)
+    pool.get_unified_swa_ring_buf_infos = lambda: (
+        _buf_infos(12) if unified else ([], [], [])
+    )
+    pool.get_c128_state_buf_infos = lambda: ([], [], [])
+    return pool
+
+
+def _make_dsv4_draft(*, unified, mapping=None):
+    pool = object.__new__(DeepSeekV4TokenToKVPool)
+    pool._unified_kv = unified
+    pool.compression_ratios = [0]
+    pool.page_size = 256
+    pool.sliding_window = 128
+    pool.full_to_swa_index_mapping = mapping
+    pool.unified_swa_window = 128
+    pool.unified_swa_ring_size = 131
+    pool.unified_swa_pages = 524
+    pool.compress_state_pools = [None]
+    pool.indexer_compress_state_pools = [None]
+    if unified:
+        pool.unified_kv_pool = SimpleNamespace(
+            swa_pages=524,
+            kv_buffer=[torch.empty((524, 16), dtype=torch.uint8)],
+        )
+    else:
+        pool.swa_kv_pool = SimpleNamespace(
+            kv_buffer=[torch.empty((2, 16), dtype=torch.uint8)]
+        )
+    return pool
+
+
+class TestDSV4DraftStateRegistration(unittest.TestCase):
+    def test_draft_state_is_a_separate_component(self):
+        mapping = torch.arange(16)
+        cases = [
+            (
+                "paged",
+                _make_dsv4_target(unified=False, mapping=mapping),
+                _make_dsv4_draft(unified=False, mapping=mapping),
+                [StateType.SWA, StateType.SWA],
+                [[11]],
+            ),
+            (
+                "unified",
+                _make_dsv4_target(unified=True),
+                _make_dsv4_draft(unified=True),
+                [StateType.SWA, StateType.SWA_RING, StateType.SWA_RING],
+                [[11], [12]],
+            ),
+        ]
+
+        for name, target, draft, expected_types, target_ptrs in cases:
+            with self.subTest(name=name):
+                if draft._unified_kv:
+                    expected_infos = draft.get_unified_swa_ring_buf_infos()
+                else:
+                    expected_infos = draft.get_state_buf_infos()
+                kv_args = KVArgs()
+
+                setup_state_kv_args(kv_args, target, draft)
+
+                self.assertEqual(kv_args.state_types, expected_types)
+                self.assertEqual(kv_args.state_data_ptrs[:-1], target_ptrs)
+                self.assertEqual(kv_args.state_data_ptrs[-1], expected_infos[0])
+                self.assertEqual(kv_args.state_data_lens[-1], expected_infos[1])
+                self.assertEqual(kv_args.state_item_lens[-1], expected_infos[2])
+
+
 if __name__ == "__main__":
     unittest.main()