minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
weireweire
2026-07-09 13:18:34 +08:00
committed by GitHub
co-authored by weireweire
parent 69ddbf9ef6
commit 1c9eb6bb0b
2 changed files with 167 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/utils.py
+77
View File
@@ -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
+90 -1
View File
@@ -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()