[PD] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks (#35762)
Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
co-authored by
Cursor
Claude Opus 5
Cursor
Claude Opus 5
parent
5f216fc33f
commit
3760296be8
@@ -0,0 +1,41 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.kvcache.pd_dcp_gather import copy_mla_rows_into_pack
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestPdDcpGather(CustomTestCase):
|
||||
def test_gathers_strided_rows_layer_major(self):
|
||||
dim = 8
|
||||
kv0 = torch.arange(32 * dim, dtype=torch.float32, device="cuda").view(
|
||||
32, 1, dim
|
||||
)
|
||||
kv1 = torch.arange(32 * 5, dtype=torch.float16, device="cuda").view(32, 1, 5)
|
||||
row_indices = torch.tensor([0, 4, 9, 12], dtype=torch.int64, device="cuda")
|
||||
item_lens = [int(kv0[0].nbytes), int(kv1[0].nbytes)]
|
||||
pack = torch.zeros(
|
||||
row_indices.numel() * sum(item_lens), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
|
||||
copy_mla_rows_into_pack(
|
||||
[kv0.data_ptr(), kv1.data_ptr()],
|
||||
row_indices,
|
||||
pack,
|
||||
item_lens,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
split = row_indices.numel() * item_lens[0]
|
||||
packed0 = pack[:split].view(torch.float32).view(4, 1, dim)
|
||||
packed1 = pack[split:].view(torch.float16).view(4, 1, 5)
|
||||
torch.testing.assert_close(packed0, kv0[row_indices], rtol=0, atol=0)
|
||||
torch.testing.assert_close(packed1, kv1[row_indices], rtol=0, atol=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,120 @@
|
||||
import unittest
|
||||
from contextlib import nullcontext
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.common.dcp_pack import (
|
||||
dcp_pack_buffer_bytes,
|
||||
try_pack_dcp_src,
|
||||
)
|
||||
from sglang.srt.disaggregation.common.utils import (
|
||||
build_dcp_token_transfer_plan,
|
||||
group_concurrent_contiguous,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class TestPackedDcpGrouping(CustomTestCase):
|
||||
def test_packed_groups_collapse_cyclic_src(self):
|
||||
page_size = 64
|
||||
dcp_size = 4
|
||||
src_pages = np.arange(4, dtype=np.int32)
|
||||
dst_pages = np.array([7], dtype=np.int32)
|
||||
plan = build_dcp_token_transfer_plan(
|
||||
src_pages,
|
||||
dst_pages,
|
||||
physical_page_size=page_size,
|
||||
dcp_size=dcp_size,
|
||||
dcp_rank=0,
|
||||
num_kv_tokens=256,
|
||||
)
|
||||
raw_src, _ = group_concurrent_contiguous(
|
||||
plan.src_token_indices, plan.dst_token_indices
|
||||
)
|
||||
self.assertEqual(len(raw_src), 64)
|
||||
self.assertTrue(all(len(group) == 1 for group in raw_src))
|
||||
|
||||
packed_src = np.arange(plan.dst_token_indices.size, dtype=np.int64)
|
||||
packed_groups, _ = group_concurrent_contiguous(
|
||||
packed_src, plan.dst_token_indices
|
||||
)
|
||||
self.assertEqual(len(packed_groups), 1)
|
||||
self.assertEqual(len(packed_groups[0]), 64)
|
||||
|
||||
|
||||
class TestDcpPackBufferBytes(CustomTestCase):
|
||||
def test_sizes_fixed_regions_for_each_dcp_rank(self):
|
||||
self.assertEqual(
|
||||
dcp_pack_buffer_bytes(
|
||||
[64 * 16, 64 * 16],
|
||||
page_size=64,
|
||||
max_tokens=10,
|
||||
dcp_size=4,
|
||||
),
|
||||
4 * 3 * (16 + 16),
|
||||
)
|
||||
|
||||
def test_rejects_invalid_item_lens(self):
|
||||
with self.assertRaisesRegex(ValueError, "at least one page"):
|
||||
dcp_pack_buffer_bytes([0], page_size=64, max_tokens=8)
|
||||
with self.assertRaisesRegex(ValueError, "page-aligned"):
|
||||
dcp_pack_buffer_bytes([100], page_size=64, max_tokens=8)
|
||||
|
||||
|
||||
class TestTryDcpPack(CustomTestCase):
|
||||
def test_try_pack_uses_requested_region_and_dense_indices(self):
|
||||
dim = 4
|
||||
kv = torch.arange(16 * dim, dtype=torch.float32).view(16, 1, dim)
|
||||
item_len = int(kv[0].nbytes)
|
||||
pack = torch.zeros(8 * item_len, dtype=torch.uint8)
|
||||
gather_stream = Mock()
|
||||
buf = type(
|
||||
"Buf",
|
||||
(),
|
||||
{
|
||||
"buffer": pack,
|
||||
"fits": lambda self, n: n <= pack.numel(),
|
||||
"get_ptr": lambda self: 0x1000,
|
||||
"get_size": lambda self: pack.numel(),
|
||||
"get_gather_stream": lambda self: gather_stream,
|
||||
},
|
||||
)()
|
||||
src = np.array([1, 5, 9, 13], dtype=np.int64)
|
||||
pack_offset = 2 * item_len
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.disaggregation.common.dcp_pack.torch.cuda.default_stream"
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.disaggregation.common.dcp_pack.torch.cuda.stream",
|
||||
return_value=nullcontext(),
|
||||
),
|
||||
patch(
|
||||
"sglang.srt.disaggregation.common.dcp_pack.copy_mla_rows_into_pack"
|
||||
) as copy_mock,
|
||||
):
|
||||
packed = try_pack_dcp_src(
|
||||
pack_buffer=buf,
|
||||
kv_data_ptrs=[kv.data_ptr()],
|
||||
src_token_indices=src,
|
||||
token_item_lens=[item_len],
|
||||
pack_offset_bytes=pack_offset,
|
||||
)
|
||||
|
||||
gather_stream.synchronize.assert_called_once_with()
|
||||
self.assertIsNotNone(packed)
|
||||
ptrs, indices = packed
|
||||
self.assertEqual(ptrs, [0x1000 + pack_offset])
|
||||
np.testing.assert_array_equal(indices, np.arange(4))
|
||||
pack_view = copy_mock.call_args.args[2]
|
||||
self.assertEqual(pack_view.storage_offset(), pack_offset)
|
||||
self.assertEqual(pack_view.numel(), src.size * item_len)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -535,6 +535,92 @@ class TestNixlTransferWorker(CustomTestCase):
|
||||
self.assertIn(room, mgr.req_to_decode_prefix_len)
|
||||
mgr.send_kvcache.assert_called_once()
|
||||
|
||||
def test_dcp_destinations_use_disjoint_pack_regions_before_chunk_barrier(self):
|
||||
room = 23
|
||||
mgr = self._make_manager(room)
|
||||
agents = ("agent0a", "agent0b", "agent1")
|
||||
dcp_ranks = (0, 0, 1)
|
||||
mgr.transfer_infos[room] = {
|
||||
agent: TransferInfo(
|
||||
room=room,
|
||||
endpoint="127.0.0.1",
|
||||
dst_port=5555 + i,
|
||||
agent_name=agent,
|
||||
dst_kv_indices=np.array([2 + i], dtype=np.int32),
|
||||
dst_aux_index=0,
|
||||
required_dst_info_num=len(agents),
|
||||
dst_state_indices=[],
|
||||
)
|
||||
for i, agent in enumerate(agents)
|
||||
}
|
||||
mgr.decode_kv_args_table = {
|
||||
agent: SimpleNamespace(
|
||||
decode_tp_size=len(agents),
|
||||
dst_kv_ptrs=[0x3000 + i * 0x100],
|
||||
dst_aux_ptrs=[0],
|
||||
gpu_id=0,
|
||||
staging_base_ptr=0,
|
||||
staging_total_size=0,
|
||||
kv_xfer_segments=None,
|
||||
dst_homogeneous_mem_kind="VRAM",
|
||||
requires_dcp_relayout=True,
|
||||
dst_dcp_size=2,
|
||||
dst_dcp_rank=dcp_rank,
|
||||
dcp_dst_region_indices=[0],
|
||||
dcp_token_item_lens=[4],
|
||||
)
|
||||
for i, (agent, dcp_rank) in enumerate(zip(agents, dcp_ranks))
|
||||
}
|
||||
mgr.kv_args = SimpleNamespace(
|
||||
engine_rank=0,
|
||||
kv_data_ptrs=[0x1000],
|
||||
page_size=4,
|
||||
)
|
||||
mgr._dcp_pack_buffers = [SimpleNamespace(get_size=lambda: 16)]
|
||||
|
||||
packed_rank0 = ([0x9000], np.arange(2, dtype=np.int64))
|
||||
packed_rank1 = ([0x9008], np.arange(2, dtype=np.int64))
|
||||
try_pack = MagicMock(side_effect=[packed_rank0, packed_rank1])
|
||||
dcp_pack_module = types.ModuleType("sglang.srt.disaggregation.common.dcp_pack")
|
||||
dcp_pack_module.try_pack_dcp_src = try_pack
|
||||
submitted = []
|
||||
|
||||
def send_kvcache_dcp(*args, **kwargs):
|
||||
submitted.append((args[0], args[-1]))
|
||||
return f"handle-{args[0]}"
|
||||
|
||||
mgr.send_kvcache_dcp = MagicMock(side_effect=send_kvcache_dcp)
|
||||
submitted_counts_at_poll = []
|
||||
|
||||
def check_xfer_state(_handle):
|
||||
submitted_counts_at_poll.append(len(submitted))
|
||||
return "DONE"
|
||||
|
||||
mgr.agent = SimpleNamespace(check_xfer_state=check_xfer_state)
|
||||
chunk = self._make_chunk(room, [1], is_last_chunk=False)
|
||||
chunk.num_kv_tokens = 4
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{"sglang.srt.disaggregation.common.dcp_pack": dcp_pack_module},
|
||||
):
|
||||
self._run_worker_once(mgr, chunk)
|
||||
|
||||
self.assertEqual(try_pack.call_count, 2)
|
||||
self.assertEqual(
|
||||
[call.kwargs["pack_offset_bytes"] for call in try_pack.call_args_list],
|
||||
[0, 8],
|
||||
)
|
||||
self.assertEqual(
|
||||
submitted,
|
||||
[
|
||||
("agent0a", packed_rank0),
|
||||
("agent0b", packed_rank0),
|
||||
("agent1", packed_rank1),
|
||||
],
|
||||
)
|
||||
self.assertEqual(submitted_counts_at_poll, [3, 3, 3])
|
||||
|
||||
|
||||
class TestNixlNotifications(CustomTestCase):
|
||||
def _make_manager(self, messages, required=None):
|
||||
@@ -788,16 +874,16 @@ class TestNixlStaging(CustomTestCase):
|
||||
agent = StagingFakeAgent(register_result=["staging"])
|
||||
mgr = self._make_manager(agent)
|
||||
|
||||
mgr._register_staging_memory(0x1000, 4096, 3)
|
||||
mgr._register_staging_memory(0x1000, 4096)
|
||||
|
||||
self.assertEqual(
|
||||
agent.register_memory_calls,
|
||||
[([(0x1000, 4096, 3, "")], "VRAM")],
|
||||
[([(0x1000, 4096, 1, "")], "VRAM")],
|
||||
)
|
||||
|
||||
mgr = self._make_manager(StagingFakeAgent(register_result=[]))
|
||||
with self.assertRaisesRegex(RuntimeError, "staging buffer"):
|
||||
mgr._register_staging_memory(0x1000, 4096, 3)
|
||||
mgr._register_staging_memory(0x1000, 4096)
|
||||
|
||||
def test_prefetch_staging_reqs_noops_when_disabled_or_missing_kv_buffers(self):
|
||||
mgr = self._make_manager()
|
||||
|
||||
Reference in New Issue
Block a user