[WIP] Multinode v1 by merceod · Pull Request #156 · mstar-project/mstar

merceod · 2026-07-03T11:22:40Z

What does this PR do?

This is a significant/important PR to here is a rather long description.

This PR lets one M* deployment span multiple machines. A config can now carry an optional cluster: section that maps hosts and their GPUs onto the global ranks node_groups already uses. The head host runs mstar-serve exactly as before; every other host runs a small agent (mstar-node --config cfg.yaml --node-rank k) that joins the conductor, receives a launch spec, and spawns/supervises that host's workers (building the model locally so weights load from its own cache). The ZMQ control plane switches from ipc to per-entity TCP endpoints derived from the cluster spec, tensors and KV keep moving over
the Mooncake engine (sessions now register routable addresses, with per-host device filters), and the NCCL bootstrap already was a TCP store, so TP across hosts just works. Configs without a cluster: section are untouched and single-host behavior is identical, which I verified.

I also fixed up lifecycle handling along the way: startup fails fast with a precise error instead of hanging when a host never joins or a worker dies while loading; at runtime a dead worker fails its in-flight requests with a real HTTP 500 and DP replicas route around it; replica choice now prefers keeping a request's stages on one host; and the conductor prints which graph edges will cross hosts at startup so a bad placement is visible before the first request.

A few preeexisting bugs got fixed en route to my finishing this implementation - graceful conductor shutdown called a method mp.Process doesn't have (this is where all my zombie conductors on coriander were came from), Mooncake init failures were silently swallowed, and models that specialize themselves from the config (BAGEL only registers its CFG walks inside get_worker_graphs()) now get the same setup calls on agent hosts as on the head.

How was it tested?

A lot of test have been added - around 50 or 60 new unit tests (cluster spec, endpoint resolution, launcher/agent handshake, replica assignment and cross-host cut detection, the model config-warming contract). The full existing suite has an identical pass/fail set to main throughout.

Also, verified single-host parity ie. served BAGEL with pinned request ids (pinned id --> pinned sampling seed) on this branch and on main and outputs are bit-identical.

A loopback two-"host" rig, checked in under test/multinode/, runs the real path on one machine: head + agent, TCP control plane, Mooncake TCP data plane. BAGEL prefill/decode split across the hosts is bit-identical to the same config on a single host with two GPUs, and CFG-parallel with both branches on the remote host produces bit-identical images. Orpheus (streaming split, and TP2 straddling hosts) serves correctly across the split.

Failure paths: kill -9 a remote worker → the next request gets a fast 500 ("no live replica remains…") instead of hanging; a missing agent or a worker dying during load aborts startup with a message naming the culprit.

I also verified ooverhead - basically, same pinned requests, single-host vs split, we get 2.222s vs 2.229s end-to-end (within run-to-run noise).

Checklist

ruff check . passes
Added or updated tests / docs where relevant

…er spec

stephen-dwq · 2026-07-03T16:38:02Z

In Multinode v1, SHM appears to be unusable, even for intra-host tensor transport, once multiple machines are involved. Since a substantial fraction of tensor transfers may still occur between workers on the same host, does this have a measurable impact on throughput? I'd expect SHM to be both lower-overhead and higher-throughput than RDMA or TCP for local transfers.

If the impact is non-negligible, would it make sense to register both an inter-host and an intra-host tensor manager on each worker? If GraphEdges were classified as intra-host or inter-host, the worker could dispatch each transfer through the appropriate tensor manager.

merceod · 2026-07-03T20:48:12Z

@stephen-dwq Yeah you made an excellent point. Good call! I measured it, and the impact is palpable for transfer-latency-bound walks, so I implemented the dispatch you describe (see my recently pushed commit 3c3691f). Sorry for the diatribute but I have detailed everything I noticed and what I implemented below:

I did some measurements (same single-host config served twice, --tensor-comm-protocol SHM vs forced TCP; identical seeded request sets, 16 requests at concurrency 8, image workload 6 at 3):

Orpheus PD, streaming audio: SHM rps / p50 was 3.33/s, 2.27s; forced-TCP rps / p50 was 2.69/s, 2.83s (so −19% rps and +25% p50)
BAGEL PD, text: SHM rps / p50 was 5.31/s, 0.30s; forced-TCP rps / p50 was 4.44/s, 0.45s (so −16% rps and +49% p50)
bAGEL CFG-parallel, image: SHM rps / p50 was 0.145/s, 20.3s; forced-TCP rps / p50 was 0.146/s, 20.3s (so ~0% impact)

