feat(tp): run one tensor-parallel group across nodes

atom/config.py

@@ -1069,6 +1069,14 @@ class ParallelConfig:
 
     data_parallel_master_ip: str = "127.0.0.1"
 
+    tensor_parallel_size_local: int | None = None
+    """TP ranks (model-runner workers) this node runs. None means all of them,
+    the single-node case. Set it below tensor_parallel_size to spread one TP
+    group over several nodes; see atom/model_engine/tp_multinode.py."""
+    tensor_parallel_rank: int = 0
+    """First TP rank this node runs. The node that runs rank 0 also runs the
+    engine and the API server; the others run only their workers."""
+
     @property
     def is_multinode_dp(self) -> bool:
         """Whether this node owns only part of the global DP group.
@@ -2207,6 +2215,61 @@ def tp_world_size(self) -> int:
         visible = torch.cuda.device_count()
         return visible if 0 < visible < tp else tp
 
+    @property
+    def tp_size_local(self) -> int:
+        """TP worker processes on this node: fewer than `tp_world_size` when
+        one TP group spans several nodes."""
+        local = self.parallel_config.tensor_parallel_size_local
+        return self.tp_world_size if local is None else local
+
+    @property
+    def tp_rank_start(self) -> int:
+        """First TP rank on this node; it runs `tp_size_local` ranks from here."""
+        return self.parallel_config.tensor_parallel_rank
+
+    @property
+    def is_multinode_tp(self) -> bool:
+        return self.tp_size_local < self.tp_world_size or self.tp_rank_start > 0
+
+    def _check_multinode_tp(self) -> None:
+        pc = self.parallel_config
+        if pc.tensor_parallel_size_local is None and pc.tensor_parallel_rank == 0:
+            return
+        local, start, tp = self.tp_size_local, self.tp_rank_start, self.tp_world_size
+        if local < 1 or tp % local or start % local or not 0 <= start < tp:
+            raise ValueError(
+                f"--tensor-parallel-size-local {local} and --tensor-parallel-rank "
+                f"{start} must cut -tp {tp} into whole, equal per-node slices"
+            )
+        if not self.is_multinode_tp:
+            return
+        unsupported = [
+            name
+            for name, used in (
+                ("data parallelism", pc.data_parallel_size > 1),
+                ("--enable-dp-attention", self.enable_dp_attention),
+                ("pipeline parallelism", self.pipeline_parallel_size > 1),
+                ("prefill context parallelism", self.prefill_context_parallel_size > 1),
+                ("--fake-eplb", self.fake_eplb),
+                ("RapidServe", self.enable_rapidserve),
+            )
+            if used
+        ]
+        if unsupported:
+            raise ValueError(
+                "tensor parallelism across nodes does not support "
+                + ", ".join(unsupported)
+            )
+        if pc.data_parallel_base_port == 0 or pc.data_parallel_master_ip in (
+            "127.0.0.1",
+            "localhost",
+        ):
+            raise ValueError(
+                "tensor parallelism across nodes needs --data-parallel-master-ip "
+                "set to the address of the node that runs TP rank 0, and the "
+                "same nonzero --data-parallel-base-port on every node"
+            )
+
     @property
     def capture_sizes(self) -> list[int]:
         """The declared CUDAGraph capture ladder, in batch sizes.
@@ -2380,6 +2443,7 @@ def __post_init__(self):
             )
 
         assert 1 <= self.tensor_parallel_size <= 8
