feat(setup): install ATOM on bare metal and fold it into the 12h demos; fix the failures of its 12h release run

docs/compatibility.rst

@@ -54,7 +54,8 @@ The following table lists the minimum requirements for running Hyperloom.
 | Python              | >= 3.10                                                |
 +---------------------+--------------------------------------------------------+
 | Inference Framework | SGLang (>= 0.5.12), vLLM (>= 0.21.0),                  |
-|                     | ATOM (preinstalled; see below), plus ``custom``        |
+|                     | ATOM (preinstalled or setup-installed; see below),     |
+|                     | plus ``custom``                                        |
 |                     | benchmark script                                       |
 +---------------------+--------------------------------------------------------+
 | Kernel Languages    | HIP, Triton, FlyDSL                                    |
@@ -131,7 +132,7 @@ The following inference frameworks are supported:
        frameworks within one session
    * - ATOM
      - 7.2.4 (recorded image stack)
-     - AMD out-of-tree engine, launched as ``python3 -m atom.entrypoints.openai_server``. Supports Docker or direct execution in a preinstalled ATOM/ROCm torch environment. The recorded image-based stack used ``rocm/atom-dev:v0.1.7-rc0`` on MI355X, not a universal ATOM pip-version minimum. Other builds require local validation. An unset or empty ``KERNEL_OPT_BACKEND_ORDER`` resolves to GEAK here as it does for every framework, and explicit values are preserved.
+     - AMD out-of-tree engine, launched as ``python3 -m atom.entrypoints.openai_server``. Supports Docker, direct execution in a preinstalled ATOM/ROCm torch environment, or a bare-metal source install with ``--install-framework atom``. The recorded image-based stack used ``rocm/atom-dev:v0.1.7-rc0`` on MI355X, not a universal ATOM pip-version minimum. Other builds require local validation. An unset or empty ``KERNEL_OPT_BACKEND_ORDER`` resolves to GEAK here as it does for every framework, and explicit values are preserved.
    * - ``custom``
      - Host-defined
      - Escape hatch for your own benchmark script; Hyperloom does not manage the server lifecycle. Requires ``HYPERLOOM_BENCHMARK_BACKEND=bypass`` plus ``--framework-path`` (or ``FRAMEWORK_REPO_PATH``) and ``--benchmark-scripts-dir`` (or ``HYPERLOOM_BYPASS_SCRIPTS_DIR``); the CLI exits with status 2 when any of the three is missing.
@@ -188,8 +189,11 @@ Hyperloom does not install ROCm or torch itself.
 
 For ATOM, ``baremetal`` means running directly in the development machine's
 selected Python environment, including when the development platform itself is
-a container; it does not start an additional Docker container. ATOM must already
-be installed. Verify a real ``import atom`` and a non-empty ``torch.version.hip``
+a container; it does not start an additional Docker container. ATOM is either
+already installed there or installed by setup with ``--install-framework atom``,
+which builds AITER and then ATOM (default ``ATOM_REF``: the commit
+``rocm/atom-dev:v0.1.7-rc0`` was built from) against the existing ROCm torch in
+that Python. Verify a real ``import atom`` and a non-empty ``torch.version.hip``
 with that Python, keep its executable first on ``PATH`` for Magpie's ``python3``
 launch, and use ``PYTHON`` with ``INFERENCE_OPTIMIZER_FORCE_PYTHON=1`` if pinning
 the interpreter. Keep any activated venv consistent; ``/opt/venv`` is not required.
@@ -199,9 +203,9 @@ Run ``python -m hyperloom.inference_optimizer.setup --check-only --
 selected interpreter first. Only after approval, repeat without ``--check-only``:
 ``none`` skips framework installation but can still write configuration and
 apply ROCm hotfixes. Preserve the selected ``USER_DATA_PATH``; an existing setup
-need not be repeated. Hyperloom does not install ATOM and does not assert a
-minimum ATOM package version derived from a Docker tag. Import checks establish
-local prerequisites, not end-to-end validation of every stack.
+need not be repeated. Hyperloom does not assert a minimum ATOM package version
+derived from a Docker tag. Import checks establish local prerequisites, not
+end-to-end validation of every stack.
 
 .. list-table::
    :header-rows: 1

docs/install/install.md

@@ -132,7 +132,7 @@ The backend runs `install_baremetal.sh` in five phases:
 
 1. **Base preflight**: Checks ROCm, GPU arch, ROCm torch, torch/triton alignment,
    and serving framework imports.
