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

[Qubitium]

This PR seems large but it all centered around the oom crash I encountered on a single vm with 1 slice of MI350X that randomly killed vm processes (tmux).

Basically there are lots of places of where disjointed fan-out was happening. The fix was a centralize fan-out control for the entire AITER pkg. No more everyone who needs a threadpool doing their own, and wrong, calculations.

The estimated fan-out in this PR is still a best-telemetry based estimate but is much better than the previous static control with sometimes incorrect upper/lower bounds based on out-dated or inaccurate metrics.

I have reviewed every delta and designed the refractor.

https://x.com/qubitium/status/2094621657226985567?s=20

Summary

This draft fixes nine validated bugs present at the origin/main merge base and adds seven worker-policy improvements. The final automatic policy uses process-available CPUs and the lower of host-available and container-available memory; process-pool children always force nested compiler fanout to one.

Validated bugs present in origin/main and fixed here (9, ordered by severity)

1. Top-level AOT limits did not bound nested compiler fanout

• Failure: each AOT process could launch additional AITER, Ninja, CMake, Make, OpenMP, BLAS, and NumExpr work. Top-level workers multiplied by nested compiler workers/threads, reproducing the compiler-process explosion and OOM.
• Fix: each process-pool child forces all nested controls to one at aiter_worker_limits.py:224-240, covered at tests/test_worker_awareness.py:203-217 and tests/test_worker_awareness.py:233-245.

2. Nested Ninja division could produce -j0, meaning unlimited jobs

• Failure: the copied C++ extension divided the global job count by PREBUILD_THREAD_NUM and truncated the result. One job divided across five parents became zero; Ninja interprets -j0 as unlimited jobs.
• Fix: nested division is integer-safe and never returns below one at aiter_worker_limits.py:213-216, covered at tests/test_worker_awareness.py:219-226.

3. Automatic memory sizing ignored container limits

• Failure: host MemAvailable alone is unsafe in a memory-limited container. An 8 GB container on a host with 256 GB free could size from 256 GB and be OOM-killed.
• Fix: host availability comes from psutil.virtual_memory().available; cgroup v2/v1 discovery and finite ancestor limit/usage handling are at aiter_worker_limits.py:33-178. v2, v1, ancestor, and 256 GB-host/8 GB-container regressions are covered at tests/test_worker_awareness.py:53-149.

4. The memory-per-worker estimate materially under-sized compiler RSS

• Failure: setup and JIT assumed 0.5 GB per job, while the cold MI350X run observed approximately 1.43 GB normalized peak RSS per worker. Other direct AOT paths had no shared memory budget.
• Fix: the policy uses a readable, rounded 1.5 GB estimate and selects the lower CPU/memory budget at aiter_worker_limits.py:6-8 and aiter_worker_limits.py:181-210. Memory capping is covered at tests/test_worker_awareness.py:194-201.

5. CPU sizing could use CPUs unavailable to the process

• Failure: several paths used os.cpu_count(), which may report host CPUs outside the process affinity/cpuset. Other call sites used unrelated static defaults, allowing CPU oversubscription.
• Fix: CPU detection prefers os.process_cpu_count(), falls back to affinity, and never returns fewer than one at aiter_worker_limits.py:14-30, covered at tests/test_worker_awareness.py:23-39.

6. Generic MAX_JOBS collided with parent frameworks

• Failure: AITER read and overwrote generic MAX_JOBS, allowing vLLM, SGLang, PyTorch, or another parent framework to accidentally control and potentially oversubscribe AITER fanout, while AITER could mutate the framework's own setting.
• Fix: only AITER_MAX_JOBS controls AITER. Namespace isolation is covered at tests/test_worker_awareness.py:167-173.

7. MLA decode AOT called the compile API with the wrong ABI

• Failure: aiter/aot/asm_mla_decode_fwd.py passed an HSACO path and integer item sizes to a compile function that expects gqa_ratio, q_dtype, and kv_dtype. Both sides diverged when introduced on May 29, 2025 in commit 01864fa8e2347421bc5c314de8b584777ea991ec.
• Fix: the driver now matches the real compile signature at aiter/aot/asm_mla_decode_fwd.py:7-28, with signature and invocation coverage at tests/test_asm_mla_decode_aot.py:16-45.

8. AOT drivers passed strings as ProcessPoolExecutor.max_workers

• Failure: MLA, PA, and PA-ragged passed os.environ.get("MAX_JOBS", "16") directly to ProcessPoolExecutor. The default is the string "16"; Python rejects it with TypeError: '<=' not supported between instances of 'str' and 'int'. PA-v1 had the same failure whenever MAX_JOBS was set. The merge-base code is visible at aiter/aot/asm_mla_decode_fwd.py:46-49, aiter/aot/pa.py:97-100, and aiter/aot/pa_v1.py:104-107.
• Fix: every pool now receives an integer from get_worker_count_for(...); MLA verifies the exact integer constructor call at tests/test_asm_mla_decode_aot.py:48-70.

9. AOT worker exceptions were silently discarded

• Failure: MLA, PA, PA-ragged, PA-v1, and sampling called executor.map(...) without consuming its iterator. Worker compilation failures could therefore remain unraised in the parent.
• Fix: all result iterators are consumed. Representative error propagation is tested at tests/test_asm_mla_decode_aot.py:72-94; sampling submits all families before consuming them at aiter/aot/sampling.py:82-106.

Improvements and hardening in the final diff (7)

1. All CPU compilation paths now use the single aiter_worker_limits.py policy instead of arbitrary per-caller defaults.
2. Automatic CPU concurrency leaves operational headroom by using at most 80% of process-available logical CPUs.
3. Config-driven pools are capped to submitted work at aiter_worker_limits.py:219-221.
4. Sampling preserves cross-family compilation concurrency while still surfacing errors, covered at tests/test_sampling_aot.py:9-42.
5. mp_tuner remains GPU-count-driven; only compiler descendants are constrained. GPU-count behavior is covered in op_tests/tuning_tests/test_mp_tuner_logic.py.
6. The shared helper is included in wheels and standalone source entry points resolve it without requiring a prior package install, covered by tests/test_worker_entrypoints.py.
7. Explicit AITER_MAX_JOBS is documented and tested as an unsafe expert override that deliberately bypasses automatic CPU/memory caps at aiter_worker_limits.py:189-210 and tests/test_worker_awareness.py:158-165.

Final policy

• AITER_MAX_JOBS is the sole AITER top-level override; generic MAX_JOBS is ignored and preserved. An explicit value bypasses automatic CPU and memory caps and can oversubscribe or OOM the process/container.
• Automatic workers are min(floor(process_available_cpus * 0.80), min(psutil_host_available, finite_cgroup_remaining) / 1_500_000_000), with a floor of one. When no finite cgroup memory limit exists, host availability is used.
• Config-driven pools are additionally capped to their work count.
• Process-pool children always force nested AITER, Ninja, CMake, Make, OpenMP, OpenBLAS, MKL, and NumExpr fanout to one.
• mp_tuner remains GPU-count-driven.

The editable installation and policy are documented at README.md:105-149, docs/installation.rst:80-89, and aiter/aot/flydsl/README.md:60-68.

Telemetry basis

The intentionally untracked cold MI350X diagnostic observed 23 live workers, 32.8 GB aggregate RSS, and approximately 1.43 GB normalized peak RSS per worker. The final fixed estimate is rounded up to a readable decimal 1.5 GB. On the current 24-logical-CPU container, no finite cgroup memory limit is configured; the final policy reports CPU 19 and memory 180 budgets and selects 19 workers.

Validation

• 47 focused tests and 7 subtests pass.
• All modified Python files pass Ruff 0.16.0 and byte-compilation.
• Modified GitHub workflows parse as YAML; git diff --check passes.
• A clean wheel builds and installs into a fresh Python 3.14 virtual environment. The installed wheel exposes only the CPU/memory policy, forces nested fanout to one, honors AITER_MAX_JOBS, and preserves inherited generic MAX_JOBS.
• Repository-wide Ruff reports only two pre-existing executable-bit warnings in untouched files: aiter/aot/flydsl/grouped_moe.py and op_tests/test_flydsl_grouped_gemm_gfx1250.py.
• No remaining review finding affects the PR diff.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5164 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.