feat: [AMD Support via SemiAnalysis][3/4] ATOM and AToMesh support

CLAUDE.md

@@ -377,6 +377,7 @@ with patch.dict(os.environ, H100Rack.slurm_env()):
 4. Add polymorphic deserialization in `BackendConfigField` in `schema.py`
 
 **Current backends:**
+- **ATOM**: Native ROCm servers behind AToMesh, with one Slurm node per logical worker and allocator-owned Mooncake handshake ports
 - **SGLang**: Per-process srun launching, supports prefill/decode/aggregated modes
 - **TRTLLM**: MPI-style launching (one srun per endpoint with all nodes), prefill/decode only
 - **vLLM**: Per-process srun launching, prefill/decode/aggregated, `per_node` DP; `frontend_type` selects Dynamo registration or a direct `vllm serve` server, and `_CONNECTOR_MAP` owns the KV connector table

README.md

@@ -1,6 +1,6 @@
 # srtctl
 
-Command-line tool for distributed LLM inference benchmarks on SLURM clusters using TensorRT LLM, SGLang and vLLM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers.
+Command-line tool for distributed LLM inference benchmarks on SLURM clusters using SGLang, vLLM, TensorRT LLM and AMD's ATOM. Replace complex shell scripts and 50+ CLI flags with a declarative `schema: 2` YAML recipe: `engine:` names the engine, `roles:` describes each worker role, and `services:` covers everything launched next to the workers.
 
 ## Quick Start
 

docs/config-reference.md

@@ -43,6 +43,19 @@ This page is the prose guide: what each block means, how the pieces interact, an
 
 ## Overview
 
+### ATOM with AToMesh
+
+Use `engine: atom` with `frontend.type: atomesh` to launch native
+`atom.entrypoints.openai_server` workers and the official AToMesh router. Both
+aggregate workers and prefill/decode topologies use static HTTP endpoints;
+disaggregated workers receive topology-owned Mooncake handshake ports.
+
+Engine flags belong under `roles.prefill.args`, `roles.decode.args`, or
+`roles.agg.args` (schema v2).
+srt-slurm owns the model path, HTTP port, tensor parallel size, and KV-transfer
+contract, so recipes cannot override those arguments. See the complete
+[ATOM/AToMesh recipe](../examples/atom/atomesh-disagg.yaml).
+
 ```yaml
 schema: 2                      # Required: recipe layout version
 name: "my-benchmark"           # Required: job name

docs/schema-reference.md

@@ -13,7 +13,7 @@ Top-level keys of a recipe YAML.
 | `name` | str | required |  |
 | `model` | [ModelConfig](#modelconfig) | required |  |
 | `resources` | [ResourceConfig](#resourceconfig) | required |  |
-| `engine` | str \| mapping | required | The engine type (`sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). |
+| `engine` | str \| mapping | required | The engine type (`atom`, `sglang`, `trtllm`, `vllm`, `mocker`) as a string, or a mapping with `type` plus the engine-wide knobs listed under [Engine types](#engine-types). |
 | `roles` | mapping of role -> [Role](#roles) | required | One block per worker role (`prefill`, `decode`, `agg`): topology, env, and engine args. |
 | `schema` | int | `2` | Recipe schema version. Write `schema: 2` for this layout. |
 | `slurm` | [SlurmConfig](#slurmconfig) | `SlurmConfig()` |  |
@@ -44,7 +44,7 @@ Three vocabularies are specific to the 2.0 layout. They are normalized into the
 
 ### engine
 
-`engine: <type>` or `engine: {type: <type>, ...}`. `type` is one of `sglang`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types).
+`engine: <type>` or `engine: {type: <type>, ...}`. `type` is one of `atom`, `sglang`, `trtllm`, `vllm`, `mocker`; the remaining keys are that engine's knobs, listed under [Engine types](#engine-types).
 
 ### roles
 
@@ -573,6 +573,18 @@ Ready when the service's log file contains a line matching the regular expressio
 
 `engine.type` selects one of the following; the remaining `engine` keys are that type's knobs.
 
+### AtomProtocol
+
+`engine.type: atom`
+
+Launch ``atom.entrypoints.openai_server`` on ROCm workers.
+
+| Key | Type | Default | Description |
+|---|---|---|---|
+| `type` | one of `'atom'` | `'atom'` |  |
+| `connector` | one of `'mooncake'` | `'mooncake'` |  |
+| `mooncake_protocol` | one of `'rdma'`, `'tcp'` \| None | `None` |  |
+
 ### SGLangProtocol
 
 `engine.type: sglang`

examples/README.md

@@ -62,3 +62,9 @@ for p in sorted(Path('examples').rglob('*.yaml')):
     print(validate_config_file(p) or f'ok {p}')
 "
 ```
+
+### ATOM and AToMesh
+
+[`atom/atomesh-disagg.yaml`](atom/atomesh-disagg.yaml) launches native ATOM prefill
+and decode workers with the AToMesh router. Each logical worker fits on one node;
+the cluster configuration selects ROCm device visibility.

examples/atom/atomesh-disagg.yaml

@@ -0,0 +1,32 @@
+# Native ATOM prefill/decode on separate Slurm nodes behind AToMesh.
+# Configure visible_devices_env: ROCR_VISIBLE_DEVICES in the cluster srtslurm.yaml.
+# The container must include both atom.entrypoints.openai_server and atomesh.
+schema: 2
+name: qwen3-0.6b-atomesh-disagg
+model:
+  path: hf:Qwen/Qwen3-0.6B
+  container: rocm/atom:latest
+  precision: bf16
+resources:
+  gpu_type: mi300x
+  gpus_per_node: 8
+engine:
+  type: atom
+  mooncake_protocol: tcp
+frontend:
+  type: atomesh
+  enable_multiple_frontends: false
+roles:
+  prefill:
+    nodes: 1
+    workers: 1
+    gpus: 1
+  decode:
+    nodes: 1
+    workers: 1
+    gpus: 1
+benchmark:
+  type: sa-bench
+  isl: 128
+  osl: 128
+  concurrencies: "4x8"

src/srtctl/backends/__init__.py

@@ -9,16 +9,20 @@
 - TRTLLM: TensorRT-LLM backend with prefill/decode disaggregation
 """
 
+from .atom import AtomProtocol, AtomServerConfig
 from .base import BackendProtocol, BackendType, SrunConfig
 from .mocker import MockerProtocol, MockerServerConfig
 from .sglang import MooncakeKVStoreConfig, SGLangProtocol, SGLangServerConfig
 from .trtllm import TRTLLMProtocol, TRTLLMServerConfig
 from .vllm import VLLMFailoverConfig, VLLMMooncakeKVStoreConfig, VLLMProtocol, VLLMServerConfig
 
 # Union type for all backend configs
-BackendConfig = SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol
+BackendConfig = AtomProtocol | SGLangProtocol | TRTLLMProtocol | VLLMProtocol | MockerProtocol
 
 __all__ = [
+    # ATOM
+    "AtomProtocol",
+    "AtomServerConfig",
     "BackendConfig",
     # Base types
     "BackendProtocol",

src/srtctl/backends/atom.py

@@ -0,0 +1,218 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""ROCm ATOM inference backend."""
+
+from __future__ import annotations
+
+import builtins
+import json
+from collections.abc import Sequence
+from dataclasses import field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
+
+from marshmallow import Schema
+from marshmallow_dataclass import dataclass
+
+from srtctl.ports import DYN_SYSTEM_PORT_BASE
+
+if TYPE_CHECKING:
+    from srtctl.backends.base import SrunConfig
+    from srtctl.core.runtime import RuntimeContext
+    from srtctl.core.schema import ProfilingConfig
+    from srtctl.core.topology import Endpoint, NodePortAllocator, Process
+
+WorkerMode = Literal["prefill", "decode", "agg"]
+
+
+@dataclass(frozen=True)
+class AtomServerConfig:
+    """Native ATOM CLI arguments for each serving role."""
+
+    prefill: dict[str, Any] | None = None
+    decode: dict[str, Any] | None = None
+    aggregated: dict[str, Any] | None = None
+
+    Schema: ClassVar[type[Schema]] = Schema
+
+
+@dataclass(frozen=True)
+class AtomProtocol:
+    """Launch ``atom.entrypoints.openai_server`` on ROCm workers."""
+
+    type: Literal["atom"] = "atom"
+    prefill_environment: dict[str, str] = field(default_factory=dict)
+    decode_environment: dict[str, str] = field(default_factory=dict)
+    aggregated_environment: dict[str, str] = field(default_factory=dict)
+    atom_config: AtomServerConfig | None = None
+    connector: Literal["mooncake"] = "mooncake"
+    mooncake_protocol: Literal["rdma", "tcp"] | None = None
+
+    Schema: ClassVar[builtins.type[Schema]] = Schema
+
+    def get_srun_config(self) -> SrunConfig:
+        from srtctl.backends.base import SrunConfig
+
+        return SrunConfig(mpi=None, oversubscribe=False, launch_per_endpoint=False)
+
+    def get_config_for_mode(self, mode: WorkerMode) -> dict[str, Any]:
+        if self.atom_config is None:
+            return {}
+        values = {
+            "prefill": self.atom_config.prefill,
+            "decode": self.atom_config.decode,
+            "agg": self.atom_config.aggregated,
+        }
+        return dict(values.get(mode) or {})
+
+    def get_environment_for_mode(self, mode: WorkerMode) -> dict[str, str]:
+        values = {
+            "prefill": self.prefill_environment,
+            "decode": self.decode_environment,
+            "agg": self.aggregated_environment,
+        }
+        return dict(values.get(mode) or {})
+
+    def get_process_environment(self, process: Process) -> dict[str, str]:
+        return {}
+
+    def get_served_model_name(self, default: str) -> str:
+        return default
+
+    @property
+    def mooncake_kv_store(self) -> None:
+        return None
+
+    @property
+    def failover(self) -> None:
+        return None
+
+    def get_mooncake_worker_env(self, infra_node_ip: str, local_hostname: str) -> dict[str, str]:
+        return {}
+
+    def get_failover_environment(self, process: Process, job_id: str) -> dict[str, str]:
+        return {}
+
+    def should_set_visible_devices(self) -> bool:
+        return True
+
+    def allocate_endpoints(
+        self,
+        num_prefill: int,
+        num_decode: int,
+        num_agg: int,
+        gpus_per_prefill: int,
+        gpus_per_decode: int,
+        gpus_per_agg: int,
+        gpus_per_node: int,
+        available_nodes: Sequence[str],
+        spread_workers: bool = False,
+    ) -> list[Endpoint]:
+        from srtctl.core.topology import allocate_endpoints
+
+        return allocate_endpoints(
+            num_prefill=num_prefill,
+            num_decode=num_decode,
+            num_agg=num_agg,
+            gpus_per_prefill=gpus_per_prefill,
+            gpus_per_decode=gpus_per_decode,
+            gpus_per_agg=gpus_per_agg,
+            gpus_per_node=gpus_per_node,
+            available_nodes=available_nodes,
+            spread_workers=spread_workers,
+        )
+
+    def endpoints_to_processes(
+        self,
+        endpoints: list[Endpoint],
+        base_sys_port: int = DYN_SYSTEM_PORT_BASE,
+        port_allocator: NodePortAllocator | None = None,
+        frontend_type: str = "atomesh",
+        dynamo_sidecar: bool = False,
+    ) -> list[Process]:
+        if dynamo_sidecar:
+            raise ValueError("ATOM does not support Dynamo sidecars")
+        from srtctl.core.topology import endpoints_to_processes
+
+        return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator)
+
+    def _kv_transfer_config(self, process: Process, worker_ip: str) -> str:
+        if process.endpoint_mode not in {"prefill", "decode"}:
+            raise ValueError("ATOM KV transfer is only valid for prefill/decode workers")
+        if process.nixl_port is None:
+            raise ValueError("ATOM P/D worker is missing its Mooncake handshake port")
+        payload = {
+            "kv_role": "kv_producer" if process.endpoint_mode == "prefill" else "kv_consumer",
+            "kv_connector": self.connector,
+            "proxy_ip": worker_ip,
+            "handshake_port": process.nixl_port,
+        }
+        if self.mooncake_protocol is not None:
+            payload["protocol"] = self.mooncake_protocol
+        return json.dumps(payload, separators=(",", ":"))
+
+    def build_worker_command(
+        self,
+        process: Process,
+        endpoint_processes: list[Process],
+        runtime: RuntimeContext,
+        frontend_type: str = "atomesh",
+        nsys_prefix: list[str] | None = None,
+        dump_config_path: Path | None = None,
+        profiling: ProfilingConfig | None = None,
+    ) -> list[str]:
+        if frontend_type != "atomesh":
+            raise ValueError(f"backend.type: atom requires frontend.type: atomesh (got {frontend_type!r})")
+        if len({item.node for item in endpoint_processes}) != 1:
+            raise ValueError("ATOM currently requires each logical endpoint to fit on one Slurm node")
+
+        from srtctl.core.slurm import get_hostname_ip
+
+        worker_ip = get_hostname_ip(process.node, runtime.network_interface)
+        config = self.get_config_for_mode(process.endpoint_mode)
+        reserved = {"model", "host", "server-port", "tp", "tensor-parallel-size", "kv-transfer-config"}
+        overlap = reserved.intersection(_canonical_arg_key(key) for key in config)
+        if overlap:
+            raise ValueError(f"ATOM config cannot override srtctl-managed argument(s): {sorted(overlap)}")
+
+        command = ["env", f"ATOM_HOST_IP={worker_ip}", *(nsys_prefix or [])]
+        command.extend(
+            [
+                "python3",
+                "-m",
+                "atom.entrypoints.openai_server",
+                "--model",
+                runtime.worker_model_arg,
+                "--host",
+                "0.0.0.0",
+                "--server-port",
+                str(process.http_port),
+                "-tp",
+                str(len(process.gpu_indices)),
+            ]
+        )
+        if process.endpoint_mode in {"prefill", "decode"}:
+            command.extend(["--kv-transfer-config", self._kv_transfer_config(process, worker_ip)])
+        command.extend(_config_to_cli_args(config))
+        return command
+
+
+def _config_to_cli_args(config: dict[str, Any]) -> list[str]:
+    """Preserve ATOM's native CLI spelling, which mixes hyphens and underscores."""
+    args: list[str] = []
+    for key, value in sorted(config.items()):
+        flag = f"--{key}"
+        if value is True:
+            args.append(flag)
+        elif value is False or value is None:
+            continue
+        elif isinstance(value, list):
+            args.extend([flag, *(str(item) for item in value)])
+        else:
+            args.extend([flag, str(value)])
+    return args
+
+
+def _canonical_arg_key(key: str) -> str:
+    return key.lstrip("-").replace("_", "-")

src/srtctl/backends/base.py

@@ -29,6 +29,7 @@ class BackendType(str, Enum):
     TRTLLM = "trtllm"
     VLLM = "vllm"
     MOCKER = "mocker"
+    ATOM = "atom"
 
 
 @dataclass

src/srtctl/core/config.py

@@ -104,6 +104,7 @@ def load_cluster_config() -> dict[str, Any] | None:
         "sbatch_directives",
         "srun_options",
         "sglang_config",
+        "atom_config",
         "vllm_config",
         "trtllm_config",
         "mocker_config",

src/srtctl/core/roles.py

@@ -54,6 +54,7 @@
 
 # backend.type -> the engine's per-mode CLI config key.
 ENGINE_CONFIG_KEY: dict[str, str] = {
+    "atom": "atom_config",
     "sglang": "sglang_config",
     "vllm": "vllm_config",
     "trtllm": "trtllm_config",