fix(aiter): rebuild the shim link line when the AITER library moves

[demandal25]

Summary

An AITER-backed module cached under .../gfx950/ can keep loading a gfx942 AITER library and segfault — and neither setting FLASHINFER_ROCM_ARCH_LIST correctly nor unsetting it recovers. Only deleting ~/.cache/flashinfer does.

Observed on an MI350X, and this is the mechanism behind a crash that first looked like a broken AITER kernel:

build.ninja  12:51:07   written with aiter_libs/gfx942__aiter-0.1.10
.so          12:53:28   relinked from that stale ninja, RUNPATH -> gfx942
test_activation_aiter_hip.py::test_silu_and_mul_aiter_vs_ref  ->  SIGSEGV

The second run had FLASHINFER_ROCM_ARCH_LIST=gfx950 set and did build a correct gfx950 library alongside — it was simply never linked.

What changed

• flashinfer/jit/aiter_source.py — new refresh_aiter_jitspec(spec): calls write_ninja() so a changed AITER library path takes effect.
• flashinfer/jit/{activation,norm,rope}.py — the three AITER-linked generators route their spec through it.

Architecture / design notes

The AITER shim libs live outside the JIT tree, under aiter_libs/<arch>__aiter-<version>/, and reach the module only as an -L/-rpath on the link line. JitSpec.build() writes build.ninja only when it is missing (core.py:323), so once a module has been built the recorded link line is never revisited, and the resulting .so carries a RUNPATH into the original directory.

CLAUDE.md already documents this staleness as a developer gotcha — "changing env vars is a silent no-op". What is new is that on this path it is not a no-op but a segfault, reachable by following our own documented workaround for a wrong-architecture build.

write_ninja funnels through write_if_different, so calling it unconditionally costs nothing when the link line is unchanged and rewrites exactly when it moves; ninja then relinks on its own.

Two refinements in the second commit. The refresh is scoped to the JIT path: an AOT-prebuilt module loads straight from aot_path, and a FLASHINFER_DISABLE_JIT run raises before ninja is consulted, so in both cases the manifest has no reader and rewriting it would be pure filesystem noise. And the write now takes spec.lock_path — the same lock JitSpec.build() holds while ninja runs — because write_if_different truncates in place: without it, a concurrent builder (pytest -n auto shares one JIT cache across processes) could have the manifest emptied under it mid-read.

Changing the spec name to include the AITER tag was considered and rejected: the name becomes TORCH_EXTENSION_NAME, and an AITER version such as 0.1.10 is not a valid C++ identifier.

Worth stating plainly, because the failure is misleading: AITER's silu_and_mul is fine on gfx950 — 33/33 pass on a clean cache. Without this fix the obvious reading of the crash is that CDNA4 support is broken, which it is not.

Relationship to #281

They compound and are best landed together. #281 stops the wrong-architecture build at the source; this PR is what lets an already-poisoned cache recover. A user upgrading to #281 alone would still segfault until they deleted ~/.cache/flashinfer.

CDNA3 impact

None. On a single-architecture host the resolved tag never changes, write_if_different makes the extra call a no-op, and no link line is rewritten. Behaviour differs only when the AITER library path changes between runs, which on a gfx942-only machine it does not.

Test plan

Verified on MI350X (gfx950), ROCm 7.2.0, torch 2.9.1, amd-aiter 0.1.10. A/B against one poisoned cache, FLASHINFER_ROCM_ARCH_LIST=gfx950 in both arms:

Arm	Result
control — unmodified amd-integration	exit 139 (SIGSEGV)
treatment — this branch	33 passed

• build.ninja rpath: gfx942__aiter-0.1.10 → gfx950__aiter-0.1.10
• .so RUNPATH: gfx942__aiter-0.1.10 → gfx950__aiter-0.1.10
• Not yet run on gfx942 hardware; a CDNA3 run is scripted and planned before merge.
• pre-commit run -a

