[PD] Add is_dummy truth-table wire tests for mooncake and nixl (#34977)

This commit is contained in:
Rishabh Sinha
2026-09-11 13:14:41 +08:00
committed by GitHub
parent 40a84d6dfc
commit 2018c64e22
2 changed files with 159 additions and 0 deletions
@@ -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):