[HIP] [OPUS] [FlyDSL] Bound AITER worker and compiler fanout

.claude/skills/opus-module-build-optimization/SKILL.md

@@ -45,7 +45,7 @@ for tu in *.cuda.o; do
 done | sort -rn | head -10
 ```
 
-Top 1-2 TUs ARE the build wall under `MAX_JOBS >= N`. Optimize those; don't touch the long tail.
+Top 1-2 TUs ARE the build wall under `AITER_MAX_JOBS >= N`. Optimize those; don't touch the long tail.
 
 ### Per-pass timing (host vs device)
 
@@ -209,7 +209,7 @@ If every per-instance device.cu redundantly emits instantiations of a shared sub
 
 **Fix**: emit one `instances/shared_kernel.device.cu` carrying all instantiations of the shared sub-kernel; remove them from each per-instance device.cu. Forward declarations in the fused host TU still resolve at link time.
 
-Per-instance TU walls drop; CPU footprint shrinks proportionally to N. End-to-end wall doesn't always move (the slowest TU may not be the shared one), but helps under constrained `MAX_JOBS`.
+Per-instance TU walls drop; CPU footprint shrinks proportionally to N. End-to-end wall doesn't always move (the slowest TU may not be the shared one), but helps under constrained `AITER_MAX_JOBS`.
 
 **When NOT to apply**: fewer than ~5 instances of the shared sub-kernel.
 

.github/scripts/aiter_prebuild_summary.py

@@ -58,7 +58,7 @@ def main() -> int:
     print(f"Runner: {os.environ.get('AITER_RUNNER_NAME', 'unknown')}")
     print(f"GPU_ARCHS: {os.environ.get('GPU_ARCHS', 'unknown')}")
     print(f"PREBUILD_KERNELS: {os.environ.get('PREBUILD_KERNELS', 'unknown')}")
-    print(f"MAX_JOBS: {os.environ.get('MAX_JOBS', 'unknown')}")
+    print(f"AITER_MAX_JOBS: {os.environ.get('AITER_MAX_JOBS', 'unknown')}")
     print(f"Build status: {args.build_status}")
     print(f"Prebuild wall time: {wall_seconds}s ({wall_seconds / 60:.1f} min)")
     print(f"Kernel count: {len(kernels)}")

.github/workflows/aiter-test.yaml

@@ -21,7 +21,7 @@ env:
   GPU_ARCH_LIST: "gfx942;gfx950"
   AITER_TEST: "op_tests"
   AITER_WHEEL_ARTIFACT_NAME: aiter-whl-${{ github.run_id }}
-  AITER_PREBUILD_MAX_JOBS: "64"
+  AITER_MAX_JOBS: "64"
 
 jobs:
   detect_aiter_test_scope:
@@ -240,7 +240,7 @@ jobs:
           set -euo pipefail
           docker run --rm \
             --network=host \
-            -e AITER_PREBUILD_MAX_JOBS="${{ env.AITER_PREBUILD_MAX_JOBS }}" \
+            -e AITER_MAX_JOBS="${{ env.AITER_MAX_JOBS }}" \
             -e AITER_RUNNER_NAME="${RUNNER_NAME:-unknown}" \
             -v "${{ github.workspace }}:/workspace" \
             -w /workspace \
@@ -257,9 +257,9 @@ jobs:
               pip install --upgrade "ninja>=1.11.1" &&
               pip install --upgrade setuptools_scm &&
               pip install tabulate &&
-              echo "Prebuilding kernels with GPU_ARCHS: ${{ env.GPU_ARCH_LIST }}, PREBUILD_KERNELS: 1, and MAX_JOBS: ${AITER_PREBUILD_MAX_JOBS}" &&
+              echo "Prebuilding kernels with GPU_ARCHS: ${{ env.GPU_ARCH_LIST }}, PREBUILD_KERNELS: 1, and AITER_MAX_JOBS: ${AITER_MAX_JOBS}" &&
               export PREBUILD_KERNELS=1 &&
-              export MAX_JOBS="${AITER_PREBUILD_MAX_JOBS}" &&
+              export AITER_MAX_JOBS="${AITER_MAX_JOBS}" &&
               export GPU_ARCHS="${{ env.GPU_ARCH_LIST }}" &&
               prebuild_start=$(date +%s) &&
               set +e &&
@@ -290,7 +290,7 @@ jobs:
           AITER_RUNNER_NAME="${RUNNER_NAME:-unknown}" \
           GPU_ARCHS="${{ env.GPU_ARCH_LIST }}" \
           PREBUILD_KERNELS=1 \
-          MAX_JOBS="${{ env.AITER_PREBUILD_MAX_JOBS }}" \
+          AITER_MAX_JOBS="${{ env.AITER_MAX_JOBS }}" \
           python3 .github/scripts/aiter_prebuild_summary.py \
             --log .aiter-prebuild.log \
             --build-status "${BUILD_STATUS}" \

.github/workflows/atom-test.yaml

@@ -257,7 +257,7 @@ jobs:
               cd /app/aiter-test && \\
               git checkout --force ${{ env.GITHUB_COMMIT_SHA }} && \\
               git submodule sync && git submodule update --init --recursive && \\
-              MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e . && \\
+              AITER_MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e . && \\
               ./.github/scripts/install_triton.sh
           RUN echo "=== triton version ===" && pip show triton || true
           RUN echo "=== Aiter version AFTER installation ===" && pip show amd-aiter || true

.github/workflows/kimi-downstream.yaml

@@ -109,7 +109,7 @@ jobs:
             pip uninstall -y amd-aiter aiter 2>/dev/null || true
             pip install --upgrade "pybind11>=3.0.1"
             cd /workspace
-            MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e .
+            AITER_MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e .
             (pip show amd-aiter || pip show aiter) | grep -E "^Version|^Location" || true
           '
 

.github/workflows/kimi-perf-downstream.yaml

@@ -93,7 +93,7 @@ jobs:
             pip uninstall -y amd-aiter aiter 2>/dev/null || true
             pip install --upgrade "pybind11>=3.0.1"
             cd /workspace
-            MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e .
+            AITER_MAX_JOBS=64 PREBUILD_KERNELS=0 GPU_ARCHS=gfx950 pip install -e .
           '
 
       - name: Run Kimi-K2.5 throughput sweep

.github/workflows/vllm-disagg-ci-smoke-workflow.yaml

@@ -42,7 +42,7 @@ env:
   DISAGG_SCRIPTS_STAGE: ${{ inputs.disagg_scripts_stage || '/data/xinhuang/vllm-disagg-logs/' }}
   SPUR_NODE_POOL: ${{ inputs.spur_node_pool || 'ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-5,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-6,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-7,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-8,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-13,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-15,ml-ai-ubuntu-gpu-mi350x8-2304gb-fabric-16' }}
   AITER_GPU_ARCHS: gfx950
-  AITER_BUILD_MAX_JOBS: "64"
+  AITER_MAX_JOBS: "64"
   AITER_WHEEL_ARTIFACT_NAME: aiter-vllm-disagg-wheel-${{ github.run_id }}
 
 jobs:
@@ -101,7 +101,7 @@ jobs:
             -e AITER_USE_SYSTEM_TRITON=1 \
             -e GPU_ARCHS="${AITER_GPU_ARCHS}" \
             -e PREBUILD_KERNELS=0 \
-            -e MAX_JOBS="${AITER_BUILD_MAX_JOBS}" \
+            -e AITER_MAX_JOBS="${AITER_MAX_JOBS}" \
             --entrypoint /bin/bash \
             "${BASE_IMAGE}" \
             -lc '
@@ -129,7 +129,7 @@ jobs:
                 "flydsl==${flydsl_version}"
               echo "Prebuilding the gfx950 CKTile A8W8 blockscale module"
               python3 .github/scripts/prebuild_vllm_disagg_cktile.py
-              echo "Building AITER wheel for GPU_ARCHS=${GPU_ARCHS}, PREBUILD_KERNELS=${PREBUILD_KERNELS}, MAX_JOBS=${MAX_JOBS}"
+              echo "Building AITER wheel for GPU_ARCHS=${GPU_ARCHS}, PREBUILD_KERNELS=${PREBUILD_KERNELS}, AITER_MAX_JOBS=${AITER_MAX_JOBS}"
               python setup.py bdist_wheel 2>&1 | tee .aiter-build.log
               ls -lh dist/*.whl
               aiter_wheels=(dist/amd_aiter-*.whl)

README.md

@@ -107,7 +107,8 @@ Release automation validates that the release tag points at the matching release
 ```bash
 git clone --recursive https://github.com/ROCm/aiter.git
 cd aiter