[Copilot]

Pull request overview

This PR addresses a ROCm/AITER JIT cache failure mode where a previously generated build.ninja can retain stale AITER link flags (notably -L/-rpath), causing a module cached under one arch directory (e.g. gfx950/) to continue loading a different-arch AITER library (e.g. gfx942/) and potentially segfault. The fix introduces an explicit “refresh” step that regenerates the ninja file so link lines track the resolved AITER library location across runs.

Changes:

• Add refresh_aiter_jitspec(spec) to force regeneration of build.ninja for AITER-backed specs.
• Route the AITER generators in activation, norm, and RoPE through refresh_aiter_jitspec(...) so link-line changes take effect without manual cache deletion.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
flashinfer/jit/aiter_source.py	Adds refresh_aiter_jitspec() to rewrite build.ninja when AITER lib paths change.
flashinfer/jit/activation.py	Wraps AITER activation JIT spec creation with the refresh helper.
flashinfer/jit/norm.py	Wraps AITER RMSNorm JIT spec creation with the refresh helper.
flashinfer/jit/rope.py	Wraps AITER RoPE JIT spec creation with the refresh helper.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

flashinfer/jit/aiter_source.py:185

• refresh_aiter_jitspec is introduced without type annotations, while the rest of this module consistently annotates public helpers (e.g. resolve_aiter_build_arch() -> str). Adding a JitSpec input/return annotation improves readability and static checking for the new helper that's now part of the JIT-generator API surface.

def refresh_aiter_jitspec(spec):

[demandal25]

Suppressed comment from review 4964930619 — read and addressed.

aiter_source.py:185 — refresh_aiter_jitspec missing type annotations. Accepted; fixed in fd60305. It was the only unannotated def in the module. Now def refresh_aiter_jitspec(spec: JitSpec) -> JitSpec:, with JitSpec imported from .core alongside the existing logger import — core does not import aiter_source, so this adds no cycle. mypy passes under pre-commit run -a.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

flashinfer/jit/aiter_source.py:224

• refresh_aiter_jitspec() introduces new behavior that is critical for correctness (rewriting build.ninja to pick up a moved AITER -L/-rpath), but there is currently no unit test coverage for it. The existing ROCm tests cover aiter_source.resolve_aiter_build_arch()/_aiter_libs_dir() (e.g. tests/rocm_tests/test_aiter_build_arch_hip.py), but nothing exercises the new refresh path (short-circuit for spec.is_aot/FLASHINFER_DISABLE_JIT, and calling spec.write_ninja() under the per-spec lock). Please add a small GPU-free test that monkeypatches a dummy spec’s write_ninja and lock_path to assert the expected call/skip behavior.

    if spec.is_aot or os.environ.get("FLASHINFER_DISABLE_JIT"):
        return spec
    with FileLock(spec.lock_path, thread_local=False):
        spec.write_ninja()
    return spec

[demandal25]

Suppressed comment from review 4965876749 — read and addressed.

aiter_source.py:224 — no test coverage for refresh_aiter_jitspec(). Accepted; added in 37db9f7 as tests/rocm_tests/test_aiter_jitspec_refresh_hip.py, following the GPU-free pattern of test_aiter_build_arch_hip.py as suggested. A StubSpec records write_ninja() calls and carries a real lock file, so nothing compiles or touches a device — 4 tests, ~1s.

Covers the JIT path (writes), both short-circuits (spec.is_aot, FLASHINFER_DISABLE_JIT), and that the write blocks while a concurrent holder owns spec.lock_path.

Checked as a real regression test rather than a vacuous one — reverting refresh_aiter_jitspec() to its pre-fix body fails three of the four:

FAILED test_refresh_skips_aot_prebuilt        - assert 1 == 0
FAILED test_refresh_skips_when_jit_disabled   - assert 1 == 0
FAILED test_refresh_waits_for_the_build_lock  - waited 0.000s

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.