Skip to content

ROCm: keep Inductor's compile pool up, so a retired pool can't hang the next one's workers - #1101

Draft
merceod wants to merge 4 commits into
users/atindra/pr-22-t4-fusionfrom
users/atindra/pr-35-t8-firststart
Draft

merceod wants to merge 4 commits into
users/atindra/pr-22-t4-fusionfrom
users/atindra/pr-35-t8-firststart

Conversation

@merceod

@merceod merceod commented Nov 1, 2026 •

Copy link
Copy Markdown
Collaborator

Stacked on #1088 (pr-22-t4-fusion). Picks the commits of #1100 first; they drop out of this diff as those merge.

Why

A Qwen3.5 9B server's first start from a fresh compile-cache seed hung in 2 of 27 tries on the week-3 text tree and 0 of 26 on the base (5 of 44 against 0 of 43 in a second set), the engine waiting in async_compile.get_result with the GPU idle. It's a fork race in torch 2.12.1's Inductor compile sidecar. The sidecar retires a pool that returned nothing for 60 s with shutdown(wait=False), and the old manager thread holds shutdown_lock until its workers exit. A worker of the next pool forked in that window copies the lock held, and at its first full GC weakref_cb waits on it forever, with the compile that gave it a job. No mstar frame is involved. The tree only exposes it: #1102's MI325X decode configs move 13 graphs out of a seeded cache, so a seeded start sends the pool 22 jobs instead of 1.

What changed

Stacked on #1088 (#1087's engine.py is under it) with #1100 picked first, since the RocmPlatform hunk sits next to #1100's env default.

  • Platform.compile_pool_overrides(): torch._inductor.config values in force while the engine builds Inductor's compile pool, {} on the base class. ROCm returns {"quiesce_async_compile_pool": False}.
  • mstar/platforms/compile.py: start_compile_pool() builds the pool under inductor_config.patch(...). The pool reads the value once when it's built and keeps it after the patch, so the setting never reaches a graph's cache key and a golden compile cache stays valid. docs/architecture.rst says so.
  • The second commit calls it in EngineManager.build before any submodule loads, not first thing in Engine.warmup: a model that compiles at load (FLUX.2 klein, Z-Image, lingbot_world) had built the pool under Inductor's defaults.
  • Cost: 8 idle workers stay up, about 795 MiB PSS plus a 100 MiB sidecar per engine process; outputs don't change.
  • CUDA: the base platform returns {}, so start_compile_pool() returns before touching Inductor. This matters for main's CUDA path only if that platform sets the override, and it does not. The same torch code runs there, so the race probably exists on CUDA too (not tested).

Verification

One MI325X per job on a shared node, so start times are direction only. First starts from a fresh seed, Qwen3.5 9B:

arm starts not ready ready s min / median / max graphs recompiled cache keys kept
tree without the fix 19 2 202 / 217 / 253 13 -
this fix 21 0 196 / 223 / 241 13 yes
this fix, race forced 3 0 208 / 223 / 247 13 yes
TORCHINDUCTOR_COMPILE_THREADS=1 10 0 163 / 197 / 226 13 yes
TORCHINDUCTOR_WORKER_START=spawn 10 0 184 / 236 / 301 13 yes
TORCHINDUCTOR_QUIESCE_ASYNC_COMPILE_POOL=0 8 0 265 / 302 / 376 167 no
  • 0 of 21 alone would happen 7.9 % of the time at the 5-of-44 rate, so the evidence is structural: with the fix every traced start (9 of 9) has one pool and no quiesce, and the forced race can't start. Forcing it without the fix (the old lock held 3 s longer) hangs 2 of 2 servers, and leaves a pool job unfinished on the base 2 of 2. A CPU repro on torch's SubprocPool alone hangs 13 of 13 at fork gaps of 0-2 ms. A CPU test asserts pool.timer is None.
  • Correctness, text-chat-256: 9B PASS (18/32 exact, 14 near ties, 0 unexplained), 0.8B PASS (15/32, 17, 0). Clean decode 9B 4.38 ms/token, 0.8B 1.07, level with the tree without the fix.
  • CPU Core, one process per file: 10732 tests, 34 failing, the same set as the tree under it; test_platforms.py 63 passed, and the new test fails with no override and with the env var. ruff clean.
  • The second commit, on the week-3 tip: a klein 4B fp8 server's first request at an uncaptured size hung on 2 of 10 starts with the call in warmup (every start retired its first pool at +73-75 s), 0 of 30 with it in build (4B fp8, 9B, 9B mtp2; no start retired a pool); ready medians level (4B 78 vs 81 s, 9B 179 vs 178); correctness PASS; CPU Core failing set equal to the tip's; test/rocm on GPU 1965 passed. These counts come from those trees, not this branch.

Not in this PR

  • The env var is the obvious fix and the wrong one: quiesce_async_compile_pool is part of every graph's cache key, so it recompiles 167 graphs on every first start of every ROCm model. A job without this fix can set TORCHINDUCTOR_COMPILE_THREADS=1 instead.
  • Upstream, shutdown(wait=True) in SubprocMain._quiesce is the one-line fix (0 of 5 hung in the CPU repro). Not reported to torch yet; Python's note on the hazard is gh-90622.
  • Not run: test/rocm on this commit alone, TP, NVIDIA. The 0.9 GiB per engine process is unmeasured on Omni's several engines.

amd-flashinfer 0.6.18+amd.1's generated AITER prefill config names __hip_bfloat16 before anything declares it, because the JIT's -DUSE_ROCM=1 turns off the hip_bf16.h include in torch's BFloat16.h. Every bf16 model on FlashInfer's auto backend dies at setup with a ninja log. Give the JIT the header through FLASHINFER_EXTRA_CUDAFLAGS on ROCm; 0.6.18+amd.2 includes it itself.
… pool's workers

The pool is built under a config patch at the top of warmup, since the setting is part of a graph's cache key and an env default would recompile everything a golden cache holds.
A submodule can compile at load, and klein does, so the pool built at the top of warmup came too late: the first compile had already built it under inductor's defaults, quiesce included. The call moves to the top of EngineManager.build, so no model has to ask for it.

@vasilevklart vasilevklart left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. The mechanism matches torch 2.12.1 subproc_pool.py (SubprocMain._quiesce does self.pool.shutdown(wait=False)),

# without waiting for it. A worker of the next pool forked meanwhile
# inherits the old pool's shutdown lock held by a thread it doesn't
# have, and hangs at its first full GC with the first job it took.
return {"quiesce_async_compile_pool": False}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The CPU repro hangs on SubprocPool alone, so CUDA should hit this on any start that retires a pool too, and the base class keeps the quiesce there

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants