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

flashinfer/jit/activation.py

@@ -153,15 +153,17 @@ def gen_act_and_mul_module(act_func_name: str) -> JitSpec:
 
 
 def gen_silu_and_mul_aiter_module() -> JitSpec:
-    from .aiter_source import aiter_jitspec_flags
+    from .aiter_source import aiter_jitspec_flags, refresh_aiter_jitspec
 
     extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_activation")
-    return gen_jit_spec(
-        "silu_and_mul_aiter",
-        [
-            jit_env.FLASHINFER_CSRC_DIR / "activation_aiter.cu",
-            jit_env.FLASHINFER_CSRC_DIR / "activation_aiter_jit_pybind.cu",
-        ],
-        extra_include_paths=extra_include_paths,
-        extra_ldflags=extra_ldflags,
+    return refresh_aiter_jitspec(
+        gen_jit_spec(
+            "silu_and_mul_aiter",
+            [
+                jit_env.FLASHINFER_CSRC_DIR / "activation_aiter.cu",
+                jit_env.FLASHINFER_CSRC_DIR / "activation_aiter_jit_pybind.cu",
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_ldflags=extra_ldflags,
+        )
     )

flashinfer/jit/aiter_source.py

@@ -25,9 +25,11 @@
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 
+from filelock import FileLock
+
 from ..arch_caps import normalize_arch
 from . import env as jit_env
-from .core import logger
+from .core import JitSpec, logger
 
 
 _DEFAULT_BUILD_ARCH = "gfx942"
@@ -180,6 +182,48 @@ def _aiter_libs_dir() -> Path:
     return d
 
 
+def refresh_aiter_jitspec(spec: JitSpec) -> JitSpec:
+    """Regenerate ``build.ninja`` so a changed AITER library path takes effect.
+
+    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, so once a module has been built the recorded link
+    line is never revisited -- the module keeps loading whichever AITER lib it
+    was first built against, even after the resolved architecture changes.
+
+    That is not a stale-build annoyance: the ``.so`` retains a RUNPATH into the
+    old directory, so a module cached under ``.../gfx950/`` can go on loading a
+    gfx942 library and **segfault**, and clearing ``FLASHINFER_ROCM_ARCH_LIST``
+    or setting it correctly does not fix it. Only deleting the cache does.
+
+    ``write_ninja`` funnels through ``write_if_different``, so this is free when
+    nothing changed and rewrites exactly when the link line moves; ninja then
+    relinks on its own.
+
+    Only the JIT path is touched. An AOT-prebuilt module is loaded 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.
+
+    The write 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.
+
+    Args:
+        spec: The freshly created :class:`~flashinfer.jit.core.JitSpec`.
+
+    Returns:
+        The same spec, for use as ``return refresh_aiter_jitspec(gen_jit_spec(...))``.
+    """
+    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
+
+
 @functools.lru_cache(maxsize=1)
 def _aiter_csrc_include_dir() -> Path:
     """The aiter_meta C++ public header dir (rmsnorm.h / activation.h / rope.h)."""

flashinfer/jit/norm.py

@@ -34,15 +34,17 @@ def gen_norm_module() -> JitSpec:
 
 
 def gen_norm_aiter_module() -> JitSpec:
-    from .aiter_source import aiter_jitspec_flags
+    from .aiter_source import aiter_jitspec_flags, refresh_aiter_jitspec
 
     extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_rmsnorm")
-    return gen_jit_spec(
-        "norm_aiter",
-        [
-            jit_env.FLASHINFER_CSRC_DIR / "norm_aiter.cu",
-            jit_env.FLASHINFER_CSRC_DIR / "norm_aiter_jit_pybind.cu",
-        ],
-        extra_include_paths=extra_include_paths,
-        extra_ldflags=extra_ldflags,
+    return refresh_aiter_jitspec(
+        gen_jit_spec(
+            "norm_aiter",
+            [
+                jit_env.FLASHINFER_CSRC_DIR / "norm_aiter.cu",
+                jit_env.FLASHINFER_CSRC_DIR / "norm_aiter_jit_pybind.cu",
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_ldflags=extra_ldflags,
+        )
     )

flashinfer/jit/rope.py

@@ -29,15 +29,17 @@ def gen_rope_module() -> JitSpec:
 
 
 def gen_rope_aiter_module() -> JitSpec:
-    from .aiter_source import aiter_jitspec_flags
+    from .aiter_source import aiter_jitspec_flags, refresh_aiter_jitspec
 
     extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_rope_pos_fwd")
-    return gen_jit_spec(
-        "rope_aiter",
-        [
-            jit_env.FLASHINFER_CSRC_DIR / "rope_aiter.cu",
-            jit_env.FLASHINFER_CSRC_DIR / "rope_aiter_jit_pybind.cu",
-        ],
-        extra_include_paths=extra_include_paths,
-        extra_ldflags=extra_ldflags,
+    return refresh_aiter_jitspec(
+        gen_jit_spec(
+            "rope_aiter",
+            [
+                jit_env.FLASHINFER_CSRC_DIR / "rope_aiter.cu",
+                jit_env.FLASHINFER_CSRC_DIR / "rope_aiter_jit_pybind.cu",
+            ],
+            extra_include_paths=extra_include_paths,
+            extra_ldflags=extra_ldflags,
+        )
     )

tests/rocm_tests/test_aiter_jitspec_refresh_hip.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc.
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Tests for ``refresh_aiter_jitspec``, the AITER link-line refresh.
+
+GPU-free by construction: a stub spec stands in for :class:`JitSpec`, so nothing
+here compiles or touches a device.
+
+What is being protected. The AITER shim libs live outside the JIT tree and reach
+a module only as an ``-L``/``-rpath`` on the link line. ``JitSpec.build()`` writes
+``build.ninja`` only when it is missing, so without this refresh a module cached
+under ``.../gfx950/`` keeps loading whichever AITER lib it was first built
+against -- a wrong-architecture load that segfaults rather than failing cleanly.
+"""
+
+import threading
+import time
+
+import pytest
+from filelock import FileLock
+
+from flashinfer.jit.aiter_source import refresh_aiter_jitspec
+
+
+class StubSpec:
+    """Minimal stand-in for JitSpec: records write_ninja(), real lock on disk."""
+
+    def __init__(self, lock_path, is_aot=False):
+        self.name = "stub_aiter"
+        self.lock_path = lock_path
+        self.is_aot = is_aot
+        self.writes = 0
+
+    def write_ninja(self):
+        self.writes += 1
+
+
+@pytest.fixture
+def lock_path(tmp_path):
+    return tmp_path / "stub_aiter.lock"
+
+
+def test_refresh_rewrites_on_the_jit_path(lock_path):
+    spec = StubSpec(lock_path)
+    assert refresh_aiter_jitspec(spec) is spec
+    assert spec.writes == 1
+
+
+def test_refresh_skips_aot_prebuilt(lock_path):
+    """build_and_load() loads straight from aot_path, so ninja never reads it."""
+    spec = StubSpec(lock_path, is_aot=True)
+    refresh_aiter_jitspec(spec)
+    assert spec.writes == 0
+
+
+def test_refresh_skips_when_jit_disabled(lock_path, monkeypatch):
+    """build() raises before consulting the manifest, so writing it is noise."""
+    monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1")
+    spec = StubSpec(lock_path)
+    refresh_aiter_jitspec(spec)
+    assert spec.writes == 0
+
+
+def test_refresh_waits_for_the_build_lock(lock_path):
+    """The write must not land while a concurrent builder holds the lock.
+
+    write_if_different truncates in place, so an unlocked rewrite could empty
+    build.ninja while another process's ninja is mid-read. Reachable in practice:
+    pytest -n auto shares one JIT cache across processes.
+    """
+    hold_for = 1.0
+    acquired = threading.Event()
+
+    def holder():
+        with FileLock(lock_path, thread_local=False):
+            acquired.set()
+            time.sleep(hold_for)
+
+    t = threading.Thread(target=holder)
+    t.start()
+    try:
+        assert acquired.wait(timeout=10), "holder thread never took the lock"
+        spec = StubSpec(lock_path)
+        start = time.monotonic()
+        refresh_aiter_jitspec(spec)
+        waited = time.monotonic() - start
+    finally:
+        t.join()
+
+    assert spec.writes == 1
+    # Generous lower bound: proves it blocked, without being timing-flaky.
+    assert waited > hold_for / 2, (
+        f"refresh did not wait on the lock (waited {waited:.3f}s)"
+    )