-2. **Framework install**: Optionally installs the SGLang or vLLM framework layer.
+2. **Framework install**: Optionally installs the SGLang, vLLM or ATOM framework layer.
 3. **ROCm hotfix**: Applies the profiler hotfix when the ROCm stack is eligible,
    covering both `/opt/rocm/lib` and PyTorch's bundled `torch/lib/`.
 4. **Credentials**: Resolves LLM gateway credentials into `.env`.
@@ -396,7 +396,7 @@ set of accepted shapes, including split entrypoints and self-hosted gateways.
 Make sure the host already provides the required base environment:
 
 - ROCm runtime and a ROCm-built torch.
-- A serving framework (SGLang or vLLM) importable in the active Python.
+- A serving framework (SGLang, vLLM or ATOM) importable in the active Python.
 - `git` for the dependency checkouts the optimization skill performs.
 
 With that in place, open the repository root in the agent and paste a launch

docs/reference/environment-variables.md

@@ -721,8 +721,9 @@ The following variables configure framework source discovery and path overrides.
 | `INFERENCE_`<br>`OPTIMIZER`<br>`_STRICT_PATHS`                | `1` when CLI bootstraps                                                | When `1`, missing path env raises instead of falling back to discovery. Set by the CLI at session start; do not override unless debugging.              |
 | `HYPERLOOM_`<br>`SGLANG_PA`<br>`TCH_EXACT`<br>`_VERSIONS`           | Unset                                                                  | Pin the sglang server-patch step to specific upstream versions; advanced compatibility option.                                                          |
 | `HYPERLOOM_`<br>`ENABLE`<br>`_PATCH`                          | `1`                                                                    | Set to `0` to skip the in-place server patch step (useful when the upstream is already pre-patched).                                                    |
-| `HYPERLOOM_`<br>`SKIP_FRAME`<br>`WORK_CHECK`                  | Unset (check enabled)                                                  | Truthy skips the `optimize` preflight gate that requires the selected serving framework to be importable and a ROCm build. Last resort: when the server runs elsewhere, set `BENCHMARK_BASE_URL` instead, which exempts the check and configures the supported path. The gate already stays out of the way for `xdit`/`custom` (server-less), external multi-node, and any framework `install_baremetal.sh` cannot install (`atom`), where it warns instead of blocking. |
+| `HYPERLOOM_`<br>`SKIP_FRAME`<br>`WORK_CHECK`                  | Unset (check enabled)                                                  | Truthy skips the `optimize` preflight gate that requires the selected serving framework to be importable and a ROCm build. Last resort: when the server runs elsewhere, set `BENCHMARK_BASE_URL` instead, which exempts the check and configures the supported path. The gate already stays out of the way for `xdit`/`custom` (server-less) and external multi-node. |
 | `AITER_REF` | Unset | Optional bare-metal AITER install pin. When unset, the installer selects the newest tag compatible with the installed torch/triton stack. |
+| `ATOM_REF` | `fe1099b15ddc52e6873e1932935f722864cdeee0` | Full commit SHA that `--install-framework atom` checks out from `ATOM_REPO` (default `https://github.com/ROCm/ATOM.git`) into `ATOM_ROOT` (default `<deps root>/atom`). The default is the commit `rocm/atom-dev:v0.1.7-rc0` was built from. |
 
 ---
 

examples/README.md

@@ -124,7 +124,7 @@ The backend runs `install_baremetal.sh` in five phases:
 
 1. **Base preflight**: checks ROCm, GPU arch, ROCm torch, torch/triton alignment,
    and serving framework imports.