-python3 setup.py develop
+./.github/scripts/install_triton.sh
+AITER_USE_SYSTEM_TRITON=1 python3 -m pip install -e .
 ```
 
 If you happen to forget the `--recursive` during `clone`, you can use the following command after `cd aiter`
@@ -117,7 +118,7 @@ git submodule sync && git submodule update --init --recursive
 
 ### FlyDSL
 
-AITER uses [FlyDSL](https://github.com/ROCm/FlyDSL)-based kernels across a range of operators (e.g., GEMM and MoE). FlyDSL is a required dependency and is installed automatically when you run `python3 setup.py develop`.
+AITER uses [FlyDSL](https://github.com/ROCm/FlyDSL)-based kernels across a range of operators (e.g., GEMM and MoE). FlyDSL is a required dependency and is installed automatically when you run `python3 -m pip install -e .`.
 
 To install it manually:
 
@@ -127,18 +128,27 @@ pip install -r requirements.txt
 
 ### Triton
 
-AITER includes Triton-based operators that require triton from AMD PyPI, with the correct version selected based on your ROCm installation.
+AITER includes Triton-based operators that require Triton from AMD PyPI, with the correct version selected based on your ROCm installation. For an editable installation, run the install script before `pip install -e .` as shown above.
 
-If you install with `python3 setup.py develop`, triton is installed automatically. To skip this and keep your existing triton, set:
+To keep an existing compatible Triton installation, skip the script and set:
 
 ```bash
-AITER_USE_SYSTEM_TRITON=1 python3 setup.py develop
+AITER_USE_SYSTEM_TRITON=1 python3 -m pip install -e .
 ```
 
-If you use `pip install -e .`, run the install script manually:
+### Build parallelism
+
+`AITER_MAX_JOBS` is AITER's only top-level compilation-worker override. A generic `MAX_JOBS` inherited from vLLM, SGLang, PyTorch, or another parent framework is ignored and left unchanged. When `AITER_MAX_JOBS` is unset, AITER selects the smaller of 80% of the CPUs available to the current process and effective available memory divided by the observed 1.5 GB RSS estimate per worker. Effective available memory is the smaller of host `MemAvailable` and remaining cgroup v2/v1 memory when a finite container limit is present.
+
+An explicit `AITER_MAX_JOBS` is an unsafe expert override. It is honored verbatim except that non-positive values become one, bypasses automatic CPU and memory safety caps, and can oversubscribe CPUs or cause an OOM.
+
+Process-pool workers force nested AITER, Ninja, CMake, Make, OpenMP, BLAS, and NumExpr compilation fanout to one.
+
+Examples:
 
 ```bash
-./.github/scripts/install_triton.sh
+AITER_MAX_JOBS=8 python3 -m pip install -e .
+AITER_MAX_JOBS=8 python3 -m aiter.aot.pa
 ```
 
 ### Opus — Lightweight C++ Template for Kernel Development

aiter/aot/asm_mla_decode_fwd.py

@@ -1,17 +1,16 @@
 import concurrent.futures
-import os
 from collections import namedtuple
 
+from aiter_worker_limits import configure_worker_subprocesses, get_worker_count_for
 from csrc.cpp_itfs.mla.asm_mla_decode_fwd import compile
-from csrc.cpp_itfs.utils import AITER_CORE_DIR
 
 MLAConfig = namedtuple(
     "MLAConfig",
     [
-        "hsaco_path",
+        "gqa_ratio",
         "page_size",
-        "q_itemsize",
-        "kv_itemsize",
+        "q_dtype",
+        "kv_dtype",
         "num_kv_splits",
         "v_head_dim",
     ],
@@ -20,10 +19,10 @@
 
 def process_config(config):
     return compile(
-        config.hsaco_path,
+        config.gqa_ratio,
         config.page_size,
-        config.q_itemsize,
-        config.kv_itemsize,
+        config.q_dtype,
+        config.kv_dtype,
         config.num_kv_splits,
         config.v_head_dim,
     )
@@ -34,19 +33,21 @@ def main():
     for num_kv_splits in range(1, 17):
         configs.append(
             MLAConfig(
-                hsaco_path=f"{AITER_CORE_DIR}/hsa/mla_stage1_a16w16_bf16.co",
+                gqa_ratio=16,
                 page_size=1,
-                q_itemsize=2,
-                kv_itemsize=2,
+                q_dtype="__hip_bfloat16",
+                kv_dtype="__hip_bfloat16",
                 num_kv_splits=num_kv_splits,
                 v_head_dim=512,
             )
         )
 
     with concurrent.futures.ProcessPoolExecutor(
-        os.environ.get("MAX_JOBS", "16")
+        max_workers=get_worker_count_for(len(configs)),
+        initializer=configure_worker_subprocesses,
     ) as executor:
-        executor.map(process_config, configs)
+        # Consume the iterator so worker compilation errors reach the caller.
+        list(executor.map(process_config, configs))
 
 
 if __name__ == "__main__":

aiter/aot/flydsl/README.md

@@ -63,8 +63,7 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 | --- | --- | --- |
 | `AITER_AOT_IMPORT` | Set to `1` so `import aiter` only loads the lightweight JIT core and skips the full top-level op namespace — faster and avoids heavy import side effects during AOT compilation (this is what `setup.py` sets while pre-compiling). | `0` |
 | `FLYDSL_RUNTIME_CACHE_DIR` | Cache directory | `~/.flydsl/cache` |
-| `AITER_FLYDSL_AOT_WORKERS` | Max concurrent worker processes. Set explicitly to honor it verbatim (bypasses the memory cap below); `0`/negative clamps to 1. Each worker uses ~1.5–2.5 GB RSS. | `min(affinity-aware CPUs, 64)`, then capped by available memory |
-| `AITER_FLYDSL_AOT_MEM_PER_WORKER_GB` | Assumed GiB/worker for the **auto memory cap** that keeps the OOM-killer from firing. Only applies when `AITER_FLYDSL_AOT_WORKERS` is **not** set; `0` disables the cap. | `2.0` |
+| `AITER_MAX_JOBS` | AITER-local maximum concurrent CPU compilation workers. When unset, the shared policy selects the minimum of 80% of process-available CPU cores and effective host/container available memory divided by the observed 1.5 GB worker RSS estimate. An explicit value is an unsafe expert override: it bypasses those CPU/memory caps, may oversubscribe or OOM, and only normalizes non-positive values to 1. | auto |
 | `AITER_FLYDSL_AOT_TIMEOUT` | Per-kernel wall-clock cap (seconds). A worker stuck *alive* past this is killed (and retried); `0` disables. | `1200` |
 | `AITER_FLYDSL_AOT_MAX_RETRIES` | Retries for a worker that **died abnormally** (OOM-kill / segfault / timeout-kill). A clean compile error is never retried. `0` disables. | `2` |
 | `AITER_CONFIGS` | Resolves the default CSV lookup path (same as the runtime JIT) | repo built-in |
@@ -80,7 +79,7 @@ python -m aiter.aot.flydsl.chunk_gdn_h --csv /path/to/tuned.csv
 Example:
 
 ```bash
-AITER_FLYDSL_AOT_WORKERS=16 python -m aiter.aot.flydsl.moe
+AITER_MAX_JOBS=16 python -m aiter.aot.flydsl.moe
 ```
 
 ---
@@ -135,8 +134,9 @@ python op_tests/test_moe_2stage.py
   is refreshed to the right version.
 - **Worker OOM / killed (exitcode -9)**: abnormal exits are auto-retried
   (`AITER_FLYDSL_AOT_MAX_RETRIES`) and the default worker count is already
-  memory-capped (`AITER_FLYDSL_AOT_MEM_PER_WORKER_GB`). If it still happens,
-  lower `AITER_FLYDSL_AOT_WORKERS` or raise the assumed GiB/worker.
+  bounded by the shared CPU and memory policy. If it still happens,
+  lower `AITER_MAX_JOBS`. Nested compiler fanout is fixed at one inside each
+  AOT worker.
 - **A kernel hangs / never finishes**: it is killed once it exceeds
   `AITER_FLYDSL_AOT_TIMEOUT` (default 1200 s) and then retried. Lower the timeout
   to fail faster, or raise it for genuinely slow kernels.

aiter/aot/flydsl/common.py

@@ -18,10 +18,10 @@
 from multiprocessing.connection import wait as wait_for_sentinels
 from typing import Any
 
+from aiter_worker_limits import configure_worker_subprocesses, get_worker_count_for
+
 _DEFAULT_KERNEL_TIMEOUT = 1200.0
-_DEFAULT_MAX_WORKERS = 64
 _DEFAULT_MAX_RETRIES = 2
-_DEFAULT_MEM_PER_WORKER_GB = 2.0
 _MAX_ERRORS_IN_MSG = 10
 
 
@@ -174,24 +174,14 @@ def _compile_one_config_for(kind: OpKind) -> Callable[..., dict[str, Any]]:
 def _run_one_to_file(
     worker: Callable[..., dict[str, Any]], kwargs: dict[str, Any], out_path: str
 ) -> None:
+    configure_worker_subprocesses()
     result = worker(**kwargs)
     tmp_path = out_path + ".tmp"
     with open(tmp_path, "w") as f:
         json.dump(result, f)
     os.replace(tmp_path, out_path)
 
 
-def _affinity_aware_cpu_count() -> int:
-    """Number of CPUs this process may actually use (respects cgroup /
-    cpuset limits via ``sched_getaffinity``, unlike ``os.cpu_count()``).
-    Falls back to ``cpu_count`` and is clamped to >=1."""
-    try:
-        n = len(os.sched_getaffinity(0))
-    except (AttributeError, OSError):
-        n = os.cpu_count() or 0
-    return max(n, 1)
-
-
 def get_kernel_timeout() -> float:
     env = os.environ.get("AITER_FLYDSL_AOT_TIMEOUT")
     if env is None:
@@ -216,39 +206,8 @@ def get_max_retries() -> int:
         ) from e
 
 
-def _memory_worker_cap(default_workers: int) -> int:
-    env = os.environ.get("AITER_FLYDSL_AOT_MEM_PER_WORKER_GB")
-    try:
-        per_gb = float(env) if env else _DEFAULT_MEM_PER_WORKER_GB
-    except ValueError as e:
-        raise ValueError(
-            f"AITER_FLYDSL_AOT_MEM_PER_WORKER_GB must be a number, got {env!r}"
-        ) from e
-    if per_gb <= 0:
-        return default_workers
-    try:
-        import psutil
-
-        avail_gb = psutil.virtual_memory().available / (1024**3)
-    except Exception:  # noqa: BLE001
-        return default_workers
-    return min(default_workers, max(1, int(avail_gb / per_gb)))
-
-
 def get_max_workers(num_jobs: int) -> int:
-    workers_env = os.environ.get("AITER_FLYDSL_AOT_WORKERS")
-    if workers_env is not None:
-        try:
-            max_workers = max(int(workers_env), 1)
-        except ValueError as e:
-            raise ValueError(
-                f"AITER_FLYDSL_AOT_WORKERS must be an integer, got {workers_env!r}"
-            ) from e
-    else:
-        max_workers = min(_affinity_aware_cpu_count(), _DEFAULT_MAX_WORKERS)
-        # Auto path only: also bound by memory so we never trip the OOM-killer.
-        max_workers = _memory_worker_cap(max_workers)
-    return min(max_workers, num_jobs)
+    return get_worker_count_for(num_jobs)
 
 
 def _run_file_pool(

aiter/aot/pa.py

@@ -1,7 +1,7 @@
 import concurrent.futures
-import os
 from collections import namedtuple
 
+from aiter_worker_limits import configure_worker_subprocesses, get_worker_count_for
 from csrc.cpp_itfs.pa.pa import compile
 
 PAConfig = namedtuple(
@@ -95,9 +95,10 @@ def main():
                         )
 
     with concurrent.futures.ProcessPoolExecutor(
-        os.environ.get("MAX_JOBS", "16")
+        max_workers=get_worker_count_for(len(configs)),
+        initializer=configure_worker_subprocesses,
     ) as executor:
-        executor.map(process_config, configs)
+        list(executor.map(process_config, configs))
 
 
 if __name__ == "__main__":

aiter/aot/pa_ragged.py

@@ -1,7 +1,7 @@
 import concurrent.futures
-import os
 from collections import namedtuple
 
+from aiter_worker_limits import configure_worker_subprocesses, get_worker_count_for
 from csrc.cpp_itfs.pa.pa_ragged import compile
 
 PAConfig = namedtuple(
@@ -95,9 +95,10 @@ def main():
                         )
 
     with concurrent.futures.ProcessPoolExecutor(
-        os.environ.get("MAX_JOBS", "16")
+        max_workers=get_worker_count_for(len(configs)),
+        initializer=configure_worker_subprocesses,
     ) as executor:
-        executor.map(process_config, configs)
+        list(executor.map(process_config, configs))
 
 
 if __name__ == "__main__":