[gbyu-amd]

Problem

ATOM's tensor parallelism runs on one node: -tp 8 needs 8 GPUs on the host. On nodes with 4 GPUs (MI455X), Kimi-K3 with TP8 attention needs one TP group over two nodes.

On gfx1250, DCP decode also could not run persistent. The non-persistent ASM decode it fell back to takes only 64-token pages, so -tp 8 -dcp 8 for Kimi-K3 failed its first decode graph capture with mla_decode_gfx1250_dispatch: only support page_size == 64.

Changes

[TP] Run one tensor-parallel group across nodes. New flags --tensor-parallel-size-local (ranks on this node) and --tensor-parallel-rank (first rank on this node).

• The node that runs TP rank 0 runs the engine, the API server and its own ranks. Every other node runs only its ranks' workers (atom/model_engine/tp_multinode.py).
• The engine reaches the remote ranks through AsyncIOProcManager's existing channels: the broadcast queue gets remote readers (aiter's MessageQueue already serves them over ZMQ), and their reply channels listen on TCP.
• The manager publishes those endpoints and the runner's constructor arguments in the model-runner TCPStore that CoreManager serves, waits for every remote reader to subscribe (a PUB socket drops what it sends before that), then releases them. The remote node spawns its workers from that, and stops them if one fails or the engine's store goes away.
• ModelRunner places each worker on its node's GPU (rank - tp_rank_start) and takes the minimum KV block count over the TP group.
• Barrier RPCs are refused, because a multiprocessing barrier does not reach another node. DP, DP attention, PP, PCP, --fake-eplb and RapidServe are rejected with a cross-node TP group.
• New env var ATOM_MULTINODE_TP_TIMEOUT_S (default 1800 s): how long the nodes wait for each other at startup.

[MLA][DCP] Run DCP decode persistent on gfx1250 with the PS1 kernel. dcp_persistent_supported() is now true on gfx1250 when AITER_MLA_DECODE_PS1_FLYDSL=1 selects the PS1 FP8 decode. That kernel is persistent and returns the LSE that DCP's merge needs (ROCm/aiter#5844 added DCP to it).

Usage

# node0: TP ranks 0-3, the engine and the API server
python -m atom.entrypoints.openai_server --model <model> -tp 8 --tensor-parallel-size-local 4 \
  --data-parallel-master-ip <node0 IP> --data-parallel-base-port <port> ...
# node1: TP ranks 4-7
python -m atom.entrypoints.openai_server --model <model> -tp 8 --tensor-parallel-size-local 4 \
  --tensor-parallel-rank 4 --data-parallel-master-ip <node0 IP> --data-parallel-base-port <port> ...

Validation

Kimi-K3 on two MI455X nodes (4 GPUs each): TP8 + DCP8 attention, MoE TP8 with ROCm/aiter#6706, RCCL between the nodes. GSM8K 5-shot:

RCCL between the nodes	Mode	strict-match
scale-up fabric (NCCL_MNNVL_ENABLE=1)	CUDA graphs	0.9621 ± 0.0053
TCP sockets (NCCL_MNNVL_ENABLE=0)	--enforce-eager	0.9583 ± 0.0055

• With MoE EP8 instead of TP8: 0.9629 ± 0.0052.
• Over TCP sockets, capturing the cross-node all-reduce in a CUDA graph fails with hipErrorInvalidValue, so that row runs eager.
• These runs used the two commits on a 2026-10-16 base (d82f906). They cherry-pick onto main (53e4fb5) without conflicts, but have not been rerun there. Of what landed in between, feat(mono): GLM-5.2 and Kimi-K3 mono decode, with the kernels vendore… #2535 turns on Kimi-K3 mono decode by default (ATOM_MONO_ENABLE=1); it has not been tried with a cross-node TP group.
• black is clean. ruff reports nothing new: its three findings in engine_core_mgr.py are on main too. The non-GPU unit tests have not been run locally.

Notes

• Draft.
• MoE TP8 needs [FlyDSL] [MoE] Fit the gfx1250 grouped GEMM k-tile to K, and refuse one that does not divide it aiter#6706: without it the gfx1250 grouped MoE drops a third of the expert slice's K at TP8, and GSM8K falls to 0.8923.
• Not included: speculative decode (DSpark) with DCP on gfx1250, and a launch script for the two-node setup.

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2647 --add-label <label>

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[zejunchen-zejun]

Review: #2647 — "feat(tp): run one tensor-parallel group across nodes"

• Head reviewed: c59ffd5d (2026-10-21)
• Scope: the six issues below are the most severe ones: hangs, GPUs left held, crashes, and one network-security exposure.
• How they were found: every issue was checked against the code at that commit. The failure scenarios follow from reading the code; none was reproduced on GPU.
• What still works: the configuration the PR validated is not affected: 2 nodes x 4 GPUs, TP8 + DCP8, openai_server, master IP and port given as CLI flags.

Line numbers refer to c59ffd5d.

---

1. A dead TP rank on another node is invisible to the engine, which then hangs

Where:

• atom/model_engine/async_proc.py:1038-1040: monitor_procs watches self.procs only.
• async_proc.py:957: call_func(..., wait_out=True) blocks on outputs_queue.get() with no timeout.

What happens: rank 5 on node1 dies, for example from an OOM.

• node1: _supervise stops ranks 4-7 and exits, but nothing reaches the engine.
• node0: ranks 0-3 block in the next cross-node collective, and EngineCore blocks forever waiting for the forward output.
• Clients: the API server stays up and every request hangs.
• collective_rpc: it only reports a generic timeout for ranks 4-7. Its liveness check (rank < len(procs)) never covers remote ranks.

Why: self.procs holds only the local_proc_num local workers. Remote workers are spawned by serve_remote_tp_ranks on the other node, and the engine has no liveness signal for them.

Suggested fix:

• Report failures: when a worker node stops, have it report the failure in the shared TCPStore, e.g. atom/tp_multinode/failed/<start_rank>.
• Watch for them: give the engine a monitor thread that polls those keys, or a heartbeat per remote node, and calls AsyncIOProcManager.exit() when one appears. This mirrors what monitor_procs already does for local ranks.

---

2. Engine shutdown never reaches the other node; it keeps its GPUs, and a restarted engine is mistaken for the old one

Where:

• async_proc.py:618-628: exit() stops self.procs only.
• atom/model_engine/tp_multinode.py:190: _supervise calls watchdog.check([BOOTSTRAP_KEY]) and discards the result.

What happens:

• Engine stops, server stays up: a local rank dies, or the engine shuts down for any reason. exit() sends nothing to ranks 4-7. node1 only notices when the TCPStore becomes unreachable, and that store lives in the API-server process, not in the engine. As long as the server process stays up, node1's workers stay blocked in a collective or a dequeue and hold all four GPUs indefinitely.
• Engine restarts on the same port: the watchdog reconnects to the new TCPStore. check() returns False because the new store has no bootstrap key yet, but the miss counter resets anyway. node1's workers stay attached to the dead engine's broadcast socket, still holding the GPUs. The new engine then waits ATOM_MULTINODE_TP_TIMEOUT_S (30 min) for subscribers that never arrive.

Suggested fix:

• Publish shutdown: on exit(), write a shutdown key to the store, e.g. atom/tp_multinode/shutdown, and have _supervise stop its workers when it sees that key.
• Count a failed check: treat check() returning False as a miss, the same as a connection error.
• Tie the store to the engine: consider keying the bootstrap by an engine instance id, so a restarted engine is never mistaken for the old one.

---

3. Unauthenticated pickle deserialization over the network (security)

Where:

• tp_multinode.py:50, 121, 128: the engine publishes the bootstrap blob with pickle.dumps into a TCPStore that has no authentication. The worker node reads it back and pickle.loads it.
• async_proc.py:612-616: for remote ranks, the per-rank reply channels are now tcp://<data_parallel_master_ip>:<random port>.
• async_proc.py:685, 704: those channels are served by process_kv_output_sockets and process_rpc_output_sockets, which pickle.loads every frame they receive.

What happens: any host that can reach node0 on the data-plane network can exploit this. Before this PR the reply channels were node-local IPC paths. Now they are TCP sockets bound on a routable address.

• Reply ports: a host can connect a PUSH socket to one of the reply ports and send a crafted pickle. That runs arbitrary code in the EngineCore process.
• Bootstrap key: a host can write atom/tp_multinode/bootstrap before the engine does. That runs arbitrary code in the worker node's server process when it loads the blob.

Suggested fix:

• Minimum:
  • document that the data-plane network must be trusted;
  • bind the reply sockets only on the interface the peer node uses.
• Better:
  • authenticate the channel, e.g. ZMQ CURVE or a shared secret/HMAC on each frame;
  • stop unpickling data that arrives from the network: send the bootstrap as plain data such as JSON, and rebuild Config on the worker node from its own CLI.

---

4. Rank slices are not cross-checked, and subscribers are counted without knowing which rank they are

Where:

• tp_multinode.py:130: the worker node only checks ranks[-1] < proc_num.
• tp_multinode.py:78: release_remote_readers counts XPUB subscription messages with no identity attached.

What happens:

• Slices overlap: say the engine runs ranks 0-3, and node1 is launched with --tensor-parallel-rank 2 --tensor-parallel-size-local 2. Ranks 2-3 are local readers in the engine's queue handle. create_from_handle then tries to open node0's /dev/shm segment on node1, and their reply addresses are node0's IPC paths. The result is a crash with a confusing error.
• Duplicate slices: say two worker nodes are both launched with --tensor-parallel-rank 4. Eight subscribers arrive, the engine releases after any four, and duplicate global ranks 4-7 join init_process_group, which hangs or fails.

Suggested fix:

• Validate the slice: put the engine's local_proc_num in the bootstrap, and have the worker node reject start < local_proc_num.
• Claim ranks explicitly: each remote rank registers itself, e.g. with a store add/set on atom/tp_multinode/rank/<r>, before it subscribes. The engine then waits for exactly the expected rank set, and rejects a duplicate claim.

---

5. Kimi-K3 mono decode is not refused when the TP group spans nodes

Where:

• atom/config.py:2177: the unsupported list in _check_multinode_tp.
• atom/mono/runtime/deployment.py:60-85: common_refusals.

What happens:

• Mono is on by default: ATOM_MONO_ENABLE defaults to 1 (feat(mono): GLM-5.2 and Kimi-K3 mono decode, with the kernels vendore… #2535).
• It needs GPU peer memory: mono builds peer buffers across the TP group with HIP IPC (atom/mono/runtime/hip_ipc.py, peer_memory.py), and HIP IPC cannot reach a GPU on another node.
• Nothing turns it off across nodes: common_refusals checks TP size, EP, DP, PP, decode/prefill CP and others, but not whether the TP group is on one node. _check_multinode_tp does not refuse mono either.
• Where it breaks: the validated runs all used DCP8, which mono already refuses ("decode CP"). A cross-node TP8 run without DCP enables mono, and IPC handle exchange or peer access fails across nodes.

Suggested fix: add (not atom_config.is_multinode_tp, "TP group spans nodes (HIP IPC peer memory)") to common_refusals. Mono then turns itself off with a logged reason, instead of cross-node TP refusing to start.

What happens:

• (a) The nodes can disagree. The env var is not part of the pickled Config, so each node decides on its own. If node0 sets it and node1 does not, ranks 0-3 size heads and work metadata for persistent mode and ranks 4-7 for non-persistent mode. The DCP Q all-gather and LSE merge then disagree across the group: a hang, garbage output, or the original only support page_size == 64 abort on half the ranks.
• (b) Some widths have no kernel. aiter's FlyDSL PS1 gate takes only nhead in (16, 32, 64, 96, 128), and before [ASM] [HIP] [FlyDSL] Optimize gfx1250 persistent MLA decode, add BF16 and TP4/TP2 MTP aiter#6563 only an FP8 cache. A gathered width of 48, 80 or 112 (e.g. TP8 with DCP4 gives 12 x 4 = 48) still selects persistent mode, but PS1 rejects it. The decode then falls through to aiter's non-gfx1250 branches and aborts. Kimi-K3 TP8 + DCP8 gathers 96 heads, so the validated run does not hit this.

Suggested fix:

• Agree across nodes: decide once on the engine and carry the decision in Config/bootstrap. At minimum, have every rank compare its value through the store at startup.
• Pad to a supported width: on gfx1250 in persistent mode, round the gathered width up to the next value in {16, 32, 64, 96, 128}.

[zejunchen-zejun]

LGTM

[zejunchen-zejun]

@gbyu-amd please help resolve the conflict

[yitingw1]

Re-checked at head 24049fed (the two 2026-10-22 commits). Both new commits look right to me; most of this comment is about what they do and do not cover.

The two new commits

24049fed (mono refusal) resolves finding 5. Putting it in common_refusals rather than _check_multinode_tp is the right layer: first_refusal returns a reason string and the caller falls back off mono, whereas _check_multinode_tp would hard-error — and mono is an optimisation, not a feature, so a silent fallback is what you want across nodes. The getattr(..., False) default is safe too, since plugin configs are refused by the "plugin mode" check two lines down regardless.

c3cc2367 (store timeout) is a defect you found yourselves, not one of the six. Sibling check is clean: of the three TCPStore sites on this branch, engine_core_mgr.py:330 is the is_master=True server, tp_multinode.py:39 already used ATOM_MULTINODE_TP_TIMEOUT_S, and model_runner.py:1054 is the one this commit fixes. Nothing of that shape is left.

Status of the six findings

I re-read each one against 24049fed rather than assuming:

#	finding	status at head
1	dead remote rank hangs the engine	open — async_proc.py:1045 is still connection.wait(sentinels)
2	shutdown does not reach the other node; misses always resets	open — tp_multinode.py:190-191 still sets misses = 0 unconditionally after check()
3	unauthenticated pickle over the network	open — tp_multinode.py:50 / :128 unchanged
4	no cross-validation of rank slices; subscriptions counted, not identified	open — tp_multinode.py:78 is still subscribed += 1, :130 still only checks ranks[-1] >= proc_num
5	mono not refused across nodes	fixed by 24049fed
6	AITER_MLA_DECODE_PS1_FLYDSL read per process	open — dcp_utils.py:62 unchanged

@zejunchen-zejun — does the LGTM cover the PR as a whole, or just the two new commits? Asking because five of the six are still live, including 1-3, which your original review called out as hangs / GPUs left held / a network RCE surface. Happy to be told they are deliberate follow-ups; I just do not want to read the approval as "those are resolved" if they are not.

One new issue, from c3cc2367

ATOM_MULTINODE_TP_TIMEOUT_S has a 30-minute ceiling it does not control. The commit routes it into the store client, but the next statement passes no timeout of its own:

store = torch.distributed.TCPStore(..., **managed_store_client_kwargs(config))
torch.distributed.init_process_group(
    backend="nccl", store=store, rank=global_rank, world_size=world_size
)

Measured on this image (torch 2.10.0+rocm7.2.4): default_pg_timeout is 0:30:00 and init_process_group's own default is None, so it falls back to that.

The 13-minute skew the commit is about is the early rank sitting inside the first collective waiting for the late one, and that wait is timed by the PG/NCCL watchdog, not by the store. It works today only because 1800 s and 30 min happen to be the same number. Raise ATOM_MULTINODE_TP_TIMEOUT_S to, say, 2400 for a tray with a slower disk and the store will wait longer while the collective still aborts at 30 minutes — the knob looks like it took effect right up until someone needs it to.

Suggest passing it through when the group spans nodes:

pg_kwargs = {}
if config.is_multinode_tp:
    pg_kwargs["timeout"] = timedelta(seconds=envs.ATOM_MULTINODE_TP_TIMEOUT_S)
torch.distributed.init_process_group(..., **pg_kwargs)

Minor, same commit: the two new tests assert the mapping (managed_store_client_kwargs returns the right dict) but not that TCPStore receives it, so deleting the **managed_store_client_kwargs(config) at the call site leaves both green. The whole value of the commit is that one call site; a test that monkeypatches TCPStore and asserts on the kwargs would hold it.

Two notes

On finding 6 — if it is not going to be fixed soon, it is at least a documentation item. Two nodes disagreeing about AITER_MLA_DECODE_PS1_FLYDSL is not a code defect so much as an undocumented deployment constraint, and right now nothing tells the operator that the variable has to match on both nodes or that the failure mode is a hang / wrong output rather than an error.

Finally, an interaction worth noting rather than a defect: c3cc2367 raises the startup wait from 300 s to 1800 s, and finding 1 means the failure path has no timeout at all. Together they make a failed cross-node start take considerably longer to show itself.

[gbyu-amd]

@valar-cop review this pr

[valar-cop]

some changes needed..

gsm8k for cross-node tp8 without dcp?

ci green.

[valar-cop]

env var works on node0 but fails on node1? resolve through ParallelConfig so both nodes agree..

[valar-cop]

local 16 with tp 8 is silently ignored here, gate on flags set and reject local > tp?

[valar-cop]

ipv6 breaks this, use get_tcp_uri instead of string format..

[valar-cop]

why is this helper here? move it to tp_multinode so tests run without aiter..

[valar-cop]

only in api_server? other entrypoints accept the flag and build a partial engine, handle it once in a shared place..

[valar-cop]

a few small things..

gsm8k for cross-node tp8 without dcp? and a run on current head?

ci green.

[valar-cop]

why importorskip aiter? profile.py only needs dataclasses, so this never runs in ci

[valar-cop]

node1 never builds ParallelConfig, so ATOM_DP_BASE_PORT/MASTER_IP are ignored here? resolve the endpoint the same way both nodes do

[valar-cop]

local > tp is accepted and silently ignored here? validate whenever either flag was set

[valar-cop]

no brackets for ipv6, use get_tcp_uri instead of formatting by hand

[valar-cop]

model_runner imports aiter so the tests skip in ci, move this to tp_multinode.py?

[valar-cop]

atomesh server builds a full local engine with this flag and no error, handle it once where EngineArgs becomes an engine

[valar-cop]

some changes needed..

gsm8k for cross-node tp8 without dcp? and a run on current head?

ci green.

[valar-cop]

test bypasses post_init, so reverting this line still passes. pin the call site?