From 1c9eb6bb0b02cf24f696213dd6f7644d419c1583 Mon Sep 17 00:00:00 2001 From: weireweire Date: Thu, 9 Jul 2026 13:18:34 +0800 Subject: [PATCH] [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 | 77 ++++++++++++++++ .../test_disaggregation_wire.py | 91 ++++++++++++++++++- 2 files changed, 167 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 9ffceda77..1ab8ec16b 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 76208a5cc..d9e5fbe81 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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()