The per-edge stats (--log-stats) was helpful in explaining the split. Orpheus pushes ~350 tiny new_token tensors per request from the decode worker to the SNAC worker: 0.12 ms per read via tmpfs vs 1.62 ms via Mooncake TCP. So basicallay, per-message latency dominates, and it sits on the audio-chunk critical path. CFG-parallel moves ~100 MiB of latents per image 6× slower over TCP loopback (150–180 MiB/s vs 800–1000 MiB/s through /dev/shm), but fifty denoise steps of compute hide it almost completely. So chatty walks lose ~15–20% throughput; compute-dominated walks don't care.

I think your proposal was a good solution and I implemented a version of that. Now the classification is done per tensor
rather than per static GraphEdge (replica assignment changes an edge's endpoints per request, and one output can fan out to several consumers). This is a summary of how it goes (you will see this in the code in more detail):

At routing time the producer classifies each output tensor by its consumer set: all consumers co-hosted, pick shared memory; any cross-host consumer, pick transfer engine. Persisted tensors always stay on the engine (a later walk can consume them from any host), and mixed local+remote fan-outs stay on the engine too. Co-hosted consumers can read the engine as well, while preparing both transports would double the producer's send work (SHM'ss "send" is a GPU to CPU copy plus a tmpfs write, so registering both would be a net loss I thik).
TensorPointerInfo.via_shm tells the consumer which transport the producer prepared HybridCommunicationManager dispatches per tensor on both ends. It's one manager with one TensorStore/ack lifecycle rather than two stacked managers. A TP-sharded edge can fan in shards over both transports at once, which two independent managers would deadlock on (each holding half the shards)
Single-host deployments are untouched: without a multi-host cluster spec the factory returns exactly the old single-transport managers. MSTAR_NO_INTRA_HOST_SHM=1 is an escape hatch that restores engine-only behavior in multi-host mode.

Some verification unit tests are added (classifier semantics, wire format, factory selection, shm round-trip/dedup/cleanup). I also did some BAGEL and Orpheus tests but please test thoroughly (whenever you have the time to) yourself. I might have missed things.

@NSagan271 ^^^^ might be interest to you as well.

NSagan271

First round of comments; I'll make another round later tonight. Overall, the abstractions look good/clean, and my comments so far are mainly nitpicks except for the one about adding a --intra-host-tensor-protocol flag in place of the MSTAR_NO_INTRA_HOST_SHM environment variable.

NSagan271 · 2026-07-04T01:48:38Z

+  tensor whose consumers are all co-hosted with its producer goes through
+  tmpfs, and only tensors with a cross-host (or not-yet-known, e.g.
+  persisted) consumer travel through the transfer engine. Set
+  ``MSTAR_NO_INTRA_HOST_SHM=1`` to force everything through the engine.


I think it would be more generalizable to specify a --intra-host-tensor-protocol (default SHM) instead of having a binary "env var not set => SHM, env var set => RDMA" flow. Especially if we end up having more options for tensor transport in the future, it would be nice to be able to just specify two different protocols.

Replaced the env var with--intra-host-tensor-protocol (default SHM) in commit 7c83902. Any non-SHM value routes co-hosted transfers through the transfer engine alongside cross-host traffic, single-host deployments ignore it, and it gives CUDA-IPC a natural slot later.

Do we have a direct D2D tensor communication? Everything seems at least D2H2D. I think there is a NCCL python wrapper that might be useful.

@merceod I was envisioning the HybridCommunicationManager being able to take in two arbitrary types of communication, instead of being hardcoded to the local communication being SHM. I guess right now the only option is TCP/RDMA for inter-node and SHM for local, but the abstraction/generalizability would be useful if, e.g., we add the potential NCCL tensor communication Steven is talking about.

NSagan271 · 2026-07-04T01:55:01Z

-            enable_prof=self.log_stats
+            enable_prof=self.log_stats,
+            endpoints=endpoints,
+            # The preprocess worker never picks shm for its own sends (a


Aside/note: if data worker -> initial worker or cross-graph-walk transfers end up being a bottleneck, we could try out double-registering initial or persist tensors (assuming it is properly configurable and that the overhead of doing so is not too much), and the receiver can choose which way to read the tensor.

Agreed it's a clean shape if those edges ever profile hot. One caveat from the measurements: the shm "send" is a D2H copy + tmpfs write, so double-registering pays that even when the reader chooses the engine which isworth keeping opt-in and profiling-driven. We should note this as a follow-up (or maybe add it to this pr).

NSagan271

More comments, also take a look at #156 (comment).

NSagan271 · 2026-07-05T20:48:48Z

        model=model,
        model_name=model_name,
