feat(zmq): TokenSpeed DP>1 as grouped engine workers

.github/workflows/pr-test-rust.yml

@@ -583,7 +583,7 @@ jobs:
     secrets: inherit
 
   e2e-2gpu-chat-zmq-dp:
-    name: e2e-2gpu-chat-zmq-dp (vllm)
+    name: e2e-2gpu-chat-zmq-dp (${{ matrix.engine }})
     needs: [build-wheel, detect-changes]
     if: >-
       always()
@@ -593,19 +593,33 @@ jobs:
           || (needs.detect-changes.result == 'success'
               && (needs.detect-changes.outputs.common == 'true'
                   || needs.detect-changes.outputs.chat-completions == 'true')))
-    # The 1-GPU ZMQ suite with grouped workers: each vLLM ZMQ worker launches
-    # two DP engines on one socket set (tp=1 models, one GPU per engine), so
-    # the gateway's handshake, the connector's least-loaded selection, and the
-    # wave protocol all run against a real dp=2 group.
+    # The 1-GPU ZMQ suite with grouped workers: each ZMQ worker launches two
+    # DP engines on one socket set (tp=1 models, one GPU per engine), so the
+    # gateway's handshake, the connector's least-loaded selection, and
+    # per-rank output attribution all run against a real dp=2 group. The wave
+    # protocol is vLLM-only (TokenSpeed ranks run independently).
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - engine: vllm
+            timeout: 46
+          # The TokenSpeed ZMQ lane restarts the engine per model group, and
+          # DP doubles model load + warmup. The ~22-minute TokenSpeed source
+          # build eats the front of the budget, and at 50 the first dp run was
+          # killed mid-suite while progressing (28 minutes of test time); 60
+          # leaves the tests the same ~40 minutes the test_timeout allows.
+          - engine: tokenspeed
+            timeout: 60
     uses: ./.github/workflows/e2e-gpu-job.yml
     with:
-      engine: vllm
+      engine: ${{ matrix.engine }}
       # Tier selects the TEST SET (and model downloads): this lane runs the
       # tier-1 chat suite; the second GPU serves the second engine of the
       # group, not 2-GPU-marked tests.
       gpu_tier: "1"
       runner: 2-gpu-h100
-      timeout: 46
+      timeout: ${{ matrix.timeout }}
       test_timeout: 40
       test_dirs: e2e_test/chat_completions
       connection_mode: zmq

bindings/python/src/smg/serve.py

@@ -61,6 +61,13 @@ def _backend_arg_int(backend_args: list[str], flag: str, default: int) -> int:
     return default
 
 
+# The band `derive_handshake_port` folds every worker's handshake port into.
+# Other launcher-derived tcp ports must stay out of it: a port in this band can
+# collide with SOME worker's handshake listener for the right ipc path.
+_ZMQ_HANDSHAKE_PORT_BASE = 20000
+_ZMQ_HANDSHAKE_PORT_SPAN = 10000
+
+
 def _zmq_handshake_port(ipc_url: str) -> int:
     """The tcp handshake port SMG derives from an ipc:// URL.
 
@@ -73,7 +80,7 @@ def _zmq_handshake_port(ipc_url: str) -> int:
     for b in path.encode():
         h ^= b
         h = (h * 0x100000001B3) & 0xFFFFFFFFFFFFFFFF
-    return 20000 + (h % 10000)
+    return _ZMQ_HANDSHAKE_PORT_BASE + (h % _ZMQ_HANDSHAKE_PORT_SPAN)
 
 
 def _reject_handshake_port_collisions(ports: list[int]) -> None:
@@ -360,14 +367,41 @@ def build_command(
     def _build_zmq_command(
         self, args: argparse.Namespace, backend_args: list[str], port: int
     ) -> list[str]:
-        """Launch a headless TokenSpeed scheduler that dials SMG's ZMQ handshake.
+        """Launch a headless TokenSpeed scheduler group that dials SMG's ZMQ handshake.
 
         SMG (the router) binds the tcp handshake + ipc data-plane sockets it
-        derives from the ipc:// worker URL; this engine connects in. Each worker
-        is a standalone engine (`--zmq-engine-index 0`); running several is
-        dense data parallelism as N independent ZMQ workers.
+        derives from the ipc:// worker URL; the engines connect in. An
+        engine-level ``--data-parallel-size N`` (after ``--``) launches a
+        grouped worker: N ranks on one socket set, each dialing with its own
+        identity (``--zmq-engine-index`` is the group's base). The default
+        stays one standalone engine per worker.
+
+        Two ports are the launcher's to own, or co-located workers collide:
+
+        - ``--port`` seeds TokenSpeed's whole derived control-plane port
+          cluster (torch.distributed store at ``port + 233``, and neighbors).
+          Left at the engine default, every worker on the host derives the
+          same cluster, and back-to-back engine restarts race the previous
+          process's teardown (EADDRINUSE on the distributed store).
+        - ``--dist-init-addr`` pins that store explicitly. TokenSpeed derives
+          it from ``--port`` at dp==1 but refuses to guess for dp>1; passing
+          the same derivation it would use keeps one port layout for both.
         """
         rpc_port = _zmq_handshake_port(_zmq_ipc_url(port))
+        # Mirrors TokenSpeed's own dp==1 derivation (ZMQ_TCP_PORT_DELTA);
+        # reflected below the u16 ceiling instead of wrapping into low ports.
+        dist_port = port + 233 if port + 233 <= 65535 else port - 233
+        # Hop over the SMG handshake band: a dist port inside it can land on
+        # a worker's rpc listener (this worker's included — the band is a hash
+        # of the ipc path). The +233 branch enters the band only from below,
+        # so one span-wide hop exits it for good; the -233 branch starts far
+        # above the band.
+        if (
+            _ZMQ_HANDSHAKE_PORT_BASE
+            <= dist_port
+            < _ZMQ_HANDSHAKE_PORT_BASE + _ZMQ_HANDSHAKE_PORT_SPAN
+        ):
+            dist_port += _ZMQ_HANDSHAKE_PORT_SPAN
         cmd = [
             sys.executable,
             "-m",
@@ -376,6 +410,10 @@ def _build_zmq_command(
             "--headless",
             "--model",
             getattr(args, "model", ""),
+            "--port",
+            str(port),
+            "--dist-init-addr",
+            f"127.0.0.1:{dist_port}",
             "--data-parallel-address",
             "127.0.0.1",
             "--data-parallel-rpc-port",
@@ -406,6 +444,8 @@ def _build_zmq_command(
                 [
                     "--model",
                     "--headless",
+                    "--port",
+                    "--dist-init-addr",
                     "--data-parallel-address",
                     "--data-parallel-rpc-port",
                     "--zmq-engine-index",

bindings/python/tests/test_serve.py

@@ -858,9 +858,53 @@ def test_build_zmq_command(self):
         assert str(expected_port) in cmd
         assert "--zmq-engine-index" in cmd
         assert "0" in cmd
+        # The launcher owns the engine's control-plane port layout: --port
+        # seeds the derived cluster per worker, --dist-init-addr pins the
+        # torch.distributed store at TokenSpeed's own dp==1 derivation
+        # (port + 233) so dp==1 and dp>1 share one layout.
+        assert cmd[cmd.index("--port") + 1] == "31000"
+        assert cmd[cmd.index("--dist-init-addr") + 1] == "127.0.0.1:31233"
         for arg in backend_args:
             assert arg in cmd
 
+    def test_build_zmq_command_passes_dp_size_through(self):
+        # DP is the engine's flag: the launcher forwards it untouched and the
+        # ranks each dial the shared socket set with their own identity.
+        launcher = TokenspeedWorkerLauncher()
+        args = argparse.Namespace(model="/tmp/model", connection_mode="zmq")
+        cmd = launcher.build_command(args, ["--data-parallel-size", "2"], "127.0.0.1", 31000)
+
+        assert cmd[cmd.index("--data-parallel-size") + 1] == "2"
+        # dp>1 hard-requires the explicit store address the launcher always passes.
+        assert cmd[cmd.index("--dist-init-addr") + 1] == "127.0.0.1:31233"
+
+    def test_build_zmq_command_reflects_dist_port_below_u16_ceiling(self):
+        launcher = TokenspeedWorkerLauncher()
+        args = argparse.Namespace(model="/tmp/model", connection_mode="zmq")
+        cmd = launcher.build_command(args, [], "127.0.0.1", 65500)
+
+        assert cmd[cmd.index("--dist-init-addr") + 1] == f"127.0.0.1:{65500 - 233}"
+
+    def test_dist_port_never_enters_the_handshake_band(self):
+        # The handshake port is a hash of the ipc path folded into
+        # 20000..=29999, so a dist port inside that band can land on some
+        # worker's rpc listener — which port collides depends on the socket
+        # dir (uid). The launcher must keep the dist port out of the band
+        # entirely; that also implies it never equals this worker's own
+        # rpc port. Sweep every worker port whose naive +233 derivation
+        # lands in the band, plus the band edges and the u16 reflection.
+        launcher = TokenspeedWorkerLauncher()
+        args = argparse.Namespace(model="/tmp/model", connection_mode="zmq")
+        ports = [*range(19767, 29767), 19766, 29767, 31000, 65500]
+        for port in ports:
+            cmd = launcher.build_command(args, [], "127.0.0.1", port)
+            dist_port = int(cmd[cmd.index("--dist-init-addr") + 1].rsplit(":", 1)[1])
+            rpc_port = _zmq_handshake_port(_zmq_ipc_url(port))
+            assert not 20000 <= dist_port <= 29999, (port, dist_port)
+            assert dist_port != rpc_port, (port, dist_port)
+            assert dist_port != port, (port, dist_port)
+            assert 1 <= dist_port <= 65535, (port, dist_port)
+
     def test_build_zmq_command_filters_launcher_owned_flags(self):
         launcher = TokenspeedWorkerLauncher()
         args = argparse.Namespace(model="/tmp/model", connection_mode="zmq")

crates/engine_zmq_client/src/connector.rs

@@ -1208,6 +1208,7 @@ mod tests {
             cached_tokens: vec![0],
             output_token_logprobs_val: vec![vec![]],
             output_token_logprobs_idx: vec![vec![]],
+            engine_index: 0,
         };
         let done = BatchTokenIDOutSlim {
             rids: vec!["ts-1".into()],
@@ -1218,6 +1219,7 @@ mod tests {
             cached_tokens: vec![0],
             output_token_logprobs_val: vec![vec![]],
             output_token_logprobs_idx: vec![vec![]],
+            engine_index: 0,
         };
         engine
             .send_output(vec![Bytes::from(encode_msgpack(&chunk).unwrap())])

crates/engine_zmq_client/src/protocol/tokenspeed/mod.rs

@@ -99,8 +99,10 @@ impl EngineProtocol for TokenSpeedProtocol {
     }
 
     fn data_parallel_rank(_request: &Self::Request) -> Option<u32> {
-        // The TokenSpeed generate request carries no DP-rank field, so requests
-        // route to the sole engine (single-engine ZMQ). DP fan-out is future work.
+        // The TokenSpeed generate request carries no DP-rank field and needs
+        // none: rank routing is purely by ZMQ identity — the connector selects
+        // a rank and sends on that rank's socket identity. `None` means "never
+        // pinned", so every request goes through least-loaded selection.
         None
     }
 
@@ -139,18 +141,19 @@ impl EngineProtocol for TokenSpeedProtocol {
         }
         let payload = frames.first().map(AsRef::as_ref).unwrap_or_default();
         let batch: BatchTokenIDOutSlim = decode_msgpack(payload)?;
+        let engine_index = batch.engine_index;
         let outputs = batch.into_outputs()?;
         let finished_request_ids = outputs
             .iter()
             .filter(|output| output.finish_reason.is_some())
             .map(|output| output.request_id.clone())
             .collect();
         Ok(EngineBatch {
-            // Single-engine ZMQ: TokenSpeed batches carry no engine index and
-            // no piggybacked scheduler load.
-            engine_index: 0,
+            engine_index,
             outputs,
             finished_request_ids,
+            // The slim batch piggybacks no scheduler load, so DP selection
+            // scores TokenSpeed ranks on the gateway's own in-flight counts.
             load: None,
             wave: None,
         })
@@ -171,6 +174,7 @@ mod tests {
             cached_tokens: vec![0, 0],
             output_token_logprobs_val: vec![vec![], vec![]],
             output_token_logprobs_idx: vec![vec![], vec![]],
+            engine_index: 1,
         }
     }
 
@@ -181,6 +185,9 @@ mod tests {
         assert_eq!(decoded.outputs.len(), 2);
         assert_eq!(decoded.finished_request_ids, vec!["b".to_string()]);
         assert!(decoded.load.is_none());
+        // The batch names its producing DP rank; the connector routes in-flight
+        // release and scoring by it.
+        assert_eq!(decoded.engine_index, 1);
     }
 
     #[test]