+        self._check_multinode_tp()
         if self.decode_context_parallel_size > 1:
             assert self.tensor_parallel_size % self.decode_context_parallel_size == 0, (
                 f"tp_size ({self.tensor_parallel_size}) must be divisible by "

atom/distributed/dcp_utils.py

@@ -57,14 +57,21 @@ def get_dcp_rank() -> int:
 def dcp_persistent_supported() -> bool:
     """Whether DCP decode can run in *persistent* mode on this GPU.
 
-    Needs an lse-emitting ASM decode kernel; only gfx950 ships one (gfx942
-    falls back to the triton stage2 reduce instead). Cache the result once
-    per ``__init__`` — calling this per-forward would graph-break on
-    ``get_gfx()``.
+    Needs an lse-emitting persistent decode kernel: gfx950's ASM one, and on
+    gfx1250 the PS1 FP8 one that AITER_MLA_DECODE_PS1_FLYDSL=1 selects (aiter
+    #5844 added DCP to it); without that, gfx1250's non-persistent decode takes
+    only 64-token pages. gfx942 falls back to the triton stage2 reduce instead.
+    Cache the result once per ``__init__`` — calling this per-forward would
+    graph-break on ``get_gfx()``.
     """
+    import os
+
     from aiter.jit.utils.chip_info import get_gfx
 
-    return get_gfx() == "gfx950"
+    gfx = get_gfx()
+    if gfx == "gfx1250":
+        return os.environ.get("AITER_MLA_DECODE_PS1_FLYDSL", "0") == "1"
+    return gfx == "gfx950"
 
 
 def mla_dcp_decode_is_persistent(

atom/model_engine/arg_utils.py

@@ -119,6 +119,8 @@ class EngineArgs:
     model: str = "Qwen/Qwen3-0.6B"
     trust_remote_code: bool = False
     tensor_parallel_size: int = 1
+    tensor_parallel_size_local: int | None = None
+    tensor_parallel_rank: int = 0
     decode_context_parallel_size: int = 1
     pipeline_parallel_size: int = 1
     prefill_context_parallel_size: int = 1
@@ -212,6 +214,28 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             default=1,
             help="Tensor parallel size.",
         )
+        parser.add_argument(
+            "--tensor-parallel-size-local",
+            type=int,
+            default=None,
+            help=(
+                "Number of tensor-parallel ranks to run on THIS node. Defaults "
+                "to --tensor-parallel-size (single-node). Set it lower to "
+                "spread one TP group over several nodes; every node then needs "
+                "--data-parallel-master-ip of the node that runs TP rank 0 and "
+                "the same --data-parallel-base-port."
+            ),
+        )
+        parser.add_argument(
+            "--tensor-parallel-rank",
+            type=int,
+            default=0,
+            help=(
+                "First tensor-parallel rank owned by this node. Node 0 uses 0 "
+                "and serves the API; the second node of a 2x4 run uses 4 and "
+                "runs only its model-runner workers."
+            ),
+        )
         parser.add_argument(
             "--pipeline-parallel-size",
             "-pp",
@@ -929,22 +953,46 @@ def _get_engine_kwargs(self) -> dict:
 
         # DP topology -> ParallelConfig. `data_parallel_size` stays in kwargs
         # too: LLMEngine still reads the loose kwarg on the legacy path.
-        parallel_config_kwargs = {
-            "data_parallel_size": kwargs["data_parallel_size"],
-            "data_parallel_size_local": kwargs.pop("data_parallel_size_local"),
-            "data_parallel_rank": kwargs.pop("data_parallel_rank"),
-            "data_parallel_master_ip": kwargs.pop("data_parallel_master_ip"),
-            "data_parallel_master_port": kwargs.pop("data_parallel_master_port"),
-        }
-        base_port = kwargs.pop("data_parallel_base_port")
-        if base_port is not None:
-            parallel_config_kwargs["data_parallel_base_port"] = base_port
-        kwargs["parallel_config"] = ParallelConfig(**parallel_config_kwargs)
+        for name in self._PARALLEL_CONFIG_FIELDS:
+            if name != "data_parallel_size":
+                kwargs.pop(name)
+        kwargs["parallel_config"] = self.parallel_config()
 
         logger.info(f"Engine kwargs: {kwargs}")
 
         return kwargs
 
+    _PARALLEL_CONFIG_FIELDS = (
+        "data_parallel_size",
+        "data_parallel_size_local",
+        "data_parallel_rank",
+        "data_parallel_master_ip",
+        "data_parallel_master_port",
+        "data_parallel_base_port",
+        "tensor_parallel_size_local",
+        "tensor_parallel_rank",
+    )
+
+    def parallel_config(self) -> ParallelConfig:
+        """This node's ParallelConfig, with the ATOM_DP_* environment applied.
+
+        A node that runs only TP workers reads the engine's address from it as
+        well, so both nodes resolve that address the same way.
+        """
+        kwargs = {name: getattr(self, name) for name in self._PARALLEL_CONFIG_FIELDS}
+        if kwargs["data_parallel_base_port"] is None:
+            del kwargs["data_parallel_base_port"]
+        return ParallelConfig(**kwargs)
+
     def create_engine(self, tokenizer=None) -> LLMEngine:
-        """Create and return an LLMEngine instance with the configured parameters."""
+        """Create and return an LLMEngine instance with the configured parameters.
+
+        A node whose TP slice starts past rank 0 (--tensor-parallel-rank) has no
+        engine to create: it runs its TP workers for the engine on another node
+        and exits the process with their status once they stop.
+        """
+        if self.tensor_parallel_rank > 0:
+            from atom.model_engine.tp_multinode import serve_remote_tp_ranks
+
+            raise SystemExit(serve_remote_tp_ranks(self))
         return LLMEngine(self.model, tokenizer=tokenizer, **self._get_engine_kwargs())

atom/model_engine/async_proc.py

@@ -46,7 +46,9 @@
 from atom.utils import (
     envs,
     get_mp_context,
+    get_open_port,
     get_open_zmq_ipc_path,
+    get_tcp_uri,
     init_exit_handler,
     make_zmq_socket,
     resolve_obj_by_qualname,
@@ -155,7 +157,7 @@ def __init__(
             dp_local_rank = cfg.parallel_config.data_parallel_rank_local
             if dp_local_rank is None:
                 dp_local_rank = cfg.parallel_config.data_parallel_rank
-            gpu = dp_local_rank * cfg.tp_world_size + rank
+            gpu = dp_local_rank * cfg.tp_world_size + rank - cfg.tp_rank_start
             numa_bind_to_node(gpu, label)
         except Exception as e:  # noqa: BLE001 - binding is an optimization
             # NUMA binding only affects locality, never correctness, so any
@@ -178,6 +180,10 @@ def __init__(
         self.rpc_queue: queue.Queue | None = None
 
         self.rpc_broadcast_mq = MessageQueue.create_from_handle(input_shm_handle, rank)
+        if self.rpc_broadcast_mq._is_remote_reader:
+            from atom.model_engine.tp_multinode import await_release
+
+            await_release(self.rpc_broadcast_mq)
         import atexit
 
         atexit.register(self._cleanup_shared_memory)
@@ -464,11 +470,22 @@ class AsyncIOProcManager:
     )
 
     def __init__(
-        self, finalizer, proc_num: int, runner: str, *args, on_worker_failure=None
+        self,
+        finalizer,
+        proc_num: int,
+        runner: str,
+        *args,
+        local_proc_num: int | None = None,
+        on_worker_failure=None,
     ):
         self.parent_finalizer = finalizer
         self.on_worker_failure = on_worker_failure
         self.proc_num = proc_num
+        # Ranks from local_proc_num on run on other nodes (tp_multinode.py) and
+        # reach this node over TCP.
+        self.local_proc_num = proc_num if local_proc_num is None else local_proc_num
+        n_remote = proc_num - self.local_proc_num
+        host = args[0].parallel_config.data_parallel_master_ip if n_remote else None
 
         io_addrs = [get_open_zmq_ipc_path(), get_open_zmq_ipc_path()]
         self.procs: list[multiprocessing.Process] = []
@@ -477,7 +494,10 @@ def __init__(
         self.label = f"AsyncIOProcManager({self.runner_label})"
 
         self.rpc_broadcast_mq = MessageQueue(
-            proc_num, proc_num, max_chunk_bytes=16 * 1024 * 1024
+            proc_num,
+            self.local_proc_num,
+            max_chunk_bytes=16 * 1024 * 1024,
+            connect_ip=host,
         )
         self._block_table_encoder = BlockTableDeltaEncoder()
         scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
@@ -486,20 +506,22 @@ def __init__(
         import atexit
 
         atexit.register(self._cleanup_shared_memory)
-        self.all_ranks_barrier = ctx.Barrier(proc_num)
+        # A multiprocessing barrier cannot reach another node; barrier RPCs are
+        # refused instead (see _collective_rpc).
+        self.all_ranks_barrier = None if n_remote else ctx.Barrier(proc_num)
         init_exit_handler(self)
 
         # KV output aggregation infrastructure
         self.kv_output_aggregator: KVOutputAggregator | None = None
-        self.kv_output_addrs = [get_open_zmq_ipc_path() for _ in range(proc_num)]
+        self.kv_output_addrs = [self._reply_addr(i, host) for i in range(proc_num)]
         self.kv_outputs_queues: list[queue.Queue] = [
             queue.Queue() for _ in range(proc_num)
         ]
         self._pending_kv_aggregation: _PendingKvAggregation | None = None
         self.kv_output_threads: list[threading.Thread] = []
 
         # Generic collective-RPC reply channels, one per rank
-        self.rpc_output_addrs = [get_open_zmq_ipc_path() for _ in range(proc_num)]
+        self.rpc_output_addrs = [self._reply_addr(i, host) for i in range(proc_num)]
         self.rpc_outputs_queues: list[queue.Queue] = [
             queue.Queue() for _ in range(proc_num)
         ]
@@ -509,7 +531,7 @@ def __init__(
         # rank -> exit code, written by the monitor before it tears down.
         self._dead_ranks: dict[int, int | None] = {}
 
-        for i in range(proc_num):
+        for i in range(self.local_proc_num):
             label = f"ModelRunner{i}/{proc_num}"
             # Only rank 0 gets the primary output address
             addrs = [None, io_addrs[1]] if i == 0 else [None, None]
@@ -569,6 +591,32 @@ def __init__(
             self.rpc_output_threads.append(t)
 
         self.monitor_procs()
+        if n_remote:
+            from atom.model_engine.tp_multinode import (
+                publish_bootstrap,
+                release_remote_readers,
+            )
+
+            publish_bootstrap(
+                args[0],
+                {
+                    "proc_num": proc_num,
+                    "handle": scheduler_output_handle,
+                    "runner": runner,
+                    "args": args,
+                    "kv_output_addrs": self.kv_output_addrs,
+                    "rpc_output_addrs": self.rpc_output_addrs,
+                },
+            )
+            release_remote_readers(
+                self.rpc_broadcast_mq, n_remote, lambda: self.still_running
+            )
+
+    def _reply_addr(self, rank: int, host: str | None) -> str:
+        """Where `rank` sends its replies: IPC from this node, TCP from another."""
+        if rank < self.local_proc_num:
+            return get_open_zmq_ipc_path()
+        return get_tcp_uri(host, get_open_port())
 
     def exit(self):
         if not self.still_running:
@@ -756,6 +804,18 @@ def _collective_rpc(
                 for rank in range(self.proc_num)
             ]
 
+        if payload.barrier and getattr(self, "local_proc_num", self.proc_num) < (
+            self.proc_num
+        ):
+            error = (
+                f"cannot run barrier RPC {func_name!r}: the TP ranks on other "
+                f"nodes share no barrier with this one"
+            )
+            return [
+                RpcResult(payload.request_id, rank, error=error)
+                for rank in range(self.proc_num)
+            ]
+
         logger.debug(
             f"{self.label}: collective_rpc {func_name} id={payload.request_id}"
         )

atom/model_engine/engine_core.py

@@ -133,6 +133,7 @@ def __init__(self, config: Config, input_address: str, output_address: str):
                 config.tp_world_size * config.prefill_context_parallel_size,
                 config.runner_qualname,
                 config,
+                local_proc_num=config.tp_size_local if config.is_multinode_tp else None,
                 on_worker_failure=self._send_engine_dead,
             )
             self._post_model_load_hook()

atom/model_engine/engine_core_mgr.py

@@ -382,8 +382,8 @@ def _start_distributed_store(
         pc = config.parallel_config
         if multinode and pc.data_parallel_base_port == 0:
             raise ValueError(
-                "Multi-node DP requires the same nonzero --data-parallel-base-port "
-                "(or ATOM_DP_BASE_PORT) on every node."
+                "Multi-node DP or TP requires the same nonzero "
+                "--data-parallel-base-port (or ATOM_DP_BASE_PORT) on every node."
             )
         if not multinode or pc.data_parallel_rank == 0:
             store = TCPStore(
@@ -482,7 +482,9 @@ def __init__(self, config: Config):
         local_dp_ranks = []
 
         try:
-            self._start_distributed_store(config, multinode=multinode)
+            self._start_distributed_store(
+                config, multinode=multinode or config.is_multinode_tp
+            )
             for engine_index in range(self.local_engine_count):
                 assignment_index = engine_index // self.pp_size
                 dp_rank, local_dp_rank = rank_assignments[assignment_index]