-        tcp_transfer_device=args.tcp_transfer_device,
+        tcp_transfer_device=cluster_spec.head.rdma_device or args.tcp_transfer_device,


I would unify the config for the RDMA device and TCP transfer device to both come from the transfer spec, instead of one from the transfer spec and one from the CLI args.

NSagan271 · 2026-07-05T22:03:46Z

+                kwargs["worker_id"], kwargs.get("device"), p.pid,
+            )
+
+    def _check_children(self) -> None:


Is the node agent itself dying of concern / should the conductor be periodically polling the health of the node agents themselves (via periodic ZMQ messages) to see if they're still alive? I could see this being useful if, e.g., someone does CTRL+C on one of the node agent processes.

NSagan271 · 2026-07-05T22:05:32Z

+        is only valid when every worker shares one machine.
+        """
+        name = getattr(protocol, "value", protocol)
+        if self.is_multi_host() and str(name).upper() == "SHM":


For future generalizability, I would give the protocol enum a function def multinode_support(self) -> bool instead of checking if the protocol is SHM here.

NSagan271 · 2026-07-05T22:15:39Z

+        if self._resolver is not None:
+            self.protocol = CommProtocol.TCP
+        else:
+            transport = os.getenv("MSTAR_ZMQ_TRANSPORT", protocol.value).upper()


Nit: I know this isn't your change, but I'd prefer to wire the ZMQ transport type through the CLI instead of as an environment variable. Or to have clear documentation of which config options are via environment variable and which are through the CLI.

NSagan271 · 2026-07-05T22:17:00Z

            raise NotImplementedError(f"Unknown protocol {protocol} for mooncake")
+        # Device filter for the engine: an HCA name (or comma list) under RDMA,
+        # a segment/interface string under TCP. Empty means auto-discovery.
+        transfer_device = tcp_transfer_device


Nit: if the transfer device argument is used for both TCP and RDMA, rename it transfer_device instead of tcp_transfer_device

NSagan271 · 2026-07-05T22:18:43Z

    def register_for_send(
        self, request_id: str, uuids: list[str],
        skip_cuda_sync: bool = False,
+        shm_uuids: set[str] | None = None,


Nit: I'd name this local_uuids instead of tying it to a specific communication protocol, e.g., if we end up having SHM and NCCL as options for local tensor communication.

NSagan271 · 2026-07-05T22:21:22Z

            shm_dir=shm_dir,
            enable_prof=enable_prof
        )
+    if same_host_entities is not None:


I'd add the intra-node communication protocol as an argument to the factory, and then also gate the hybrid communication manager on protocol != local_protocol, instead of leaving that gating to the caller (that seems to currently be happening by passing in None for same_host_entities, which is a bit un-intuitive).

NSagan271 · 2026-07-05T22:25:11Z

+        tcp_transfer_device="",
+        intra_host_tensor_protocol=CommProtocol.SHM,
+        model_cache_dir: str | None = None,
+        startup_timeout_s: float = 600.0,


I'd make this timeout configurable by the user as a CLI argument or environment variable, just like vllm-omni and sglang-omni do. I remember we had issues with vllm-omni timing out on the very first startup and having to set the timeout to be higher.

merceod added 6 commits July 3, 2026 07:26

Map global worker ranks onto hosts and local CUDA devices via a clust…

0fd1780

…er spec

Resolve control-plane endpoints per entity for multi-host TCP

ad5eda7

Launch multi-host workers through per-host node agents

e1cdab0

Add per-host transfer devices, multi-host docs, and a loopback test rig

3512b1a

Warm agent-built models with the deployment config

e632156

Prefer co-located replicas and report cross-host edges at startup

3254f90

merceod requested review from NSagan271, kamahori and stephen-dwq July 3, 2026 11:22

Use shared memory for co-hosted tensor transfers in multi-host mode

3c3691f

NSagan271 requested changes Jul 4, 2026

View reviewed changes

merceod added 2 commits July 4, 2026 12:05

Add --intra-host-tensor-protocol and drop the env-var hatch

7c83902

Break replica-assignment ties toward workers already in use

cc90d0d

NSagan271 requested changes Jul 5, 2026

View reviewed changes

This was referenced Jul 13, 2026

Scope a migration toward a Rust-native runtime (Python stays the model layer) #130

Open

SHM tensor transport over a Rust shared-memory arena (RFC #130, Step 2) #168

Merged

Uh oh!

Conversation

merceod commented Jul 3, 2026

What does this PR do?

How was it tested?

Checklist

Uh oh!

stephen-dwq commented Jul 3, 2026

Uh oh!

merceod commented Jul 3, 2026

Uh oh!

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

NSagan271 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants