[PD] Add is_dummy truth-table wire tests for mooncake and nixl (#34977)
This commit is contained in:
@@ -29,6 +29,7 @@ from sglang.srt.disaggregation.decode_schedule_batch_mixin import (
|
|||||||
from sglang.srt.disaggregation.mooncake.conn import (
|
from sglang.srt.disaggregation.mooncake.conn import (
|
||||||
KVArgsRegisterInfo,
|
KVArgsRegisterInfo,
|
||||||
MooncakeKVManager,
|
MooncakeKVManager,
|
||||||
|
TransferInfo,
|
||||||
)
|
)
|
||||||
from sglang.srt.disaggregation.utils import (
|
from sglang.srt.disaggregation.utils import (
|
||||||
MetadataBuffers,
|
MetadataBuffers,
|
||||||
@@ -188,6 +189,65 @@ class TestCPReplicatedStateTransfer(unittest.TestCase):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class TestMooncakeTransferInfoIsDummy(unittest.TestCase):
|
||||||
|
"""Truth table for mooncake's payload-inferred is_dummy, with frames built
|
||||||
|
as KVSender sends them: kv and aux are empty iff dummy, state indices are
|
||||||
|
gated on dummy, decode_prefix_len and required_dst_info_num are sent
|
||||||
|
unconditionally."""
|
||||||
|
|
||||||
|
def _frames(self, kv, aux, state, prefix):
|
||||||
|
return [
|
||||||
|
b"7",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"1234",
|
||||||
|
b"session",
|
||||||
|
kv,
|
||||||
|
aux,
|
||||||
|
state,
|
||||||
|
b"1",
|
||||||
|
prefix,
|
||||||
|
b"",
|
||||||
|
]
|
||||||
|
|
||||||
|
def test_real_transfer_is_not_dummy(self):
|
||||||
|
kv = np.array([3, 5], dtype=np.int32)
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
self._frames(kv.tobytes(), b"4", pack_int_lists([[1]], "i"), b"0")
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(info.is_dummy)
|
||||||
|
np.testing.assert_array_equal(info.dst_kv_indices, kv)
|
||||||
|
self.assertEqual(info.dst_aux_index, 4)
|
||||||
|
self.assertEqual(info.dst_state_indices, [[1]])
|
||||||
|
|
||||||
|
def test_full_prefix_hit_with_empty_kv_is_not_dummy(self):
|
||||||
|
# Empty kv indices serialize to an empty frame, so only the non-empty
|
||||||
|
# aux frame distinguishes a full-prefix-hit transfer from a dummy one.
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
self._frames(np.array([], dtype=np.int32).tobytes(), b"4", b"", b"128")
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(info.is_dummy)
|
||||||
|
self.assertEqual(info.dst_aux_index, 4)
|
||||||
|
self.assertEqual(info.decode_prefix_len, 128)
|
||||||
|
|
||||||
|
def test_dummy_parses_dummy_and_clears_payload_fields(self):
|
||||||
|
info = TransferInfo.from_zmq(self._frames(b"", b"", b"", b"0"))
|
||||||
|
|
||||||
|
self.assertTrue(info.is_dummy)
|
||||||
|
self.assertEqual(info.dst_kv_indices.size, 0)
|
||||||
|
self.assertIsNone(info.dst_aux_index)
|
||||||
|
self.assertEqual(info.dst_state_indices, [])
|
||||||
|
|
||||||
|
def test_dummy_with_prefix_hit_still_parses_dummy(self):
|
||||||
|
# decode_prefix_len is sent unconditionally and the inference ignores
|
||||||
|
# it, so a dummy rank with a decode-side prefix hit stays dummy.
|
||||||
|
info = TransferInfo.from_zmq(self._frames(b"", b"", b"", b"128"))
|
||||||
|
|
||||||
|
self.assertTrue(info.is_dummy)
|
||||||
|
self.assertEqual(info.decode_prefix_len, 128)
|
||||||
|
|
||||||
|
|
||||||
class TestGroupConcurrentContiguous(unittest.TestCase):
|
class TestGroupConcurrentContiguous(unittest.TestCase):
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def _arr(values):
|
def _arr(values):
|
||||||
|
|||||||
@@ -184,6 +184,105 @@ class TestNixlTransferInfo(CustomTestCase):
|
|||||||
|
|
||||||
self.assertTrue(info.is_dummy())
|
self.assertTrue(info.is_dummy())
|
||||||
|
|
||||||
|
def test_explicit_dummy_frame_true_is_dummy(self):
|
||||||
|
# msg[9] is the explicit is_dummy frame the sender writes
|
||||||
|
# (str(int(is_dummy))); it wins over payload inference.
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
[
|
||||||
|
b"11",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"12349",
|
||||||
|
b"agent",
|
||||||
|
np.array([], dtype=np.int32).tobytes(),
|
||||||
|
b"2",
|
||||||
|
b"1",
|
||||||
|
b"",
|
||||||
|
b"0",
|
||||||
|
b"1",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(info.is_dummy())
|
||||||
|
|
||||||
|
def test_explicit_dummy_frame_true_with_prefix_hit_stays_dummy(self):
|
||||||
|
# A dummy rank whose request also has a decode-side prefix hit: the
|
||||||
|
# sender sends decode_prefix_len unconditionally, so only the explicit
|
||||||
|
# frame distinguishes this from a real full-prefix-hit transfer.
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
[
|
||||||
|
b"12",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"12350",
|
||||||
|
b"agent",
|
||||||
|
np.array([], dtype=np.int32).tobytes(),
|
||||||
|
b"2",
|
||||||
|
b"1",
|
||||||
|
b"",
|
||||||
|
b"128",
|
||||||
|
b"1",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertTrue(info.is_dummy())
|
||||||
|
|
||||||
|
def test_explicit_dummy_frame_false_with_empty_indices_is_real(self):
|
||||||
|
# Full prefix hit as the sender encodes it: empty kv indices,
|
||||||
|
# decode_prefix_len > 0, explicit is_dummy 0.
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
[
|
||||||
|
b"13",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"12351",
|
||||||
|
b"agent",
|
||||||
|
np.array([], dtype=np.int32).tobytes(),
|
||||||
|
b"2",
|
||||||
|
b"1",
|
||||||
|
b"",
|
||||||
|
b"128",
|
||||||
|
b"0",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(info.is_dummy())
|
||||||
|
|
||||||
|
def test_explicit_dummy_frame_false_for_real_transfer(self):
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
[
|
||||||
|
b"14",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"12352",
|
||||||
|
b"agent",
|
||||||
|
np.array([3, 5], dtype=np.int32).tobytes(),
|
||||||
|
b"2",
|
||||||
|
b"1",
|
||||||
|
b"",
|
||||||
|
b"0",
|
||||||
|
b"0",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(info.is_dummy())
|
||||||
|
|
||||||
|
def test_fallback_without_dummy_frame_reads_prefix_hit_dummy_as_real(self):
|
||||||
|
# Old-peer fallback: without msg[9], a dummy rank with a decode-side
|
||||||
|
# prefix hit is indistinguishable from a real full-prefix-hit transfer
|
||||||
|
# and parses as real. The explicit frame above exists for this case.
|
||||||
|
info = TransferInfo.from_zmq(
|
||||||
|
[
|
||||||
|
b"15",
|
||||||
|
b"127.0.0.1",
|
||||||
|
b"12353",
|
||||||
|
b"agent",
|
||||||
|
np.array([], dtype=np.int32).tobytes(),
|
||||||
|
b"2",
|
||||||
|
b"1",
|
||||||
|
b"",
|
||||||
|
b"128",
|
||||||
|
]
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(info.is_dummy())
|
||||||
|
|
||||||
|
|
||||||
class TestNixlKVArgsRegisterInfo(CustomTestCase):
|
class TestNixlKVArgsRegisterInfo(CustomTestCase):
|
||||||
def test_from_zmq_preserves_unsigned_pointers_and_optional_fields(self):
|
def test_from_zmq_preserves_unsigned_pointers_and_optional_fields(self):
|
||||||
|
|||||||
Reference in New Issue
Block a user