-2. **Framework install**: optionally installs the SGLang or vLLM framework layer.
+2. **Framework install**: optionally installs the SGLang, vLLM or ATOM framework layer.
 3. **ROCm hotfix**: applies the profiler hotfix when the ROCm stack is eligible.
 4. **Credentials**: resolves LLM gateway credentials into `.env`.
 5. **Runtime env**: persists bare-metal runtime vars (framework, ROCm/venv roots,

examples/hyperloom-qwen3-14b-fp8-12h-atom-forge/SKILL.md

@@ -140,8 +140,11 @@ Both modes use the following steps in the shell where ATOM will run.
 
 ### Framework
 
-ATOM must already be installed with ROCm torch. Hyperloom does not install ATOM:
-use `--install-framework none --frameworks atom --require-frameworks` below.
+ATOM runs on the environment's ROCm torch. When ATOM is already installed, use
+`--install-framework none --frameworks atom --require-frameworks` below. On a
+baremetal environment without ATOM, setup can install it with
+`--install-framework atom` (see the setup skill); do not do this in Docker mode,
+where ATOM comes from the image.
 This example is single-node, fixes `FRAMEWORK=atom`, and lets Hyperloom detect the
 GPU rather than passing `--gpu-type`.
 
@@ -194,6 +197,16 @@ changes and obtain approval before running:
   --user-data-path "${USER_DATA_PATH:?USER_DATA_PATH missing}" --yes
 ```
 
+If the check reports ATOM missing in a baremetal environment and the operator
+approves installing it, run setup with `atom` instead of `none`. It installs
+AITER and ATOM from source into the selected Python, keeping its ROCm torch:
+
+```bash
+"$PYTHON" -m hyperloom.inference_optimizer.setup -- \
+  --install-framework atom \
+  --user-data-path "${USER_DATA_PATH:?USER_DATA_PATH missing}" --yes
+```
+
 `none` skips framework installation only: actual setup writes `.env` and may
 apply ROCm hotfixes; `--yes` is not user consent. Do not bypass base checks, repeat
 onboarding unnecessarily, or install SGLang/vLLM to compensate for missing ATOM.

examples/hyperloom-qwen3-14b-fp8-12h-atom/SKILL.md

@@ -134,8 +134,11 @@ Both modes use the following steps in the shell where ATOM will run.
 
 ### Framework
 
-ATOM must already be installed with ROCm torch. Hyperloom does not install ATOM:
-use `--install-framework none --frameworks atom --require-frameworks` below.
+ATOM runs on the environment's ROCm torch. When ATOM is already installed, use
+`--install-framework none --frameworks atom --require-frameworks` below. On a
+baremetal environment without ATOM, setup can install it with
+`--install-framework atom` (see the setup skill); do not do this in Docker mode,
+where ATOM comes from the image.
 This example is single-node, fixes `FRAMEWORK=atom`, and lets Hyperloom detect the
 GPU rather than passing `--gpu-type`.
 
@@ -179,6 +182,16 @@ changes and obtain approval before running:
   --user-data-path "${USER_DATA_PATH:?USER_DATA_PATH missing}" --yes
 ```
 
+If the check reports ATOM missing in a baremetal environment and the operator
+approves installing it, run setup with `atom` instead of `none`. It installs
+AITER and ATOM from source into the selected Python, keeping its ROCm torch:
+
+```bash
+"$PYTHON" -m hyperloom.inference_optimizer.setup -- \
+  --install-framework atom \
+  --user-data-path "${USER_DATA_PATH:?USER_DATA_PATH missing}" --yes
+```
+
 `none` skips framework installation only: actual setup writes `.env` and may
 apply ROCm hotfixes; `--yes` is not user consent. Do not bypass base checks, repeat
 onboarding unnecessarily, or install SGLang/vLLM to compensate for missing ATOM.

src/hyperloom/inference_optimizer/assets/install_baremetal.sh

@@ -6,10 +6,10 @@
 #
 # Runs Hyperloom directly on a host that already provides the ROCm framework
 # base (ROCm runtime + a ROCm-built torch + a serving framework). For bare-metal
-# installs, the script can optionally install SGLang or vLLM ROCm framework layers.
+# installs, the script can optionally install SGLang, vLLM or ATOM ROCm framework layers.
 #
 # Phase 1  base preflight  — ROCm / GPU arch / ROCm torch / serving framework
-# Phase 2  framework       — optional bare-metal SGLang/vLLM install
+# Phase 2  framework       — optional bare-metal SGLang/vLLM/ATOM install
 # Phase 3  ROCm hotfix     — install ROCclr HIP runtime + roctracer profiler fix
 # Phase 4  credentials     — resolve Anthropic/DeepSeek LLM creds into .env
 # Phase 5  runtime env     — persist bare-metal runtime vars into .env
@@ -81,6 +81,10 @@ VLLM_SOURCE_REF="${VLLM_SOURCE_REF:-98dff2a81d747d1dba01a47f939f48c3526d4206}"
 # The source checkout is a depth-1 fetch of a commit SHA and carries no tags, so
 # setuptools_scm would otherwise stamp the build 0.1.dev1 instead of the release.
 VLLM_PRETEND_VERSION="${VLLM_PRETEND_VERSION:-${VLLM_VERSION}}"
+# ATOM publishes no release tags. ATOM_REF is the ATOM_COMMIT build argument of
+# rocm/atom-dev:v0.1.7-rc0, the only ATOM stack docs/compatibility.rst records.
+ATOM_REPO="${ATOM_REPO:-https://github.com/ROCm/ATOM.git}"
+ATOM_REF="${ATOM_REF:-fe1099b15ddc52e6873e1932935f722864cdeee0}"
 VLLM_ROOT="${VLLM_ROOT:-/opt/hyperloom/vllm}"
 # Index publishing the TheRock ROCm SDK wheels these images are built from;
 # rocm-sdk-devel is pulled from here to supply source-build headers.
@@ -101,8 +105,8 @@ usage() {
 Usage: src/hyperloom/inference_optimizer/assets/install_baremetal.sh [options]
 
 Set up a bare-metal host with ROCm + ROCm torch for Hyperloom. Verifies the base,
-optionally installs SGLang/vLLM, resolves credentials, and writes the combined
-runtime env. Stops BEFORE launching.
+optionally installs SGLang/vLLM/ATOM, resolves credentials, and writes the
+combined runtime env. Stops BEFORE launching.
 
 Options:
   --user-data-path PATH  Writable artifact root (default: /workspace/hyperloom)
@@ -111,7 +115,7 @@ Options:
                          sglang,vllm,atom). Phase 1 passes when at least one
                          entry imports.
   --install-framework FW Install a missing bare-metal framework layer.
-                         Supported: none, sglang, vllm. Default: none.
+                         Supported: none, sglang, vllm, atom. Default: none.
   --framework-env MODE   Install target for framework packages: shared or
                          isolated. Default: shared, except vLLM which defaults
                          to isolated so it never replaces the shared ROCm
@@ -141,7 +145,7 @@ SGLANG_REPO, SGLANG_REF, SGLANG_ROOT, SGLANG_ROCM_PYPI_VERSION,
 SGLANG_ROCM_EXTRA, SGLANG_BUILD_RUST_EXTS, AITER_REPO, AITER_REF, AITER_ROOT, ROCM_PATH, HIP_PATH,
 LD_LIBRARY_PATH, ROCM_SDK_INDEX_URL, VLLM_VERSION, VLLM_ROCM_VARIANT, VLLM_ROCM_INDEX,
 VLLM_INSTALL_METHOD, VLLM_REPO, VLLM_SOURCE_REF, VLLM_ROOT, VLLM_VENV_ROOT,
-HYPERLOOM_WHEEL_REPO, HYPERLOOM_WHEEL_TAG.
+ATOM_REPO, ATOM_REF, ATOM_ROOT, HYPERLOOM_WHEEL_REPO, HYPERLOOM_WHEEL_TAG.
 EOF
 }
 
@@ -155,8 +159,8 @@ while [ "$#" -gt 0 ]; do
       shift
       INSTALL_FRAMEWORK="${1:-}"
       case "$INSTALL_FRAMEWORK" in
-        none|sglang|vllm) ;;
-        *) echo "[install-baremetal] ERROR: --install-framework must be one of: none, sglang, vllm" >&2; exit 2 ;;
+        none|sglang|vllm|atom) ;;
+        *) echo "[install-baremetal] ERROR: --install-framework must be one of: none, sglang, vllm, atom" >&2; exit 2 ;;
       esac
       ;;
     --framework-env)
@@ -233,9 +237,9 @@ framework_probe_python() {
 # none is importable. Walks $FRAMEWORKS in order — the same list Phase 1 probes
 # — so an engine that passes preflight is always the one written to .env.
 resolve_installed_framework() {
-  if [ "$INSTALL_FRAMEWORK" = "sglang" ] || [ "$INSTALL_FRAMEWORK" = "vllm" ]; then
-    printf '%s' "$INSTALL_FRAMEWORK"; return 0
-  fi
+  case "$INSTALL_FRAMEWORK" in
+    sglang|vllm|atom) printf '%s' "$INSTALL_FRAMEWORK"; return 0 ;;
+  esac
   local py fw probe_py _rif_arr
   py="$(resolve_python)" || return 0
   IFS=',' read -r -a _rif_arr <<< "$FRAMEWORKS"
@@ -956,6 +960,80 @@ PY
   log "SGLang framework install complete (SGLANG_USE_AITER=${SGLANG_USE_AITER})"
 }
 
