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

[PD] Pack draft KV head slices for DCP transfers (#40500)

Browse Source 
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
This commit is contained in:
Khoa Pham
2026-09-21 21:11:27 -07:00
committed by GitHub
co-authored by Qiaolin Yu
parent b44e248682
commit 018b73c7a0
6 changed files with 273 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/kvcache/test_pd_dcp_gather.py
+134
View File
@@ -1,8 +1,13 @@
import concurrent.futures
import unittest
from types import SimpleNamespace
import numpy as np
import torch
from sglang.kernels.ops.kvcache.pd_dcp_gather import copy_mla_rows_into_pack
from sglang.srt.disaggregation.common.staging_buffer import StagingBuffer
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -36,6 +41,135 @@ class TestPdDcpGather(CustomTestCase):
torch.testing.assert_close(packed0, kv0[row_indices], rtol=0, atol=0)
torch.testing.assert_close(packed1, kv1[row_indices], rtol=0, atol=0)
def test_packed_tp2_pp2_to_dcp4_preserves_kv(self):
"""Packing must preserve both target rows and draft head shards across PP stages."""
for custom_pool in (False, True):
for capacity in (256 * (2 * 64 + 2 * 512), 256 * 2 * 64 // 4):
for rank in range(4):
with self.subTest(
custom_pool=custom_pool, capacity=capacity, rank=rank
):
self._check_packed_transfer(rank, custom_pool, capacity)
def _check_packed_transfer(self, rank, custom_pool, capacity):
page, tokens, chunk = 64, 521, 256
src_pages = np.array([7, 1, 9, 3, 4, 11, 2, 5, 8], dtype=np.int32)
dst_pages = np.array([4, 1, 6], dtype=np.int32)
layers, widths = [3, 11, 19, 27, 28, 28], [64] * 4 + [256] * 2
logical = torch.arange(tokens, device="cuda")
src_rows = (
torch.as_tensor(src_pages, device="cuda")[logical // page] * page
+ logical % page
)
values = [
(
(logical[:, None] + 256) * 13
+ torch.arange(width, device="cuda") * 7
+ entry * 31
)
.remainder(251)
.to(torch.uint8)
for entry, width in enumerate([64] * 4 + [1024] * 2)
]
destinations = [
torch.full((2048, w), 165, dtype=torch.uint8, device="cuda") for w in widths
]
expected = [x.clone() for x in destinations]
owned = logical[rank::4]
target_rows = (
torch.as_tensor(dst_pages, device="cuda")[owned // 256] * page
+ owned % 256 // 4
)
draft_rows = (
torch.as_tensor(dst_pages, device="cuda")[logical // 256] * 256
+ logical % 256
)
for entry in range(4):
expected[entry][target_rows] = values[entry][owned]
for entry in (4, 5):
expected[entry][draft_rows] = values[entry][
:, rank * 256 : (rank + 1) * 256
]
pack = StagingBuffer(capacity, "cuda:0", 0)
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
for stage, entries in enumerate(([0, 1], [2, 3, 4, 5])):
sources = []
for entry in entries:
data = values[entry]
if entry >= 4:
start = (rank // 2) * 512
data = data[:, start : start + 512]
source = torch.full(
(1024, data.shape[1]), 165, dtype=torch.uint8, device="cuda"
)
source[src_rows] = data
sources.append(source)
buffers = sources + destinations + [pack.buffer]
def transfer(session, blocks):
def view(ptr, size):
for tensor in buffers:
offset = ptr - tensor.data_ptr()
if 0 <= offset and offset + size <= tensor.numel():
return tensor.flatten()[offset : offset + size]
raise AssertionError(
f"Transfer outside registered buffers: {ptr}, {size}"
)
for src, dst, size in blocks:
view(dst, size).copy_(view(src, size))
torch.cuda.synchronize()
return 0
manager = SimpleNamespace(
is_mla_backend=False,
kv_args=SimpleNamespace(
page_size=page,
kv_layer_ids=[layers[e] for e in entries],
kv_data_ptrs=[x.data_ptr() for x in sources],
num_draft_entries=2 if stage else 0,
engine_rank=stage * 2 + rank // 2,
),
attn_tp_size=2,
max_transfer_batch_indices=37,
enable_custom_mem_pool=custom_pool,
enable_deferred_decode_kv_release=False,
_transfer_data=transfer,
)
manager._await_transfer_futures = lambda futures: (
MooncakeKVManager._await_transfer_futures(manager, futures)
)
for start in range(0, tokens, chunk):
count = min(chunk, tokens - start)
result = MooncakeKVManager.send_kvcache_dcp(
manager,
"session",
src_pages[start // page : (start + count + page - 1) // page],
[x.data_ptr() for x in destinations],
dst_pages,
dcp_token_item_lens=[x.shape[1] for x in sources],
dst_dcp_size=4,
dst_dcp_rank=rank,
src_page_offset=start // page,
decode_prefix_len=256,
num_kv_tokens=count,
executor=executor,
dst_layer_ids=layers,
pack_buffer=pack,
dst_kv_item_lens=[
page * w * (4 if e >= 4 else 1)
for e, w in enumerate(widths)
],
dst_tp_rank=rank,
dst_attn_tp_size=4,
)
self.assertEqual(result, 0)
for entry in entries:
torch.testing.assert_close(
destinations[entry], expected[entry], rtol=0, atol=0
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/disaggregation/test_mooncake_transfer_batching.py
+65 -1
View File
@@ -1,12 +1,14 @@
import concurrent.futures
import unittest
from threading import Event
from types import SimpleNamespace
from unittest.mock import MagicMock, call
from unittest.mock import MagicMock, call, patch
import numpy as np
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -286,5 +288,67 @@ class TestDcpDraftHeadTransfer(unittest.TestCase):
self.assertFalse(dst_buffers[1000000].any())
class TestDcpPackLifetime(CustomTestCase):
def test_failed_transfer_drains_before_pack_buffer_reuse(self):
"""A failed layer must not release the pack buffer while another transfer reads it."""
manager = TestMooncakeTransferBatching._make_manager(
enable_custom_mem_pool=True
)
manager.kv_args = SimpleNamespace(
page_size=1, kv_layer_ids=[], kv_data_ptrs=[1000, 2000], num_draft_entries=0
)
source = np.array([11], dtype=np.uint8)
observed = []
running, release = Event(), Event()
def transfer(session, blocks):
if blocks[0][0] == 1000:
self.assertTrue(running.wait(10))
return 17
running.set()
self.assertTrue(release.wait(10))
observed.append(int(source[0]))
return 0
def send(executor):
result = MooncakeKVManager.send_kvcache_dcp(
manager,
"session",
np.array([0, 1], dtype=np.int32),
[5000, 6000],
np.array([0], dtype=np.int32),
dcp_token_item_lens=[1, 1],
dst_dcp_size=2,
dst_dcp_rank=0,
src_page_offset=0,
decode_prefix_len=0,
num_kv_tokens=2,
executor=executor,
dst_layer_ids=[],
pack_buffer=object(),
)
source[0] = 22
return result
manager._transfer_data = transfer
with (
patch(
"sglang.srt.disaggregation.common.dcp_pack.try_pack_dcp_src",
return_value=([1000, 2000], np.array([0], dtype=np.int64)),
),
concurrent.futures.ThreadPoolExecutor(max_workers=2) as transfers,
concurrent.futures.ThreadPoolExecutor(max_workers=1) as worker,
):
future = worker.submit(send, transfers)
try:
self.assertTrue(running.wait(10))
with self.assertRaises(concurrent.futures.TimeoutError):
future.result(timeout=1)
finally:
release.set()
self.assertEqual(future.result(timeout=10), 17)
self.assertEqual(observed, [11])
if __name__ == "__main__":
unittest.main()