[Spec][1/N] Decoupled speculative decoding: IPC protocol + cross-process request id + server flags (#27634)

test/registered/unit/server_args/test_server_args.py

@@ -518,6 +518,52 @@ class TestPortArgs(unittest.TestCase):
         self.assertTrue(port_args.detokenizer_ipc_name.startswith("ipc://"))
         self.assertIsInstance(port_args.nccl_port, int)
 
+    @patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
+    def test_init_new_builds_decoupled_spec_ipc_config(self, mock_temp_file):
+        mock_temp_file.return_value.name = "temp_file"
+
+        server_args = ServerArgs(model_path="dummy")
+        server_args.nccl_port = None
+        server_args.enable_dp_attention = False
+        server_args.decoupled_spec_role = "verifier"
+        server_args.decoupled_spec_bind_endpoint = "ipc:///tmp/v"
+        server_args.decoupled_spec_connect_endpoints = ["ipc:///tmp/d"]
+        server_args.decoupled_spec_rank = 0
+
+        port_args = PortArgs.init_new(server_args)
+
+        self.assertIsNotNone(port_args.decoupled_spec_ipc_config)
+        self.assertEqual(port_args.decoupled_spec_ipc_config.rank, 0)
+        self.assertEqual(
+            port_args.decoupled_spec_ipc_config.bind_endpoint, "ipc:///tmp/v"
+        )
+        self.assertEqual(
+            port_args.decoupled_spec_ipc_config.connect_endpoints, ("ipc:///tmp/d",)
+        )
+
+    @patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
+    def test_init_new_no_decoupled_config_when_role_null(self, mock_temp_file):
+        mock_temp_file.return_value.name = "temp_file"
+
+        server_args = ServerArgs(model_path="dummy")
+        server_args.nccl_port = None
+        server_args.enable_dp_attention = False
+        # decoupled_spec_role defaults to "null"
+
+        port_args = PortArgs.init_new(server_args)
+
+        self.assertIsNone(port_args.decoupled_spec_ipc_config)
+
+    def test_init_new_decoupled_role_requires_endpoints(self):
+        server_args = ServerArgs(model_path="dummy")
+        server_args.nccl_port = None
+        server_args.enable_dp_attention = False
+        server_args.decoupled_spec_role = "drafter"
+        # endpoints intentionally left as their None defaults
+
+        with self.assertRaises(ValueError):
+            PortArgs.init_new(server_args)
+
     def test_init_new_with_single_node_dp_attention(self):
 
         server_args = ServerArgs(model_path="dummy")
@@ -919,6 +965,54 @@ class TestNgramExternalSamArgs(CustomTestCase):
         self.assertIn("external-corpus-max-tokens", str(context.exception))
 
 
+class TestDecoupledSpecArgs(CustomTestCase):
+    """Decoupled speculative-decoding CLI flags.
+
+    These flags are auto-derived from the ``A[...]`` field metadata on
+    ``ServerArgs``; a bare annotation is silently skipped by
+    ``add_cli_args_from_dataclass``. This guards against the regression where
+    the flags went missing (e.g. after rebasing onto the auto-gen
+    ``add_cli_args``), which the direct-attribute ``PortArgs`` tests cannot
+    catch because they never exercise the CLI.
+    """
+
+    def test_decoupled_spec_cli_flags_round_trip(self):
+        server_args = prepare_server_args(
+            [
+                "--model-path",
+                "dummy",
+                "--decoupled-spec-role",
+                "verifier",
+                "--decoupled-spec-bind-endpoint",
+                "ipc:///tmp/v",
+                "--decoupled-spec-connect-endpoints",
+                '["ipc:///tmp/d"]',
+                "--decoupled-spec-rank",
+                "0",
+                "--spec-trace-dir",
+                "/tmp/tr",
+            ]
+        )
+        self.assertEqual(server_args.decoupled_spec_role, "verifier")
+        self.assertEqual(server_args.decoupled_spec_bind_endpoint, "ipc:///tmp/v")
+        self.assertEqual(server_args.decoupled_spec_connect_endpoints, ["ipc:///tmp/d"])
+        self.assertEqual(server_args.decoupled_spec_rank, 0)
+        self.assertEqual(server_args.spec_trace_dir, "/tmp/tr")
+
+    def test_decoupled_spec_role_defaults_to_null(self):
+        server_args = prepare_server_args(["--model-path", "dummy"])
+        self.assertEqual(server_args.decoupled_spec_role, "null")
+        self.assertIsNone(server_args.decoupled_spec_bind_endpoint)
+        self.assertIsNone(server_args.decoupled_spec_connect_endpoints)
+        self.assertIsNone(server_args.decoupled_spec_rank)
+
+    def test_decoupled_spec_role_rejects_invalid_choice(self):
+        with self.assertRaises(SystemExit):
+            prepare_server_args(
+                ["--model-path", "dummy", "--decoupled-spec-role", "bogus"]
+            )
+
+
 class TestAdaptiveSpecArgs(CustomTestCase):
     def test_adaptive_defaults_to_config_step_when_spec_params_omitted(self):
         with tempfile.NamedTemporaryFile("w", suffix=".json") as f:

test/registered/unit/spec/test_decoupled_spec_io.py

@@ -0,0 +1,272 @@
+"""Unit tests for srt/speculative/decoupled_spec_io.
+
+decoupled_spec_io is the schema-only IPC layer for decoupled speculative
+decoding: protocol message dataclasses, the cross-process request id codec, and
+the drafter-side reconciliation helpers. These tests drive the real logic (id
+round-trip + parse errors, commit validation, segment coalescing / contiguity /
+prefix extraction, and inbox routing) on CPU; there is no GPU or transport here.
+"""
+
+import unittest
+
+from sglang.srt.speculative.decoupled_spec_io import (
+    DraftClose,
+    DraftControlBatch,
+    DraftMeshMessage,
+    DraftMeshMessageType,
+    DraftReqKey,
+    DraftSync,
+    VerifierCommitSegment,
+    VerifyCommit,
+    build_draft_scheduler_rid,
+    parse_draft_scheduler_rid,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=6, suite="base-a-test-cpu")
+
+
+def _commit(rid, *, pre, tokens, src_verifier_rank=0, drafter_rank=0) -> VerifyCommit:
+    return VerifyCommit(
+        request_id=rid,
+        src_verifier_rank=src_verifier_rank,
+        dst_drafter_rank=drafter_rank,
+        pre_verify_committed_len=pre,
+        committed_tokens=list(tokens),
+    )
+
+
+def _segment(
+    rid, *, pre=0, drafter_rank=0, src_verifier_rank=0
+) -> VerifierCommitSegment:
+    return VerifierCommitSegment(
+        draft_key=DraftReqKey(src_verifier_rank=src_verifier_rank, request_id=rid),
+        dst_drafter_rank=drafter_rank,
+        pre_verify_committed_len=pre,
+    )
+
+
+class TestDraftSchedulerRid(CustomTestCase):
+    def test_round_trip(self):
+        key = DraftReqKey(src_verifier_rank=3, request_id="req-1")
+        rid = build_draft_scheduler_rid(key)
+        self.assertEqual(rid, "draft:3:req-1")
+        self.assertEqual(parse_draft_scheduler_rid(rid), key)
+
+    def test_request_id_containing_colon_round_trips(self):
+        # request_id may contain ':' — parse splits on the first ':' only.
+        key = DraftReqKey(src_verifier_rank=1, request_id="a:b:c")
+        self.assertEqual(parse_draft_scheduler_rid(build_draft_scheduler_rid(key)), key)
+
+    def test_parse_invalid_rid_raises(self):
+        for bad in ["no-prefix", "draft:", "draft:0:", "draft:notint:r"]:
+            with self.subTest(rid=bad):
+                with self.assertRaises(ValueError):
+                    parse_draft_scheduler_rid(bad)
+
+
+class TestVerifyCommitValidation(CustomTestCase):
+    def test_empty_tokens_raises(self):
+        with self.assertRaises(ValueError):
+            _commit("r", pre=0, tokens=[]).validate_committed_tokens()
+
+    def test_negative_pre_len_raises(self):
+        with self.assertRaises(ValueError):
+            _commit("r", pre=-1, tokens=[1]).validate_committed_tokens()
+
+    def test_valid_commit_passes(self):
+        # Should not raise.
+        _commit("r", pre=0, tokens=[1, 2]).validate_committed_tokens()
+
+
+class TestVerifierCommitSegment(CustomTestCase):
+    def test_append_coalesces_contiguous_commits(self):
+        seg = _segment("r", pre=0)
+        seg.append_message(_commit("r", pre=0, tokens=[10, 11]))
+        self.assertEqual(seg.committed_tokens, [10, 11])
+        self.assertEqual(seg.end_committed_len, 2)
+        seg.append_message(_commit("r", pre=2, tokens=[12]))
+        self.assertEqual(seg.committed_tokens, [10, 11, 12])
+        self.assertEqual(seg.end_committed_len, 3)
+
+    def test_append_wrong_request_raises(self):
+        seg = _segment("r", pre=0)
+        with self.assertRaises(RuntimeError):
+            seg.append_message(_commit("other", pre=0, tokens=[1]))
+
+    def test_append_wrong_drafter_rank_raises(self):
+        seg = _segment("r", pre=0, drafter_rank=0)
+        with self.assertRaises(RuntimeError):
+            seg.append_message(_commit("r", pre=0, tokens=[1], drafter_rank=7))
+
+    def test_append_non_contiguous_raises(self):
+        seg = _segment("r", pre=0)
+        seg.append_message(_commit("r", pre=0, tokens=[10]))  # end -> 1
+        with self.assertRaises(RuntimeError):
+            seg.append_message(_commit("r", pre=5, tokens=[11]))  # gap
+
+    def test_append_runs_message_validation(self):
+        seg = _segment("r", pre=0)
+        with self.assertRaises(ValueError):
+            seg.append_message(_commit("r", pre=0, tokens=[]))  # empty -> validate
+
+    def test_extract_prefix_splits_segment(self):
+        seg = _segment("r", pre=0)
+        seg.append_message(_commit("r", pre=0, tokens=[10, 11, 12, 13]))
+        prefix = seg.extract_prefix(2)
+        self.assertEqual(prefix.committed_tokens, [10, 11])
+        self.assertEqual(prefix.pre_verify_committed_len, 0)
+        # Remainder stays in the original segment, with pre advanced by 2.
+        self.assertEqual(seg.committed_tokens, [12, 13])
+        self.assertEqual(seg.pre_verify_committed_len, 2)
+        self.assertEqual(seg.end_committed_len, 4)
+
+    def test_extract_prefix_bounds(self):
+        seg = _segment("r", pre=0)
+        seg.append_message(_commit("r", pre=0, tokens=[10, 11]))
+        with self.assertRaises(ValueError):
+            seg.extract_prefix(0)
+        with self.assertRaises(ValueError):
+            seg.extract_prefix(3)  # exceeds segment length
+
+
+class TestDraftControlInbox(CustomTestCase):
+    def _inbox(self):
+        from sglang.srt.speculative.decoupled_spec_io import DraftControlInbox
+
+        return DraftControlInbox()
+
+    def _sync(self, rid, drafter_rank=0):
+        return DraftSync(
+            request_id=rid, src_verifier_rank=0, dst_drafter_rank=drafter_rank
+        )
+
+    def _close(self, rid, drafter_rank=0):
+        return DraftClose(
+            request_id=rid,
+            src_verifier_rank=0,
+            dst_drafter_rank=drafter_rank,
+            reason="x",
+        )
+
+    def test_add_control_batch_routes_each_message_type(self):
+        inbox = self._inbox()
+        inbox.add_control_batch_locked(
+            DraftControlBatch(
+                dst_drafter_rank=0,
+                sync_messages=[self._sync("s")],
+                verify_commit_messages=[_commit("c", pre=0, tokens=[1])],
+                close_messages=[self._close("x")],
+            )
+        )
+        self.assertEqual([m.request_id for m in inbox.sync_messages], ["s"])
+        self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments)
+        self.assertIn(DraftReqKey(0, "x"), inbox.close_keys)
+
+    def test_close_drops_pending_segment_and_sync(self):
+        inbox = self._inbox()
+        inbox.add_control_batch_locked(
+            DraftControlBatch(
+                dst_drafter_rank=0,
+                sync_messages=[self._sync("r")],
+                verify_commit_messages=[_commit("r", pre=0, tokens=[1])],
+            )
+        )
+        inbox.add_close_key_locked(DraftReqKey(0, "r"))
+        self.assertEqual(inbox.sync_messages, [])
+        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)
+        self.assertIn(DraftReqKey(0, "r"), inbox.close_keys)
+
+    def test_verify_commit_for_closed_key_is_ignored(self):
+        inbox = self._inbox()
+        inbox.add_close_key_locked(DraftReqKey(0, "r"))
+        inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1]))
+        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)
+
+    def test_extract_ready_controls_full_consume(self):
+        inbox = self._inbox()
+        inbox.add_control_batch_locked(
+            DraftControlBatch(
+                dst_drafter_rank=0,
+                sync_messages=[self._sync("s")],
+                verify_commit_messages=[_commit("c", pre=0, tokens=[1, 2])],
+                close_messages=[self._close("x")],
+            )
+        )
+        ready = inbox.extract_ready_controls_locked(
+            lambda seg: len(seg.committed_tokens)
+        )
+        self.assertEqual([m.request_id for m in ready.sync_messages], ["s"])
+        self.assertEqual({k.request_id for k in ready.close_keys}, {"x"})
+        self.assertEqual(len(ready.ready_commit_segments), 1)
+        self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1, 2])
+        # Fully consumed -> the segment is gone; inbox drained.
+        self.assertTrue(inbox.is_empty())
+
+    def test_extract_ready_controls_zero_consumable_keeps_segment(self):
+        inbox = self._inbox()
+        inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2]))
+        ready = inbox.extract_ready_controls_locked(lambda seg: 0)
+        self.assertEqual(ready.ready_commit_segments, [])
+        # Segment is left buffered for a later step.
+        self.assertIn(DraftReqKey(0, "c"), inbox.verifier_commit_segments)
+
+    def test_extract_ready_controls_partial_consume_buffers_remainder(self):
+        inbox = self._inbox()
+        inbox.add_verify_commit_locked(_commit("c", pre=0, tokens=[1, 2, 3]))
+        ready = inbox.extract_ready_controls_locked(lambda seg: 1)
+        self.assertEqual(ready.ready_commit_segments[0].committed_tokens, [1])
+        # Remainder [2, 3] stays buffered with pre advanced to 1.
+        seg = inbox.verifier_commit_segments[DraftReqKey(0, "c")]
+        self.assertEqual(seg.committed_tokens, [2, 3])
+        self.assertEqual(seg.pre_verify_committed_len, 1)
+
+    def test_close_in_same_batch_drops_sync_and_commit(self):
+        # add_control_batch applies close first, so a same-key sync/commit in the
+        # same batch is dropped/ignored: close wins within one batch.
+        inbox = self._inbox()
+        inbox.add_control_batch_locked(
+            DraftControlBatch(
+                dst_drafter_rank=0,
+                sync_messages=[self._sync("r")],
+                verify_commit_messages=[_commit("r", pre=0, tokens=[1])],
+                close_messages=[self._close("r")],
+            )
+        )
+        self.assertEqual(inbox.sync_messages, [])
+        self.assertNotIn(DraftReqKey(0, "r"), inbox.verifier_commit_segments)
+        self.assertEqual({k.request_id for k in inbox.close_keys}, {"r"})
+
+    def test_two_commits_same_key_coalesce_in_inbox(self):
+        # A second commit for an existing key appends to the buffered segment.
+        inbox = self._inbox()
+        inbox.add_verify_commit_locked(_commit("r", pre=0, tokens=[1, 2]))
+        inbox.add_verify_commit_locked(_commit("r", pre=2, tokens=[3]))
+        seg = inbox.verifier_commit_segments[DraftReqKey(0, "r")]
+        self.assertEqual(seg.committed_tokens, [1, 2, 3])
+        self.assertEqual(seg.end_committed_len, 3)
+
+
+class TestDraftMeshMessageEnvelope(CustomTestCase):
+    def test_from_control_batch_sets_discriminant_and_slot(self):
+        batch = DraftControlBatch(dst_drafter_rank=0)
+        msg = DraftMeshMessage.from_control_batch(batch)
+        self.assertEqual(msg.message_type, DraftMeshMessageType.CONTROL_BATCH)
+        self.assertIs(msg.control_batch, batch)
+        self.assertIsNone(msg.tail_stream_output_batch)
+
+    def test_from_tail_stream_output_batch_sets_discriminant_and_slot(self):
+        from sglang.srt.speculative.decoupled_spec_io import DraftTailStreamOutputBatch
+
+        batch = DraftTailStreamOutputBatch()
+        msg = DraftMeshMessage.from_tail_stream_output_batch(batch)
+        self.assertEqual(
+            msg.message_type, DraftMeshMessageType.TAIL_STREAM_OUTPUT_BATCH
+        )
+        self.assertIs(msg.tail_stream_output_batch, batch)
+        self.assertIsNone(msg.control_batch)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)