+# Editable, as in the rocm/atom-dev image, so framework-agent patches land in
+# the tree the server imports. The Rust atomesh build stays off (ATOM_MESH_BUILD
+# unset): the optimizer launches atom.entrypoints.openai_server, not atomesh.
+install_atom_from_source() {
+  local py="$1" deps_root="$2" atom_root constraint_file
+  atom_root="${ATOM_ROOT:-${deps_root}/atom}"
+
+  log "installing ATOM from source at ${atom_root} (ref=${ATOM_REF})"
+  if [ ! -d "${atom_root}/.git" ]; then
+    mkdir -p "$atom_root"
+    git init -q "$atom_root"
+    git -C "$atom_root" remote add origin "$ATOM_REPO"
+  fi
+  # Fetch the ref rather than `clone --branch`: ATOM_REF is a commit, which
+  # GitHub serves directly only as a full 40-char SHA.
+  git -C "$atom_root" fetch --depth 1 origin "$ATOM_REF"
+  git -C "$atom_root" checkout -q FETCH_HEAD
+
+  constraint_file="$(mktemp)"
+  write_rocm_torch_constraints "$py" "$constraint_file"
+  "$py" -m pip install --constraint "$constraint_file" -e "$atom_root" \
+    || { rm -f "$constraint_file"; die "ATOM ${ATOM_REF} install failed under the current torch/triton constraints"; }
+  rm -f "$constraint_file"
+}
+
+# ATOM runs its kernels through AITER, so AITER goes in first, selected the same
+# way as for SGLang: the newest tag compatible with the installed torch/triton.
+install_atom_framework() {
+  local py deps_root aiter_root
+  py="$(resolve_python)" || die "no usable Python found for ATOM install"
+  deps_root="$(framework_deps_root)"
+  aiter_root="${AITER_ROOT:-${deps_root}/aiter}"
+
+  log "Phase 2: installing ATOM framework layer"
+  log "framework python: ${py}"
+  log "ATOM source: ${ATOM_REPO}@${ATOM_REF}"
+  log "AITER_ROOT=${aiter_root}"
+  if [ -n "$AITER_REF" ]; then
+    log "AITER_REF=${AITER_REF}"
+  else
+    log "AITER_REF=auto (newest tag compatible with installed torch/triton)"
+  fi
+
+  if [ "$CHECK_ONLY" -eq 1 ]; then
+    _py_has "$py" atom && log "atom import OK" \
+      || warn "atom missing (check-only; would install ${ATOM_REPO}@${ATOM_REF})"
+    _py_has "$py" aiter && log "aiter import OK" \
+      || warn "aiter missing (check-only; would install ${AITER_REF:-the newest compatible tag} from ${AITER_REPO})"
+    return 0
+  fi
+
+  if [ "$DRY_RUN" -eq 1 ]; then
+    log "would install AITER ${AITER_REF:-(newest compatible tag)} from ${AITER_REPO} at ${aiter_root} with current torch/triton constraints"
+    log "would clone ${ATOM_REPO}@${ATOM_REF} at ${ATOM_ROOT:-${deps_root}/atom} and install it editable with current torch/triton constraints"
+    return 0
+  fi
+
+  ensure_aiter_for_python "$py" "$aiter_root"
+  if _py_has "$py" atom; then
+    log "atom already importable; skipping ATOM source install"
+  else
+    install_atom_from_source "$py" "$deps_root"
+  fi
+
+  "$py" -c "import aiter" >/dev/null || die "aiter not importable after install"
+  "$py" -c "import atom" >/dev/null || die "atom not importable after install"
+  # Import the server module rather than run `--help`: ATOM's argparse help
+  # text carries unescaped "%" (e.g. at fe1099b15), so rendering help raises
+  # even though the server itself starts and serves.
+  "$py" -c "import atom.entrypoints.openai.api_server" >/dev/null \
+    || die "atom.entrypoints.openai.api_server not importable after install"
+  log "ATOM framework install complete"
+}
+
 # Verify that the installed vLLM package resolves to a ROCm runtime.
 verify_vllm_rocm() {
   local py="$1"
@@ -1451,6 +1529,7 @@ install_requested_framework() {
     none) log "Phase 2: framework install skipped (--install-framework none)" ;;
     sglang) install_sglang_framework ;;
     vllm) install_vllm_framework ;;
+    atom) install_atom_framework ;;
   esac
 }
 
@@ -2609,7 +2688,7 @@ main() {
     shared|isolated) ;;
     *) die "FRAMEWORK_ENV must be one of: shared, isolated" ;;
   esac
-  if [ "$FRAMEWORK_ENV" = "isolated" ] && [ "$INSTALL_FRAMEWORK" = "sglang" ]; then
+  if [ "$FRAMEWORK_ENV" = "isolated" ] && { [ "$INSTALL_FRAMEWORK" = "sglang" ] || [ "$INSTALL_FRAMEWORK" = "atom" ]; }; then
     die "--framework-env isolated is currently supported for vLLM only"
   fi