diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index c2e354f2e..c5cb2bdba 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -29,6 +29,7 @@ from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( from sglang.srt.disaggregation.mooncake.conn import ( KVArgsRegisterInfo, MooncakeKVManager, + TransferInfo, ) from sglang.srt.disaggregation.utils import ( 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): @staticmethod def _arr(values): diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 2bc8a2107..60d5bf3e4 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -184,6 +184,105 @@ class TestNixlTransferInfo(CustomTestCase): 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): def test_from_zmq_preserves_unsigned_pointers_and_optional_fields(self):