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

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

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()