From b352dceb3035aa2a81ed1c51518e8f15cd6da895 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 11:10:31 -0400 Subject: [PATCH 01/41] docs(skill): pick the backend by measurement, not by reputation The benchmark skill told you how to time a kernel but not which candidates to time, so "AMD ships an optimized library for this" has been read as a reason to route to it. The record here says that is a coin flip: AITER prefill beats fa2 by 3.67x on gfx942 and 4.51x on gfx950, while the in-tree HIP kernel beats AITER on rmsnorm/fused_add_rmsnorm by 1.6-1.8x and on append_paged_kv_cache at 3.62 TB/s against 2.86. Adds a section saying to enumerate the ROCm libraries that already implement an op -- AITER, CK/CK-Tile, hipBLASLt, rocBLAS, MIOpen, Triton -- benchmark each against the in-tree kernel, and route auto to the winner per arch, recording it as a Capability row with an evidence string. Also names the three things that make the measurement wrong more often than the stopwatch does: a library can win on time and still fail the contract (CK-Tile layernorm2d reads fp32 gamma/beta as the input dtype), a library may dispatch internally so "AITER" is not one number, and gfx942 and gfx950 have diverged. Also records a measurement the paged-prefill bullet needed. It already warned that the advertised native page sizes are a hint and that plan() degrades to flat-gather. On amd-aiter 0.1.21.post2 all three advertised sizes fail the probe with "no matching kernel found", on both gfx942 and gfx950: advertised native page sizes: [128, 256, 1024] page_size= 128 available=False page_size= 256 available=False page_size= 1024 available=False So mha_batch_prefill is currently unreachable and the paged path should be benchmarked as flat-gather. Co-Authored-By: Claude --- .claude/skills/benchmark-kernel/SKILL.md | 26 +++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/.claude/skills/benchmark-kernel/SKILL.md b/.claude/skills/benchmark-kernel/SKILL.md index 16de4710685..55c10a5ce4f 100644 --- a/.claude/skills/benchmark-kernel/SKILL.md +++ b/.claude/skills/benchmark-kernel/SKILL.md @@ -25,7 +25,7 @@ For the in-repo profiler wrapper, see [`profiler/rocm/rocm_profiler.py`](../../. - Explicit `backend="aiter"` + `kv_layout != "NHD"` → `ValueError`. Grep [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for `only supports kv_layout`; single prefill raises at the call, both batch wrappers at `plan()`. Not raised by auto-selection — that path silently falls back to `fa2`. - Explicit `backend="aiter"` on non-gfx942/gfx950 → `RuntimeError`. - `amd-aiter` not importable → `ImportError`. - - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py). The `{16, 1024}` arm there is unreachable now that `flashinfer.rocm.aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.20. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. + - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py). The `{16, 1024}` arm there is unreachable now that `flashinfer.rocm.aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.20. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. **Measured on amd-aiter 0.1.21.post2, all three advertised sizes fail the probe** (`no matching kernel found`) on gfx942 *and* gfx950, so every paged prefill flat-gathers and `mha_batch_prefill` is unreachable — benchmark the paged path as flat-gather until that is fixed. - Auto-selection (no explicit `backend=`) silently falls back to `fa2` for any of: `kv_layout != "NHD"`, custom mask, dtype not in `{fp16, bf16}`, `dtype_q != dtype_kv`, `head_dim_qk != head_dim_vo`, `pos_encoding_mode != "NONE"`, or `amd-aiter` not importable. See `_auto_select_prefill_backend()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for the authoritative list; it returns `(backend, reason)`, and the reason names the constraint that forced the fallback. - **Always verify numerical parity before trusting perf numbers.** Compare default-HIP vs AITER outputs with `torch.testing.assert_close(rtol=1e-2, atol=1e-2)` for BF16/FP16 first. - **`gcnArchName` is the unambiguous arch marker.** Device strings show `cuda:0` on AMD too. Record `torch.cuda.get_device_properties(0).gcnArchName` and `torch.version.hip` alongside every number — a `gfx942` / ROCm 7.2 result is not comparable to a `gfx950` / ROCm 7.0.2 result. @@ -83,6 +83,30 @@ Output (under `benchmarks/rocm/`, gitignored): - **Hang or no output:** confirm `which rocprofv3` is on `PATH`; the wrapper uses script `print()` output as a heartbeat — make sure the `if __name__ == "__main__":` block prints something. - **Use `--timing-only` first** to verify the kernel path works before involving `rocprofv3`. +## Before tuning: check whether a library already wins + +For any op, enumerate the ROCm libraries that already implement it — AITER, CK / +CK-Tile, hipBLASLt, rocBLAS, MIOpen, Triton — and benchmark each against the +in-tree `hip`/`fa2`/`native` kernel before writing or tuning one. Route `auto` to +whichever wins, on evidence, per arch. + +Both outcomes are on the record here, so assume neither: AITER prefill beats fa2 +by 3.67x (gfx942) and 4.51x (gfx950), while the in-tree HIP kernel beats AITER on +`rmsnorm`/`fused_add_rmsnorm` (1.6-1.8x) and on `append_paged_kv_cache`. + +Three things that make this measurement wrong more often than the stopwatch: + +- **Contract before speed.** A library can be faster and still unusable. CK-Tile's + `layernorm2d` reads this API's fp32 `gamma`/`beta` as the input dtype and returns + garbage with no error — it was rejected on contract, not on time. +- **Name the arm.** A library may dispatch internally between its own backends, so + "AITER" is not one number. Say which module and codegen path ran. +- **Both boards.** gfx942 and gfx950 have diverged; a win on one is not a win. + +Record the outcome as a `Capability` row in `flashinfer/rocm/arch_caps.py` with an +`evidence=` string naming board, versions and date, so the next person inherits the +measurement instead of repeating it. + ## External tuning references When optimizing a CDNA kernel, consult these in order: From 01c7b73fab62ea14af98d3e258ca17539ff53fdb Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 18:15:38 -0400 Subject: [PATCH 02/41] fix(rocm): key the AITER lib cache by ROCm as well as version MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `aiter_libs` cache tag is `__aiter-`. ROCm was never a component because the pinned wheel's version string carried it: amd-aiter 0.1.20+rocm10.1.0a20260819.3135022 A source-built AITER has no local version segment — `pip show amd-aiter` on a `v0.1.21.post2` clone reports a bare `0.1.21.post2` — so that accidental key disappears and two ROCm toolchains on one box share a cache directory. These are CK-tile objects linked into the shims by `-L`/`-rpath`, and a mismatched one loads rather than failing, so the failure mode is a wrong kernel, not an error. `_rocm_version()` already existed in aiter_variants for the #363 store tag; move it down to aiter_source, which aiter_variants already imports from, and use it in both. Its docstring claimed the aiter_libs cache "never leaves the box that built it" as the reason ROCm was omitted there — true, and beside the point, since a box's toolchain can change under a persistent cache. The composed tag now gets the path-component guard the arch already had: `torch.version.hip` is not this module's value to trust, and the arch check alone would let a separator through in the ROCm field. Co-Authored-By: Claude --- flashinfer/jit/rocm/aiter_source.py | 30 +++++++++++++--- flashinfer/jit/rocm/aiter_variants.py | 17 +-------- tests/rocm/test_aiter_source_probes.py | 48 +++++++++++++++++++++++++- 3 files changed, 74 insertions(+), 21 deletions(-) diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 12ccbfa8883..d34a390b5c8 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -201,8 +201,18 @@ def resolve_aiter_build_arch() -> str: return _DEFAULT_BUILD_ARCH +def _rocm_version() -> str: + """The ROCm/HIP version these artifacts were compiled against.""" + try: + import torch + + return torch.version.hip or "unknown" + except Exception: + return "unknown" + + def _aiter_cache_tag() -> str: - """A filesystem-safe tag keying the lib cache by target arch and AITER version. + """A filesystem-safe tag keying the lib cache by arch, AITER version and ROCm. Without the arch component, a lib built for one arch would be silently reused on a machine with a different arch. Without the version component, an AITER @@ -210,6 +220,10 @@ def _aiter_cache_tag() -> str: silently reuse the stale .so. The FlashInfer JIT dir already keys by its own version+arch, but this cache sits outside it. + ROCm is a component because a source-built AITER's version is bare + ("0.1.21.post2"), where a wheel's carried "+rocm10.1.0a..." and keyed ROCm by + accident. Without it, two ROCm toolchains share one tag. + Keyed on the *resolved* architecture -- the one actually compiled for -- so the tag cannot disagree with the contents of the directory it names.""" arch = resolve_aiter_build_arch() @@ -231,13 +245,21 @@ def _aiter_cache_tag() -> str: version = _md.version("amd-aiter") except Exception: version = "unknown" - return f"{arch}__aiter-{version}" + tag = f"{arch}__aiter-{version}__rocm-{_rocm_version()}" + # The version strings come from package metadata and torch, neither of which + # this module controls, so the composed tag gets the same check the arch did. + if tag != Path(tag).name or tag.startswith("."): + raise ValueError( + f"refusing to build a cache directory name from {tag!r}: " + f"not a single safe path component" + ) + return tag @functools.lru_cache(maxsize=1) def _aiter_libs_dir() -> Path: - # Keyed by arch + AITER version so a cached lib is never reused across an - # incompatible arch or a changed AITER ABI. + # Keyed by arch + AITER version + ROCm so a cached lib is never reused across + # an incompatible arch, a changed AITER ABI, or a different toolchain. d = jit_env.FLASHINFER_CACHE_DIR / "aiter_libs" / _aiter_cache_tag() d.mkdir(parents=True, exist_ok=True) return d diff --git a/flashinfer/jit/rocm/aiter_variants.py b/flashinfer/jit/rocm/aiter_variants.py index ec77fcd6be2..32b74e1362b 100644 --- a/flashinfer/jit/rocm/aiter_variants.py +++ b/flashinfer/jit/rocm/aiter_variants.py @@ -21,7 +21,7 @@ from typing import List, Optional, Tuple from .. import env as jit_env -from .aiter_source import resolve_aiter_build_arch +from .aiter_source import _rocm_version, resolve_aiter_build_arch __all__ = [ "Family", @@ -201,21 +201,6 @@ def builds() -> Tuple[BuildSpec, ...]: MANIFEST_NAME = "variants_manifest.json" -def _rocm_version() -> str: - """The ROCm/HIP version the artifacts were compiled against. - - Part of the store tag because these are CK-tile objects that ship between - machines in a wheel, unlike the ``aiter_libs`` cache, which never leaves the - box that built it. - """ - try: - import torch - - return torch.version.hip or "unknown" - except Exception: - return "unknown" - - def variant_store_dir(arch: Optional[str] = None) -> Path: """``/aiter_variants/__aiter-__rocm-/``. diff --git a/tests/rocm/test_aiter_source_probes.py b/tests/rocm/test_aiter_source_probes.py index ce803376829..cda4db399d5 100644 --- a/tests/rocm/test_aiter_source_probes.py +++ b/tests/rocm/test_aiter_source_probes.py @@ -53,18 +53,52 @@ def test_the_tag_carries_the_arch_and_the_aiter_version(self, monkeypatch): assert tag.startswith("gfx942__aiter-") + def test_the_tag_carries_the_rocm_version(self, monkeypatch): + """A source-built AITER's version is bare, so ROCm is no longer keyed by + the wheel's "+rocm..." local segment and has to be its own component.""" + monkeypatch.setattr(aiter_source, "resolve_aiter_build_arch", lambda: "gfx942") + monkeypatch.setattr(aiter_source, "_rocm_version", lambda: "7.15.26333") + + assert aiter_source._aiter_cache_tag().endswith("__rocm-7.15.26333") + + def test_two_rocm_toolchains_do_not_share_a_tag(self, monkeypatch): + monkeypatch.setattr(aiter_source, "resolve_aiter_build_arch", lambda: "gfx942") + + monkeypatch.setattr(aiter_source, "_rocm_version", lambda: "7.15.26333") + first = aiter_source._aiter_cache_tag() + monkeypatch.setattr(aiter_source, "_rocm_version", lambda: "6.4.43483") + second = aiter_source._aiter_cache_tag() + + assert first != second + def test_an_unreadable_version_still_produces_a_tag(self, monkeypatch): """A cache key with no version is still better than failing the build.""" import importlib.metadata monkeypatch.setattr(aiter_source, "resolve_aiter_build_arch", lambda: "gfx942") + monkeypatch.setattr(aiter_source, "_rocm_version", lambda: "unknown") monkeypatch.setattr( importlib.metadata, "version", lambda name: (_ for _ in ()).throw(RuntimeError("no metadata")), ) - assert aiter_source._aiter_cache_tag() == "gfx942__aiter-unknown" + assert aiter_source._aiter_cache_tag() == "gfx942__aiter-unknown__rocm-unknown" + + def test_an_absent_torch_still_produces_a_tag(self, monkeypatch): + """`_rocm_version` runs on the import path of a CPU-only smoke test.""" + import builtins + + real_import = builtins.__import__ + + def _no_torch(name, *args, **kwargs): + if name == "torch": + raise ImportError("no torch") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _no_torch) + + assert aiter_source._rocm_version() == "unknown" @pytest.mark.parametrize("arch", ["", ".", "..", "a/b", ".hidden"]) def test_an_unsafe_arch_never_becomes_a_directory_name(self, monkeypatch, arch): @@ -75,6 +109,18 @@ def test_an_unsafe_arch_never_becomes_a_directory_name(self, monkeypatch, arch): with pytest.raises(ValueError, match="single safe path component"): aiter_source._aiter_cache_tag() + @pytest.mark.parametrize("rocm", ["../..", "a/b", "7.15/26333"]) + def test_an_unsafe_rocm_version_never_becomes_a_directory_name( + self, monkeypatch, rocm + ): + """`torch.version.hip` is not this module's to trust; the arch guard + alone would let a separator through in the ROCm component.""" + monkeypatch.setattr(aiter_source, "resolve_aiter_build_arch", lambda: "gfx942") + monkeypatch.setattr(aiter_source, "_rocm_version", lambda: rocm) + + with pytest.raises(ValueError, match="single safe path component"): + aiter_source._aiter_cache_tag() + class TestCsrcIncludeDir: @pytest.fixture(autouse=True) From 5c0e65316a53b4ff52b8e02d175fb240c7c64732 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 18:18:10 -0400 Subject: [PATCH 03/41] fix(rocm): resolve AITER's csrc include dir without requiring aiter_meta MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_aiter_csrc_include_dir` resolved only through the `aiter_meta` package. An editable AITER install has no `aiter_meta` at all: `setup.py develop` takes an `is_develop_mode()` branch that sets `packages = ["aiter"]`, so every shim needing a C++ header — rope, rmsnorm, activation, fused MoE, batch decode — failed with "is the aiter source package installed?" against an install that plainly was. Fall back to AITER's own `AITER_CSRC_DIR`. It derives from `AITER_META_DIR`, which falls back to the repo root when `aiter_meta` has no `csrc/`, so in a develop layout it names the clone's `csrc/` and in any layout it honours an `AITER_META_DIR` override. Confirmed on the 0.1.21.post2 non-editable install that both routes resolve to the same directory: AITER_CSRC_DIR = .../site-packages/aiter_meta/csrc same dir as aiter_meta/csrc/include = True aiter_meta stays first: it needs no `import aiter`, and importing aiter runs arch detection against a live device. The fallback pays that cost only where the alternative was an exception. Co-Authored-By: Claude --- flashinfer/jit/rocm/aiter_source.py | 34 +++++++++++--- tests/rocm/test_aiter_source_probes.py | 62 +++++++++++++++++++++++++- 2 files changed, 89 insertions(+), 7 deletions(-) diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index d34a390b5c8..4773158a148 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -311,17 +311,41 @@ def refresh_aiter_jitspec(spec: JitSpec) -> JitSpec: @functools.lru_cache(maxsize=1) def _aiter_csrc_include_dir() -> Path: - """The aiter_meta C++ public header dir (rmsnorm.h / activation.h / rope.h).""" + """The AITER C++ public header dir (rmsnorm.h / activation.h / rope.h). + + Normally ``aiter_meta/csrc/include``. An editable AITER install has no + ``aiter_meta`` at all -- ``setup.py develop`` takes a branch that sets + ``packages = ["aiter"]`` -- so fall back to AITER's own ``AITER_CSRC_DIR``, + which resolves to the clone's ``csrc/`` in that layout and honours an + ``AITER_META_DIR`` override in any layout. + """ # aiter ships its C++ sources/headers in the sibling aiter_meta package, # which is a namespace package (no __file__) — resolve via __path__. - import aiter_meta + try: + import aiter_meta + + for p in aiter_meta.__path__: + inc = Path(p) / "csrc" / "include" + if inc.exists(): + return inc + except ImportError: + pass + + # Only reached when aiter_meta is missing or incomplete, which is also the + # only case where importing aiter (and needing a live GPU for it) is no + # worse than the RuntimeError below. + try: + from aiter.jit.core import AITER_CSRC_DIR - for p in aiter_meta.__path__: - inc = Path(p) / "csrc" / "include" + inc = Path(AITER_CSRC_DIR) / "include" if inc.exists(): return inc + except Exception: + pass + raise RuntimeError( - "Could not locate aiter_meta/csrc/include; is the aiter source package installed?" + "Could not locate AITER's csrc/include, via either aiter_meta or " + "aiter.jit.core.AITER_CSRC_DIR; is the aiter source package installed?" ) diff --git a/tests/rocm/test_aiter_source_probes.py b/tests/rocm/test_aiter_source_probes.py index cda4db399d5..bcb43758e3a 100644 --- a/tests/rocm/test_aiter_source_probes.py +++ b/tests/rocm/test_aiter_source_probes.py @@ -129,12 +129,70 @@ def _uncached(self): yield aiter_source._aiter_csrc_include_dir.cache_clear() - def test_a_missing_header_tree_names_the_package(self, monkeypatch): + @staticmethod + def _fake_aiter_core(monkeypatch, csrc_dir): + """Put aiter.jit.core in sys.modules so the fallback resolves without a + real `import aiter`, which runs arch detection and needs a device.""" + import sys + import types + + pkg = types.ModuleType("aiter") + pkg.__path__ = [] + jit = types.ModuleType("aiter.jit") + jit.__path__ = [] + core = types.ModuleType("aiter.jit.core") + core.AITER_CSRC_DIR = str(csrc_dir) + for name, mod in ( + ("aiter", pkg), + ("aiter.jit", jit), + ("aiter.jit.core", core), + ): + monkeypatch.setitem(sys.modules, name, mod) + + @staticmethod + def _no_aiter_meta(monkeypatch): + """`setup.py develop` ships no aiter_meta at all, so the import raises.""" + import builtins + + real_import = builtins.__import__ + + def _blocked(name, *args, **kwargs): + if name == "aiter_meta": + raise ImportError("no module named aiter_meta") + return real_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", _blocked) + + def test_aiter_meta_is_preferred_when_present(self): + found = aiter_source._aiter_csrc_include_dir() + assert found.name == "include" and found.parent.name == "csrc" + + def test_an_absent_aiter_meta_falls_back_to_aiter_csrc_dir( + self, monkeypatch, tmp_path + ): + inc = tmp_path / "csrc" / "include" + inc.mkdir(parents=True) + self._no_aiter_meta(monkeypatch) + self._fake_aiter_core(monkeypatch, tmp_path / "csrc") + + assert aiter_source._aiter_csrc_include_dir() == inc + + def test_an_incomplete_aiter_meta_falls_back_too(self, monkeypatch, tmp_path): + """Present but pointing nowhere useful is the same failure as absent.""" import aiter_meta + inc = tmp_path / "csrc" / "include" + inc.mkdir(parents=True) monkeypatch.setattr(aiter_meta, "__path__", ["/nonexistent"]) + self._fake_aiter_core(monkeypatch, tmp_path / "csrc") + + assert aiter_source._aiter_csrc_include_dir() == inc + + def test_neither_route_working_names_both(self, monkeypatch, tmp_path): + self._no_aiter_meta(monkeypatch) + self._fake_aiter_core(monkeypatch, tmp_path / "absent") - with pytest.raises(RuntimeError, match="aiter_meta/csrc/include"): + with pytest.raises(RuntimeError, match="aiter_meta.*AITER_CSRC_DIR"): aiter_source._aiter_csrc_include_dir() From fbaf4674dc144d6c6ab86f0b644b95a6b1d63aec Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 18:19:03 -0400 Subject: [PATCH 04/41] docs(rocm): AITER's build machinery needs no GPU; only this driver does MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both the module docstring and backends.md claimed prebuilding "cannot be a `docker build` step". That conflates the driver with the machinery it calls. `prebuild_aiter_variants` needs a GPU because it produces a variant by calling the op, which launches a kernel. AITER's builder does not: `aiter/setup.py:212` reaches it with sys.path.insert(0, ".../aiter/"); from jit import core a flat import that never runs `aiter/__init__.py` — which is what needs a device, for arch detection and a triton import that fails with "0 active drivers ([]). There should only be one." Measured: a GPU-less container builds `mha_fwd_bf16_nbias_mask_nlse_ndropout_nqscale` this way and `nm -D` finds the expected `aiter::mha_fwd` symbol. The claim matters because it rules out an image-build prebuild that is in fact available, so correct it before anything is designed around it. Co-Authored-By: Claude --- docs/rocm/backends.md | 9 ++++++--- flashinfer/rocm/prebuild_aiter_variants.py | 9 ++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index ef315bb66f2..be2ccf32025 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -272,9 +272,12 @@ probe has to run whether or not the artifact is in the store, and prebuilding that family costs ~32 min on gfx942 for files nothing ever saves time on. Pass `--only mha_batch_prefill` to build them anyway. -**It needs a GPU**: the only supported way to make AITER emit a variant is to -call the op. So it cannot be a `docker build` step — run it as a GPU-attached -job and copy the resulting directory into the image. +**This driver needs a GPU**: it emits a variant by calling the op, which launches +a kernel. So run it as a GPU-attached job and copy the resulting directory into +the image. AITER's build machinery itself needs no GPU — `aiter`'s own +`setup.py` drives it through a flat `sys.path` import of `jit.core` that never +runs `aiter/__init__.py` — so a differently written driver can build variants +inside `docker build`. The store lives under FlashInfer's cache directory, at `aiter_variants/__aiter-__rocm-/`. That cache is diff --git a/flashinfer/rocm/prebuild_aiter_variants.py b/flashinfer/rocm/prebuild_aiter_variants.py index 55d2f6c61d0..843bfc672ef 100644 --- a/flashinfer/rocm/prebuild_aiter_variants.py +++ b/flashinfer/rocm/prebuild_aiter_variants.py @@ -12,9 +12,12 @@ python -m flashinfer.rocm.prebuild_aiter_variants --arch gfx942 python -m flashinfer.rocm.prebuild_aiter_variants --only mha_batch_prefill -**This needs a GPU.** The only supported way to make AITER emit a variant is to -call the op, which launches a kernel -- so it cannot be a ``docker build`` step -and has to run as a GPU-attached job whose output the image copies in. +**This driver needs a GPU**, because it produces variants by calling the op, +which launches a kernel. AITER's build machinery itself does not: ``aiter``'s +own ``setup.py`` reaches it by a flat ``sys.path`` import of ``jit.core`` that +never runs ``aiter/__init__.py``, so a driver written that way can be a +``docker build`` step. This one cannot, and runs as a GPU-attached job whose +output the image copies in. One build at a time, deliberately. ``_aiter_env_scope`` mutates process-global environment that AITER reads, so two concurrent builds in one process have the From b49e9fff1275b9dfe14629f31faf4b4135f5ce99 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 23:15:48 -0400 Subject: [PATCH 05/41] feat(rocm): name and compose every AITER attention variant's CK build MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A GPU-free driver that builds the attention variant .so at image-build time, plus the two parity tests that keep it honest. No Dockerfile change yet, so nothing in this commit runs in CI; commit 5 wires it up. Why a second driver rather than reusing prebuild_aiter_variants: that one emits a variant by *calling the op*, which needs a device. AITER's builder does not — aiter/setup.py:212 reaches it by a flat `sys.path` import of `jit.core` that never runs aiter/__init__.py, the part needing arch detection and triton. Artifacts land in AITER's own jit/ dir rather than FlashInfer's store, which is what makes mha_batch_prefill worth building at all. Its bootstrap doubles as the page-size probe and runs before the loader, so only a hit AITER itself can see spares the compile. Measured on gfx942: arm=warm plan() = 0.2s .so built during plan: 0 arm=cold plan() = 101.4s .so built during plan: 1 Default set is 36 of 40: everything but mha_fwd's four LSE variants, the most expensive builds in the set at 803 s each two-arch. LSE on *single* prefill serves cascade/merge, and Blaze-O1's LUM3 — the consumer that pins this repo as a submodule — records its footprint as "fp16, head_dim 64, batch mode, causal, no soft-cap/bias/LSE". Batch mode is mha_fwd, so LUM3's exact file (mha_fwd_fp16_nbias_mask_nlse_ndropout_nqscale.so) is in the default set and a test asserts it. `--all` builds 40. Verified end to end, GPU_ARCHS="gfx942;gfx950", against the lazily built artifact as control — varlen is the control because its recipe comes from AITER's own helper: lazy : 3123528 bytes 168 symbols driver : 3123528 bytes 168 symbols 99 non-cuid symbols identical; the 69 that differ are all __hip_cuid_* both carry amdgcn-amd-amdhsa--gfx942 and --gfx950 Two traps this cost, both now guarded by tests: - get_args_of_build("mha_fwd") *warns and returns empty args* instead of raising; the real key is "module_mha_fwd". The empty args dropped every include path and surfaced minutes later as "fmha_fwd.hpp file not found". _args_of_build now asserts srcs and includes are non-empty. - compose_mha_fwd_variant_suffix_and_filter is misnamed: it always emits _logits and _skip tokens, which is the *varlen* shape. Its only callers are varlen. mha_fwd composes inline with receipt 100 against varlen's 200, and spells its mask filter "_m*" where batch_prefill spells it "_mask*". Child build output goes to a file, never a pipe: hipcc is verbose enough to fill a pipe buffer, and a child blocked on a full pipe never exits. All six parity assertions A/B'd by breaking the driver one way at a time: caught mha_fwd mask filter -> _mask* caught batch_prefill dtype filter gains a _ caught mha_fwd receipt 100 -> 200 caught batch_prefill loses its sink token caught varlen token renamed caught mha_fwd lse arm no longer skipped Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 464 ++++++++++++++++++++ tests/rocm/test_prebuild_aiter_attention.py | 219 +++++++++ 2 files changed, 683 insertions(+) create mode 100644 docker/prebuild_aiter_attention.py create mode 100644 tests/rocm/test_prebuild_aiter_attention.py diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py new file mode 100644 index 00000000000..2ff22cb7fb8 --- /dev/null +++ b/docker/prebuild_aiter_attention.py @@ -0,0 +1,464 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# SPDX-License-Identifier: Apache-2.0 +"""Build AITER's attention variant ``.so`` at image-build time, with no GPU. + +FlashInfer ``dlopen``s these by composed filename. The wheel prebuilds almost +none of them, so AITER's lazy JIT compiles each one *inside the serving process* +on the first request that needs it -- minutes of stall, possibly hours into the +process's life, under a process-global build lock. + +Runs during ``docker build``, where there is no device. Everything here is +therefore stdlib-only and reaches AITER through the flat ``sys.path`` import its +own ``setup.py`` uses; ``import aiter`` would run arch detection and a triton +import that both need a GPU. The companion GPU-attached driver +(``flashinfer.rocm.prebuild_aiter_variants``) produces variants by *calling the +op*, which is why that one cannot run here. + +Artifacts land in AITER's own ``jit/`` directory, so AITER itself finds them +already built and the bootstraps return without compiling. That is what makes +``mha_batch_prefill`` worth building by this route and not by the store route: +its bootstrap doubles as the page-size probe and runs either way, so only a hit +AITER can see spares the compile. Measured: 101.4 s cold against 0.2 s prebuilt. + + python3 docker/prebuild_aiter_attention.py --list + python3 docker/prebuild_aiter_attention.py --jobs 2 + python3 docker/prebuild_aiter_attention.py --check + +This file deliberately carries its own copy of the variant table: importing +``flashinfer.jit.rocm.aiter_variants`` would pull in torch. +``tests/rocm/test_prebuild_aiter_attention.py`` asserts the two agree, and +asserts the receipts and token spellings against the installed ``aiter/ops/mha.py`` +-- the only detector for an upstream recipe change. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple + +DTYPES = ("bf16", "fp16") + +# Receipts select which CK-tile instance set the codegen emits. They differ per +# family and a wrong one ships a *different kernel* under the right filename, +# with nothing downstream to detect it. +RECEIPT_MHA_FWD = 100 +RECEIPT_VARLEN = 200 +RECEIPT_BATCH_PREFILL = 200 + +GENERATE_PY = "example/ck_tile/01_fmha/generate.py" + + +class Variant(NamedTuple): + """One ``.so``. Mirrors ``VariantKey`` in flashinfer/jit/rocm/aiter_variants.py. + + ``has_alibi`` is absent for the same reason it is absent there: all three C++ + construction sites hard-code it false, so the other arm is unreachable. + """ + + family: str # "mha_fwd" | "mha_varlen_fwd" | "mha_batch_prefill" + dtype: str + has_logits_cap: bool + needs_mask: bool + has_lse: bool + + @property + def md_name(self) -> str: + return _MD_NAME[self.family](self) + + @property + def so_name(self) -> str: + return f"{self.md_name}.so" + + +def _md_mha_fwd(v: Variant) -> str: + """Mirrors aiter/ops/mha.py::mha_fwd. No logits or skip token in this family.""" + return ( + "mha_fwd" + + f"_{v.dtype}" + + "_nbias" + + ("_mask" if v.needs_mask else "_nmask") + + ("_lse" if v.has_lse else "_nlse") + + "_ndropout" + + "_nqscale" + ) + + +def _md_varlen(v: Variant) -> str: + """Mirrors aiter/ops/mha.py::mha_varlen_fwd. Carries logits *and* skip tokens.""" + return ( + "mha_varlen_fwd" + + f"_{v.dtype}" + + ("_logits" if v.has_logits_cap else "_nlogits") + + "_nbias" + + ("_mask" if v.needs_mask else "_nmask") + + ("_lse" if v.has_lse else "_nlse") + + "_ndropout" + + "_nskip" + + "_nqscale" + ) + + +def _md_batch_prefill(v: Variant) -> str: + """Mirrors aiter/ops/mha.py::mha_batch_prefill. Adds the sink axis.""" + return ( + "mha_batch_prefill" + + f"_{v.dtype}" + + ("_logits" if v.has_logits_cap else "_nlogits") + + "_nbias" + + ("_mask" if v.needs_mask else "_nmask") + + ("_lse" if v.has_lse else "_nlse") + + "_ndropout" + + "_nqscale" + + "_nsink" + ) + + +_MD_NAME = { + "mha_fwd": _md_mha_fwd, + "mha_varlen_fwd": _md_varlen, + "mha_batch_prefill": _md_batch_prefill, +} + +# Families that carry a logits-cap axis. mha_fwd has no _logits arm at all -- +# get_aiter_mha_fwd_handle refuses has_logits_cap, routing those calls to varlen. +_HAS_LOGITS_AXIS = { + "mha_fwd": False, + "mha_varlen_fwd": True, + "mha_batch_prefill": True, +} + + +def _filter_mha_fwd(v: Variant) -> str: + """The CK ``--filter`` for mha_fwd, mirroring aiter/ops/mha.py. + + Two spellings differ from every other family and both are load-bearing: the + dtype token takes a leading underscore, and the mask token is ``_m*``, not + ``_mask*``. + """ + return ( + "*" + + f"_{v.dtype}*" + + "_nbias*" + + ("_m*" if v.needs_mask else "_nmask*") + + ("_lse*" if v.has_lse else "_nlse*") + + "_ndropout*" + + "_nqscale*" + ) + + +def _filter_batch_prefill(v: Variant) -> str: + """The CK ``--filter`` for mha_batch_prefill, mirroring aiter/ops/mha.py. + + The dtype token has **no** leading underscore here, unlike mha_fwd. + """ + return ( + "*" + + f"{v.dtype}*" + + ("_logits*" if v.has_logits_cap else "_nlogits*") + + "_nbias*" + + ("_mask*" if v.needs_mask else "_nmask*") + + ("_lse*" if v.has_lse else "_nlse*") + + "_ndropout*" + + "_nqscale*" + + "_nsink*" + ) + + +def reachable_variants() -> Tuple[Variant, ...]: + """Every ``.so`` the loader can ask for: 8 + 16 + 16 = 40 per architecture.""" + out: List[Variant] = [] + for family, has_logits_axis in _HAS_LOGITS_AXIS.items(): + logits_values = (False, True) if has_logits_axis else (False,) + for dtype in DTYPES: + for has_logits_cap in logits_values: + for needs_mask in (False, True): + for has_lse in (False, True): + out.append( + Variant(family, dtype, has_logits_cap, needs_mask, has_lse) + ) + return tuple(out) + + +def _is_default(v: Variant) -> bool: + """The set the image builds: everything except mha_fwd's LSE arm. + + Those four are the most expensive builds in the set (803 s each against 206 s + and 140 s), and LSE on *single* prefill serves cascade/merge rather than a + plain forward pass -- the Blaze-O1 LUM3 footprint names it unused. Skipping + them costs one ~5 min stall to anyone calling single prefill with + ``return_lse=True``; it is not an error. + """ + return not (v.family == "mha_fwd" and v.has_lse) + + +def selected_variants( + only: Optional[Sequence[str]] = None, everything: bool = False +) -> Tuple[Variant, ...]: + variants = reachable_variants() + if not everything: + variants = tuple(v for v in variants if _is_default(v)) + if only: + variants = tuple(v for v in variants if v.family in set(only)) + return variants + + +# -------------------------------------------------------------------------- +# Everything below needs the installed aiter; nothing above does, so the module +# stays importable (and testable) on a box with no aiter and no GPU. +# -------------------------------------------------------------------------- + + +def _aiter_jit_core(): + """Flat-import AITER's builder without running ``aiter/__init__.py``. + + ``aiter/setup.py`` does exactly this. Going through ``aiter.jit.core`` instead + runs the package __init__, which calls ``get_gfx_runtime()`` and imports + triton -- both fatal with no GPU attached, and with no env-var escape. + """ + import importlib.util + + spec = importlib.util.find_spec("aiter") + if spec is None or not spec.submodule_search_locations: + raise RuntimeError("amd-aiter is not installed") + aiter_dir = Path(list(spec.submodule_search_locations)[0]) + if str(aiter_dir) not in sys.path: + sys.path.insert(0, str(aiter_dir)) + from jit import core # type: ignore[import-not-found] + + return core + + +def build_recipe(v: Variant, core) -> Dict[str, object]: + """``{md_name, blob_gen_cmd}`` for one variant. + + varlen goes through AITER's own helper so those builds are byte-identical to + the lazy ones; the other two families have no helper and are composed here. + """ + ck_dir = core.CK_DIR + if v.family == "mha_varlen_fwd": + from jit.utils.mha_recipes import ( # type: ignore[import-not-found] + get_mha_varlen_prebuild_variants_by_names, + ) + + recipes = get_mha_varlen_prebuild_variants_by_names( + [v.md_name], ck_dir, receipt=RECEIPT_VARLEN + ) + if len(recipes) != 1 or recipes[0]["md_name"] != v.md_name: + raise RuntimeError( + f"aiter's varlen recipe helper returned {recipes!r} for {v.md_name}; " + "its md_name composition has changed" + ) + return recipes[0] + + if v.family == "mha_fwd": + direction, receipt, filt = "fwd", RECEIPT_MHA_FWD, _filter_mha_fwd(v) + else: + direction, receipt = "batch_prefill", RECEIPT_BATCH_PREFILL + filt = _filter_batch_prefill(v) + return { + "md_name": v.md_name, + "blob_gen_cmd": [ + f"{ck_dir}/{GENERATE_PY} -d {direction} " + f"--receipt {receipt} --filter {filt} --output_dir {{}}" + ], + } + + +def jit_dir(core) -> Path: + return Path(core.get_user_jit_dir()) + + +def _args_of_build(core, family: str) -> Dict[str, object]: + """AITER's compile flags for a family, keyed as `compile_ops` keys them. + + The config key carries a ``module_`` prefix; without it ``get_args_of_build`` + *warns and returns empty args* rather than raising, which drops every include + path and fails as a missing header several minutes into the build. Assert + rather than trust, so a renamed key is loud. + """ + args = core.get_args_of_build(f"module_{family}") + if not args.get("srcs") or not args.get("extra_include"): + raise RuntimeError( + f"aiter's optCompilerConfig.json gave no srcs/includes for " + f"module_{family}; the config key has been renamed" + ) + return args + + +def build_one(v: Variant) -> Path: + """Build a single variant. Returns the artifact path; raises if absent after.""" + core = _aiter_jit_core() + recipe = build_recipe(v, core) + args = _args_of_build(core, v.family) + core.build_module( + recipe["md_name"], + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + recipe["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args.get("torch_exclude", False), + args.get("third_party", []), + ) + out = jit_dir(core) / v.so_name + if not out.is_file(): + raise RuntimeError(f"build of {v.md_name} produced no {out}") + return out + + +def _run_jobs(variants: Sequence[Variant], jobs: int) -> int: + """Build in parallel across *processes*. + + Never threads: AITER's build mutates process-global environment and takes a + process-global lock, so two concurrent builds in one process have the first + to finish restore the environment under the second. + """ + pending = list(variants) + # Child output goes to a file, never a pipe: hipcc is verbose enough to fill + # a pipe buffer, and a child blocked on a full pipe never exits, so poll() + # would spin forever. Only read on failure, and only the tail. + running: List[Tuple[Variant, subprocess.Popen, float, str]] = [] + failed: List[str] = [] + done = 0 + total = len(pending) + + while pending or running: + while pending and len(running) < jobs: + v = pending.pop(0) + fd, log_path = tempfile.mkstemp( + prefix=f"prebuild-{v.md_name}-", suffix=".log" + ) + try: + proc = subprocess.Popen( + [sys.executable, __file__, "--build-one", _encode(v)], + stdout=fd, + stderr=subprocess.STDOUT, + ) + finally: + os.close(fd) # Popen dup'd it; holding ours would exhaust the table + running.append((v, proc, time.time(), log_path)) + time.sleep(1.0) + for entry in list(running): + v, proc, started, log_path = entry + if proc.poll() is None: + continue + running.remove(entry) + done += 1 + if proc.returncode == 0: + print( + f"[{done}/{total}] ok {v.so_name} {time.time() - started:.0f}s", + flush=True, + ) + os.unlink(log_path) + else: + failed.append(v.so_name) + tail = Path(log_path).read_text(errors="replace").splitlines()[-15:] + print( + f"[{done}/{total}] FAIL {v.so_name} {time.time() - started:.0f}s" + f" (full log: {log_path})\n" + "\n".join(tail), + flush=True, + ) + + if failed: + print(f"\n{len(failed)} build(s) failed:\n " + "\n ".join(failed)) + return 1 if failed else 0 + + +def _encode(v: Variant) -> str: + return "|".join( + [ + v.family, + v.dtype, + str(int(v.has_logits_cap)), + str(int(v.needs_mask)), + str(int(v.has_lse)), + ] + ) + + +def _decode(s: str) -> Variant: + family, dtype, lc, mask, lse = s.split("|") + return Variant(family, dtype, bool(int(lc)), bool(int(mask)), bool(int(lse))) + + +def _check(variants: Sequence[Variant]) -> int: + """Assert every selected artifact is present and non-empty.""" + core = _aiter_jit_core() + d = jit_dir(core) + missing = [ + v.so_name + for v in variants + if not (d / v.so_name).is_file() or (d / v.so_name).stat().st_size == 0 + ] + print(f"{len(variants) - len(missing)}/{len(variants)} present in {d}") + if missing: + print("missing or empty:\n " + "\n ".join(missing)) + return 1 + return 0 + + +def main(argv: Optional[Sequence[str]] = None) -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument( + "--list", action="store_true", help="print the selected set and exit" + ) + p.add_argument( + "--check", action="store_true", help="verify artifacts, build nothing" + ) + p.add_argument( + "--all", + action="store_true", + help="build all 40 reachable variants, not just the default set", + ) + p.add_argument( + "--only", + action="append", + choices=sorted(_MD_NAME), + help="restrict to a family (repeatable)", + ) + p.add_argument( + "--jobs", + type=int, + default=int(os.environ.get("AITER_PREBUILD_JOBS", "1")), + help="concurrent build processes", + ) + p.add_argument("--build-one", help=argparse.SUPPRESS) + a = p.parse_args(argv) + + if a.build_one: + build_one(_decode(a.build_one)) + return 0 + + variants = selected_variants(a.only, everything=a.all) + if a.list: + for v in variants: + print(v.so_name) + print(f"\n{len(variants)} variant(s)") + return 0 + if a.check: + return _check(variants) + + if a.jobs < 1: + p.error("--jobs must be >= 1") + t0 = time.time() + rc = _run_jobs(variants, a.jobs) + print( + f"\n{len(variants)} variant(s) in {(time.time() - t0) / 60:.0f} min " + f"at --jobs {a.jobs}" + ) + return rc or _check(variants) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py new file mode 100644 index 00000000000..510bd2dfc11 --- /dev/null +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -0,0 +1,219 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Parity for docker/prebuild_aiter_attention.py. + +That driver runs inside `docker build` with no GPU and no torch, so it cannot +import flashinfer's variant table and carries its own copy. Two things can go +wrong silently, and one test guards each: + +1. Its filenames drift from what `aiter_loader.cc` asks for. The artifact is + built and then never found -- a silent latency loss, not an error. +2. Its CK receipts or filter tokens drift from AITER's. The right filename gets + a *different kernel* inside, which nothing downstream detects. + +Both read the installed AITER source rather than a checkout: the two disagree, +and only the installed one is what runs. No build, no GPU. +""" + +import importlib.util +import re +from pathlib import Path + +import pytest + +from flashinfer.jit.rocm import aiter_variants as fi_variants + +_DRIVER_PATH = ( + Path(__file__).resolve().parents[2] / "docker" / "prebuild_aiter_attention.py" +) + + +def _load_driver(): + spec = importlib.util.spec_from_file_location( + "_prebuild_aiter_attention", _DRIVER_PATH + ) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +driver = _load_driver() + + +def _aiter_ops_mha_source() -> str: + spec = importlib.util.find_spec("aiter") + if spec is None or not spec.submodule_search_locations: + pytest.skip("amd-aiter is not installed") + src = Path(list(spec.submodule_search_locations)[0]) / "ops" / "mha.py" + if not src.is_file(): + pytest.skip(f"{src} not found") + return src.read_text() + + +def _composer_body(family: str) -> str: + """The body of AITER's `cmdGenFunc_`, where md_name and filter are built.""" + text = _aiter_ops_mha_source() + start = text.find(f"def cmdGenFunc_{family}(") + assert start != -1, f"aiter/ops/mha.py has no cmdGenFunc_{family}" + nxt = text.find("\ndef ", start + 1) + return text[start : nxt if nxt != -1 else len(text)] + + +# --------------------------------------------------------------------------- +# 1. filenames match what the loader asks for +# --------------------------------------------------------------------------- + +_FAMILY_BY_NAME = { + "mha_fwd": fi_variants.Family.MHA_FWD, + "mha_varlen_fwd": fi_variants.Family.MHA_VARLEN_FWD, + "mha_batch_prefill": fi_variants.Family.MHA_BATCH_PREFILL, +} + + +def _fi_so_names() -> set: + return {fi_variants.so_name(k) for k in fi_variants.reachable_variants()} + + +def test_the_driver_enumerates_the_same_40_variants_as_flashinfer(): + assert ( + len(driver.reachable_variants()) == len(fi_variants.reachable_variants()) == 40 + ) + + +def test_every_driver_filename_is_one_the_loader_asks_for(): + """A name the loader never asks for is built and then never opened.""" + assert {v.so_name for v in driver.reachable_variants()} == _fi_so_names() + + +def test_the_default_set_is_the_full_set_minus_mha_fwd_lse(): + selected = driver.selected_variants() + skipped = set(driver.reachable_variants()) - set(selected) + + assert len(selected) == 36 + assert all(v.family == "mha_fwd" and v.has_lse for v in skipped) + assert len(skipped) == 4 + + +def test_the_blaze_lum3_variant_is_in_the_default_set(): + """Blaze-O1's LUM3 pins fp16 / batch mode / causal / no soft-cap / no LSE; + batch mode is mha_fwd, so this exact file is the one it dlopens.""" + assert "mha_fwd_fp16_nbias_mask_nlse_ndropout_nqscale.so" in { + v.so_name for v in driver.selected_variants() + } + + +@pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) +def test_only_restricts_without_renaming(family): + picked = driver.selected_variants(only=[family], everything=True) + assert picked + assert {v.so_name for v in picked} <= _fi_so_names() + assert all(v.family == family for v in picked) + + +# --------------------------------------------------------------------------- +# 2. receipts and token spellings match the installed aiter +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("family", "direction", "receipt"), + [ + ("mha_fwd", "fwd", 100), + ("mha_varlen_fwd", "fwd", 200), + ("mha_batch_prefill", "batch_prefill", 200), + ], +) +def test_the_receipt_and_direction_match_aiter(family, direction, receipt): + """A wrong receipt ships a different instance set under the right filename.""" + body = _composer_body(family) + assert f"-d {direction} " in body + assert f"--receipt {receipt} " in body + + from_driver = { + "mha_fwd": driver.RECEIPT_MHA_FWD, + "mha_varlen_fwd": driver.RECEIPT_VARLEN, + "mha_batch_prefill": driver.RECEIPT_BATCH_PREFILL, + }[family] + assert from_driver == receipt + + +def test_mha_fwd_masks_on_underscore_m_not_underscore_mask(): + """The one filter token that differs between families. Spelling `_mask*` here + selects nothing, so the module builds and dispatch then finds no kernel.""" + body = _composer_body("mha_fwd") + assert 'filter += "_m*"' in body + + v = driver.Variant("mha_fwd", "bf16", False, needs_mask=True, has_lse=False) + assert "_m*" in driver._filter_mha_fwd(v) + assert "_mask*" not in driver._filter_mha_fwd(v) + + +def test_batch_prefill_dtype_filter_has_no_leading_underscore(): + """Unlike mha_fwd's `_bf16*`, batch_prefill appends a bare `bf16*`.""" + body = _composer_body("mha_batch_prefill") + assert 'filter_fwd += "bf16*"' in body + + v = driver.Variant("mha_batch_prefill", "bf16", False, True, False) + assert "*bf16*" in driver._filter_batch_prefill(v) + assert "_bf16*" not in driver._filter_batch_prefill(v) + + +def _mha_recipes_source() -> str: + spec = importlib.util.find_spec("aiter") + if spec is None or not spec.submodule_search_locations: + pytest.skip("amd-aiter is not installed") + src = ( + Path(list(spec.submodule_search_locations)[0]) + / "jit" + / "utils" + / "mha_recipes.py" + ) + if not src.is_file(): + pytest.skip(f"{src} not found") + return src.read_text() + + +def _md_tokens(family: str) -> set: + """Every name token AITER can emit for this family, dtype excluded. + + mha_fwd and batch_prefill append string literals to md_name directly. varlen + delegates its whole suffix to `compose_mha_fwd_variant_suffix_and_filter`, so + its literals live in mha_recipes.py instead. Both spell the dtype with an + f-string, which leaves no literal to find -- hence dtype is stripped by the + caller rather than matched here. + """ + source = ( + _mha_recipes_source() if family == "mha_varlen_fwd" else _composer_body(family) + ) + return set(re.findall(r'"(_[a-z0-9]+)"', source)) + + +@pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) +def test_every_md_name_decomposes_into_aiter_token_literals(family): + """Catches a renamed or reordered token: each name we compose must be the + family base plus a concatenation of literals AITER itself appends.""" + tokens = _md_tokens(family) + assert tokens, f"no md_name += literals found in cmdGenFunc_{family}" + + for v in driver.reachable_variants(): + if v.family != family: + continue + assert v.md_name.startswith(f"{family}_{v.dtype}") + rest = v.md_name[len(f"{family}_{v.dtype}") :] + while rest: + match = max( + (t for t in tokens if rest.startswith(t)), key=len, default=None + ) + assert match, f"{v.md_name!r}: no aiter token matches at {rest!r}" + rest = rest[len(match) :] + + +def test_varlen_goes_through_aiters_own_helper(): + """varlen is the control: its builds must stay byte-identical to the lazy + ones, which only holds while we call AITER's composer rather than mirror it.""" + src = _DRIVER_PATH.read_text() + assert "get_mha_varlen_prebuild_variants_by_names" in src + assert "_md_varlen" in src # the mirror exists only to name the file we ask for From 6dd43ef72018f9699363e8d011dd8c325491b967 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 23:23:11 -0400 Subject: [PATCH 06/41] build(rocm): build AITER from source at v0.1.20 and prebuild the variants Wires up the driver from the previous commit. Deliberately still at v0.1.20 -- the ABI the image already runs -- so that a failure here is the *mechanism* and a failure after the v0.1.21.post2 bump is the *revision*. Do not collapse the two commits; the bisect depends on the split. Source rather than wheel because it is the only way to compile AITER against the ROCm in this image. No +rocm10.0 build of amd-aiter is published, so the pinned wheel is a +rocm10.1.0a retarget of the same revision. AITER_SOURCE=0 restores that wheel install verbatim for an unblock. Three things a reader cannot see by reading, so they are checked rather than asserted in prose (tmp validator, not shipped): - The clone is created and deleted in ONE RUN. Splitting them leaves it in the earlier layer, where a later `rm -rf` cannot reach it -- CK alone is over a GB. - The driver COPY sits after the install, so editing the variant list re-runs only the prebuild. - Every AITER_* the new layers reference is declared with ARG in this stage. Non-editable install on purpose: `setup.py develop` takes an is_develop_mode() branch setting packages = ["aiter"] with no aiter_meta, and aiter_meta/csrc/include is what every C++ shim compiles against. The non-editable path runs prepare_packaging(), which copies csrc and the CK tree into aiter_meta -- which is also what makes deleting the clone safe. PREBUILD_KERNELS=0 because 1 rmtree's the jit dir and does `from aiter.aot.flydsl.common import run_aot`, a real `import aiter`. Losing its prebuilt modules costs nothing measurable: FlashInfer links its own lib.so built from the same sources, and of AITER's 111 prebuilt modules only module_mla_asm is ever loaded (11.5 s, 0.2 MB). The post-install check is a hard assertion, not a print. A shallow clone that lost its tag makes setuptools_scm invent a .dev version, which falls under aiter_utils.AITER_MIN_VERSION and silently degrades every AITER path to fa2 -- a working image that is quietly several times slower on prefill. GPU_ARCHS must be explicit: there is no device here for AITER's arch detection, and both architectures go in one fat binary. Validated: 12 RUN instructions parse under `sh -n`, the embedded Python heredoc parses, and all five invariants A/B'd by breaking each one and confirming the check fires. This buildx has neither --check nor --call, and a real image build is ~2h, so the build-time number lands in the PR description rather than here. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 98 ++++++++++++++++++++++++++++++++++++++---- 1 file changed, 90 insertions(+), 8 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 68010dda823..7678862d90b 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -14,6 +14,11 @@ # detection *and* imports triton, which fails with # "0 active drivers ([]). There should only be one." when no GPU is attached, # and that has no env-var escape. Install it here; verify it at run time. +# AITER's *builder* is reachable without that import, which is what lets the +# variant prebuild below run here: aiter/setup.py:212 does +# `sys.path.insert(0, ".../aiter/"); from jit import core`, a flat import +# that never executes aiter/__init__.py. Do not "simplify" it to +# `from aiter.jit import core` -- that runs the __init__ and needs a GPU. ARG ROCM_VERSION=10.0 ARG UBUNTU_VERSION=24.04 ARG PY_VERSION=3.12 @@ -31,10 +36,22 @@ FROM rocm/pytorch:rocm${ROCM_VERSION}_ubuntu${UBUNTU_VERSION}_py${PY_VERSION}_py ARG ROCM_VERSION ARG PY_VERSION -# AITER is pinned to an exact build including the local version segment, because -# pip will not select a local version from a loose specifier. No +rocm10.0 build -# of amd-aiter exists; every 0.1.20 wheel shares source build id 3135022, so the -# +rocm10.1.0a one is the nearest retarget of the same revision. +# AITER is built from source at a pinned tag, so it compiles against the ROCm in +# this image rather than a retargeted wheel. AITER_SOURCE=0 restores the wheel +# install below for an unblock. +ARG AITER_SOURCE=1 +ARG AITER_REF=v0.1.20 +ARG AITER_REPO=https://github.com/ROCm/aiter.git +# Both architectures in one fat binary. Must be explicit: there is no device here +# for AITER's arch detection to find. +ARG AITER_GPU_ARCHS="gfx942;gfx950" +ARG AITER_PREBUILD_JOBS=2 + +# The wheel path (AITER_SOURCE=0). Pinned to an exact build including the local +# version segment, because pip will not select a local version from a loose +# specifier. No +rocm10.0 build of amd-aiter exists; every 0.1.20 wheel shares +# source build id 3135022, so the +rocm10.1.0a one is the nearest retarget of the +# same revision. ARG AITER_VERSION=0.1.20+rocm10.1.0a20260819.3135022 ARG AITER_INDEX=https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/ @@ -133,6 +150,7 @@ USER $USERNAME # PyPI and pulls a non-ROCm wheel. The torch check repeats after AITER for the # same reason -- --extra-index-url leaves PyPI visible, so a dependency of the # AITER wheel can swap in a CPU-only torch and the build still exits 0. +# Its own layer, so bumping AITER below does not reinstall all of this. RUN pip install --no-cache-dir \ ninja \ "setuptools>=80" \ @@ -144,10 +162,74 @@ RUN pip install --no-cache-dir \ pybind11 \ ruff \ einops \ - filelock && \ - pip install --no-cache-dir "amd-aiter==${AITER_VERSION}" --extra-index-url "${AITER_INDEX}" && \ - python3 -c "import importlib.metadata as m; print('amd-aiter', m.version('amd-aiter'))" && \ - python3 -c "import torch, sys; hip = torch.version.hip; print(f'torch {torch.__version__} (hip={hip})'); sys.exit('ERROR: the amd-aiter install replaced torch with a non-ROCm build' if hip is None else 0)" + filelock + +# Clone, install and delete in ONE layer. Splitting them would leave the clone +# (CK alone is over a GB) in the earlier layer, where a later `rm -rf` cannot +# reach it and the image carries it forever. +# +# Non-editable on purpose: `setup.py develop` takes an is_develop_mode() branch +# that sets packages = ["aiter"] with no aiter_meta, and aiter_meta/csrc/include +# is what every C++ shim compiles against. The non-editable install runs +# prepare_packaging(), which copies csrc and the CK tree into aiter_meta -- which +# is also what makes deleting the clone safe. +# +# PREBUILD_KERNELS=1 is unusable here: it rmtree's the jit dir and does +# `from aiter.aot.flydsl.common import run_aot`, a real `import aiter`. Its +# prebuilt modules are no loss -- FlashInfer links its own lib.so built +# from these same sources, and measured, only aiter.mla loads one of AITER's. +RUN set -eu; \ + if [ "$AITER_SOURCE" = "1" ]; then \ + git clone --recursive --depth 1 --shallow-submodules \ + --branch "$AITER_REF" "$AITER_REPO" /tmp/aiter && \ + ( cd /tmp/aiter && \ + PREBUILD_KERNELS=0 GPU_ARCHS="$AITER_GPU_ARCHS" \ + pip install --no-cache-dir --no-build-isolation . ) && \ + rm -rf /tmp/aiter; \ + else \ + pip install --no-cache-dir "amd-aiter==${AITER_VERSION}" \ + --extra-index-url "${AITER_INDEX}"; \ + fi + +# Hard assertions, not prints. A shallow clone that lost its tag makes +# setuptools_scm invent a .dev version, which would fall under +# aiter_utils.AITER_MIN_VERSION and silently degrade every AITER path to fa2 -- +# a working image that is quietly half as fast. The torch check repeats here +# because --extra-index-url leaves PyPI visible, so a dependency can swap in a +# CPU-only torch and the build still exits 0. +RUN python3 - <<'PY' +import importlib.metadata as md +import os +import sys +from pathlib import Path + +got = md.version("amd-aiter") +ref, version = os.environ["AITER_REF"], os.environ["AITER_VERSION"] +want = ref[1:] if ref.startswith("v") else ref +want = want if os.environ["AITER_SOURCE"] == "1" else version +if got != want: + sys.exit(f"ERROR: amd-aiter resolved to {got!r}, expected {want!r}") + +inc = Path(md.distribution("amd-aiter").locate_file("aiter_meta/csrc/include")) +if not inc.is_dir(): + sys.exit(f"ERROR: {inc} is missing; the shims cannot compile against this install") + +import torch + +if torch.version.hip is None: + sys.exit("ERROR: the amd-aiter install replaced torch with a non-ROCm build") +print(f"amd-aiter {got}; torch {torch.__version__} (hip={torch.version.hip})") +PY + +# Placed after the install so editing the variant list re-runs only the prebuild. +COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuild_aiter_attention.py + +# Compile the attention variants that the first request would otherwise compile +# in-process, under a process-global lock, minutes at a time. The driver ends in +# its own --check, so a missing or empty artifact fails the build here. +RUN set -eu; \ + GPU_ARCHS="$AITER_GPU_ARCHS" \ + python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" WORKDIR /home/$USERNAME From bfeab6e5f62174cb64ebd0a9cb4b2af410c97934 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 23:45:07 -0400 Subject: [PATCH 07/41] build(rocm): move the AITER pin to v0.1.21.post2, ABI and all The plan split this into three commits -- shim migration, then AITER_REF, then the floor raise -- on the claim that the POD overloads already existed at 0.1.20 and the migration could land first. They do not: 0.1.20 rmsnorm_quant.h:37 void rmsnorm(torch::Tensor& out, ...) 0.1.20 grep aiter_tensor_t {rmsnorm,rmsnorm_quant,rope}.h -> no matches 0.1.21 aiter::rmsnorm(aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&, double, bool) Either order leaves a commit whose rmsnorm/rope backend="aiter" tests fail at dlopen with `undefined symbol: _ZN5aiter7rmsnormERN2at6TensorES2_S2_db`. They are one atomic change. The floor raise joins them because a 0.1.20 user would otherwise clear the gate and then fail to load. The 5/6 bisect split the plan wanted is intact: commit 5 proved the *mechanism* at 0.1.20, this moves the *revision*. - norm_aiter.cu and rope_aiter.cu now declare the POD signatures and convert through csrc/rocm/aiter_tensor_compat.h, matching activation/page/fused_moe, which migrated at 0.1.16 and already pass. Named lvalues at the call sites: the out parameters are non-const references, so a to_aiter() temporary would not bind. - AITER_MIN_VERSION 0.1.20 -> 0.1.21, _AITER_LAST_VALIDATED -> 0.1.21.post2 (no +rocm segment; a source build has none), the vendored header pins, and kAbiPinNote, which now names the source build rather than a dead pip line. Not touched: arch_caps.py evidence strings stay at 0.1.20 because they assert a measurement that has not been repeated -- honestly stale beats quietly wrong. The soft-cap re-measurement is the next commit. Two plan items turned out to be non-issues, checked rather than assumed: _AITER_SOFTCAP_DEFECT_THROUGH appears only in message text, and test_single_prefill_kernels.py:92's skip is already driven by arch_caps.aiter_softcap_defect_arch, so neither stops firing at 0.1.21.post2. A/B by construction: all 69 rmsnorm/rope AITER tests failed at dlopen on 0.1.21.post2 before this change and pass after, measured on gfx942. The version floor table gains 0.1.20 -> False and 0.1.21.dev0 -> False. Co-Authored-By: Claude --- csrc/rocm/aiter_loader.cc | 14 ++++---- csrc/rocm/norm_aiter.cu | 35 +++++++++++++------ csrc/rocm/rope_aiter.cu | 26 +++++++++----- docker/Dockerfile.rocm | 2 +- flashinfer/rocm/aiter_utils.py | 2 +- flashinfer/rocm/prefill.py | 2 +- .../attention/aiter/mha_batch_prefill_args.h | 3 +- .../rocm/attention/aiter/mha_fwd_args.h | 2 +- tests/rocm/test_aiter_version_gate.py | 5 ++- 9 files changed, 60 insertions(+), 31 deletions(-) diff --git a/csrc/rocm/aiter_loader.cc b/csrc/rocm/aiter_loader.cc index c54f9eedb2d..226ec578ba4 100644 --- a/csrc/rocm/aiter_loader.cc +++ b/csrc/rocm/aiter_loader.cc @@ -21,12 +21,12 @@ namespace flashinfer::aiter { namespace { // Appended to failures on the paths that resolve the pinned mangled symbols -// below. Keep the version in step with _AITER_LAST_VALIDATED in prefill_rocm.py. +// below. Keep the version in step with _AITER_LAST_VALIDATED in flashinfer/rocm/prefill.py. constexpr const char* kAbiPinNote = - "\n These symbols and .so names are pinned to amd-aiter 0.1.20, which has no stable" - "\n C++ ABI. If AITER was upgraded, that is the likely cause. Reinstall the pin:" - "\n pip install amd-aiter==0.1.20+rocm10.1.0a20260819.3135022 \\" - "\n --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/" + "\n These symbols and .so names are pinned to amd-aiter 0.1.21.post2, which has no" + "\n stable C++ ABI -- 0.1.21 moved rmsnorm and rope from at::Tensor to the POD" + "\n aiter_tensor_t. If AITER was changed, that is the likely cause. The supported" + "\n install is the source build docker/Dockerfile.rocm performs at tag v0.1.21.post2," "\n or re-pin the symbols in csrc/rocm/aiter_loader.cc."; // Directories to try, in order, when resolving a variant .so. @@ -191,7 +191,7 @@ void* load_path_sym(std::shared_mutex& mu, std::unordered_map& // Mangled symbol for aiter::mha_fwd(aiter::mha_fwd_args, ck_tile::stream_config const&). // Stable across GCC/Clang Itanium ABI; verified by `nm -D` on all shipped variants. // Both the mha_fwd and mha_varlen_fwd .so files export this same dispatcher symbol. -// Pinned to amd-aiter 0.1.20. Regenerate with: nm -D | grep mha_fwd +// Pinned to amd-aiter 0.1.21.post2. Regenerate with: nm -D | grep mha_fwd constexpr const char* kMhaFwdSymbol = "_ZN5aiter7mha_fwdENS_12mha_fwd_argsERKN7ck_tile13stream_configE"; @@ -235,7 +235,7 @@ std::string batch_prefill_variant_so_name(BatchPrefillVariantKey const& key) { } // Itanium-ABI mangled symbol for aiter::mha_batch_prefill(...). -// Pinned to amd-aiter 0.1.20. Regenerate with: nm -D <.so> | grep mha_batch_prefill +// Pinned to amd-aiter 0.1.21.post2. Regenerate with: nm -D <.so> | grep mha_batch_prefill // Note: '23' encodes len("fmha_batch_prefill_args") == 23. constexpr const char* kMhaBatchPrefillSymbol = "_ZN5aiter17mha_batch_prefillE23fmha_batch_prefill_argsRKN7ck_tile13stream_configE" diff --git a/csrc/rocm/norm_aiter.cu b/csrc/rocm/norm_aiter.cu index a4344ad9755..e30edce6f3f 100644 --- a/csrc/rocm/norm_aiter.cu +++ b/csrc/rocm/norm_aiter.cu @@ -17,15 +17,18 @@ #include #include +#include "aiter_tensor_compat.h" + // AITER's public header (rmsnorm_quant.h) pulls in → full -// pybind11, which clashes with FlashInfer's -DPy_LIMITED_API. torch::Tensor is -// at::Tensor, so forward-declare the entry points; the linker resolves them -// against the symbol-visible AITER .so. 0.1.20 renamed these from the global -// rmsnorm2d/rmsnorm2d_with_add and moved them into namespace aiter. +// pybind11, which clashes with FlashInfer's -DPy_LIMITED_API, so forward-declare +// the entry points; the linker resolves them against the symbol-visible AITER +// .so. 0.1.21 moved both onto the POD aiter_tensor_t API -- declaring them with +// at::Tensor compiles and then fails at dlopen on the mangled name. namespace aiter { -void add_rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& residual_in, - at::Tensor& residual_out, at::Tensor& weight, double epsilon, bool gemma_norm); -void rmsnorm(at::Tensor& out, at::Tensor& input, at::Tensor& weight, double epsilon, +void add_rmsnorm(aiter_tensor_t& out, aiter_tensor_t& input, aiter_tensor_t& residual_in, + aiter_tensor_t& residual_out, aiter_tensor_t& weight, double epsilon, + bool gemma_norm); +void rmsnorm(aiter_tensor_t& out, aiter_tensor_t& input, aiter_tensor_t& weight, double epsilon, bool gemma_norm); } // namespace aiter @@ -42,7 +45,15 @@ void fused_add_rmsnorm_aiter(at::Tensor input, at::Tensor residual, at::Tensor w // several rows per block for small n, so a write lands on a row still unread. at::Tensor out = at::empty_like(input); at::Tensor residual_out = at::empty_like(residual); - aiter::add_rmsnorm(out, input, residual, residual_out, weight2d, eps, /*gemma_norm=*/false); + // Named lvalues: every POD parameter is a non-const reference, so a temporary + // from to_aiter() would not bind. + auto a_out = flashinfer::aiter_compat::to_aiter(out); + auto a_input = flashinfer::aiter_compat::to_aiter(input); + auto a_residual = flashinfer::aiter_compat::to_aiter(residual); + auto a_residual_out = flashinfer::aiter_compat::to_aiter(residual_out); + auto a_weight = flashinfer::aiter_compat::to_aiter(weight2d); + aiter::add_rmsnorm(a_out, a_input, a_residual, a_residual_out, a_weight, eps, + /*gemma_norm=*/false); input.copy_(out); residual.copy_(residual_out); } @@ -54,11 +65,15 @@ void rmsnorm_aiter(at::Tensor out, at::Tensor input, at::Tensor weight, double e at::Tensor weight2d = weight.reshape({1, -1}); // rmsnorm(x, w, out=x) is a documented idiom, and that alias corrupts the // same way; stage it only when the caller actually aliased. + auto a_input = flashinfer::aiter_compat::to_aiter(input); + auto a_weight = flashinfer::aiter_compat::to_aiter(weight2d); if (out.data_ptr() == input.data_ptr()) { at::Tensor staged = at::empty_like(out); - aiter::rmsnorm(staged, input, weight2d, eps, /*gemma_norm=*/false); + auto a_staged = flashinfer::aiter_compat::to_aiter(staged); + aiter::rmsnorm(a_staged, a_input, a_weight, eps, /*gemma_norm=*/false); out.copy_(staged); return; } - aiter::rmsnorm(out, input, weight2d, eps, /*gemma_norm=*/false); + auto a_out = flashinfer::aiter_compat::to_aiter(out); + aiter::rmsnorm(a_out, a_input, a_weight, eps, /*gemma_norm=*/false); } diff --git a/csrc/rocm/rope_aiter.cu b/csrc/rocm/rope_aiter.cu index b812a0a46f9..4ff842fe0df 100644 --- a/csrc/rocm/rope_aiter.cu +++ b/csrc/rocm/rope_aiter.cu @@ -18,14 +18,17 @@ #include #include +#include "aiter_tensor_compat.h" + // AITER's public header (rope.h) pulls in → full pybind11, -// which clashes with FlashInfer's -DPy_LIMITED_API. torch::Tensor is at::Tensor, -// so forward-declare the entry point instead; the linker resolves it against the -// symbol-visible AITER .so (verified mangling matches at::Tensor& signature). -void rope_cached_positions_2c_fwd_impl(at::Tensor& output_x, at::Tensor& output_y, - const at::Tensor& input_x, const at::Tensor& input_y, - const at::Tensor& cos, const at::Tensor& sin, - const at::Tensor& positions, const int32_t rotate_style, +// which clashes with FlashInfer's -DPy_LIMITED_API, so forward-declare the entry +// point instead; the linker resolves it against the symbol-visible AITER .so. +// 0.1.21 moved it onto the POD aiter_tensor_t API -- declaring it with at::Tensor +// compiles and then fails at dlopen on the mangled name. +void rope_cached_positions_2c_fwd_impl(aiter_tensor_t& output_x, aiter_tensor_t& output_y, + const aiter_tensor_t& input_x, const aiter_tensor_t& input_y, + const aiter_tensor_t& cos, const aiter_tensor_t& sin, + const aiter_tensor_t& positions, const int32_t rotate_style, const bool reuse_freqs_front_part, const bool nope_first); void apply_rope_pos_ids_cos_sin_cache_aiter(at::Tensor query, at::Tensor key, at::Tensor query_out, @@ -77,7 +80,14 @@ void apply_rope_pos_ids_cos_sin_cache_aiter(at::Tensor query, at::Tensor key, at at::Tensor q_in_rot = q_in.slice(3, 0, rotary_dim); at::Tensor k_in_rot = k_in.slice(3, 0, rotary_dim); - rope_cached_positions_2c_fwd_impl(q_out_rot, k_out_rot, q_in_rot, k_in_rot, cos, sin, pos, + // Named lvalues: the two outputs are non-const references, so a temporary + // from to_aiter() would not bind. + namespace compat = flashinfer::aiter_compat; + auto a_q_out = compat::to_aiter(q_out_rot); + auto a_k_out = compat::to_aiter(k_out_rot); + rope_cached_positions_2c_fwd_impl(a_q_out, a_k_out, compat::to_aiter(q_in_rot), + compat::to_aiter(k_in_rot), compat::to_aiter(cos), + compat::to_aiter(sin), compat::to_aiter(pos), /*rotate_style=*/is_neox ? 0 : 1, /*reuse_freqs_front_part=*/true, /*nope_first=*/false); diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7678862d90b..7feadd13696 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -40,7 +40,7 @@ ARG PY_VERSION # this image rather than a retargeted wheel. AITER_SOURCE=0 restores the wheel # install below for an unblock. ARG AITER_SOURCE=1 -ARG AITER_REF=v0.1.20 +ARG AITER_REF=v0.1.21.post2 ARG AITER_REPO=https://github.com/ROCm/aiter.git # Both architectures in one fat binary. Must be explicit: there is no device here # for AITER's arch detection to find. diff --git a/flashinfer/rocm/aiter_utils.py b/flashinfer/rocm/aiter_utils.py index 3b4f1c49637..4ccf77827a0 100644 --- a/flashinfer/rocm/aiter_utils.py +++ b/flashinfer/rocm/aiter_utils.py @@ -54,7 +54,7 @@ def _ensure_aiter_gpu_archs() -> None: # an older AITER cannot resolve the shim's symbols at all. The vendored prefill # structs travel by value through dlsym'd pointers, where a mismatch corrupts # silently -- hence a hard floor rather than a warning. -AITER_MIN_VERSION = "0.1.20" +AITER_MIN_VERSION = "0.1.21" def _aiter_installed_version() -> Optional[str]: diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index 844b45caaf0..af3b3b77cae 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -64,7 +64,7 @@ # sizes we try. The second is the newest release we have actually validated against; # bumping it must not silently move the support boundary. _AITER_NATIVE_PAGING_SINCE = "0.1.10" -_AITER_LAST_VALIDATED = "0.1.20+rocm10.1.0a20260819.3135022" +_AITER_LAST_VALIDATED = "0.1.21.post2" # Newest AITER carrying the mha_varlen_fwd soft-cap defect. Bump only after # re-measuring against an fp32 reference; the wrong answer is silent. _AITER_SOFTCAP_DEFECT_THROUGH = "0.1.21" diff --git a/include/flashinfer/rocm/attention/aiter/mha_batch_prefill_args.h b/include/flashinfer/rocm/attention/aiter/mha_batch_prefill_args.h index 41c0e197ade..be23fef26f5 100644 --- a/include/flashinfer/rocm/attention/aiter/mha_batch_prefill_args.h +++ b/include/flashinfer/rocm/attention/aiter/mha_batch_prefill_args.h @@ -1,7 +1,8 @@ // SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 // -// Vendored fmha_batch_prefill_args, following the amd-aiter 0.1.20 layout. +// Vendored fmha_batch_prefill_args, following the amd-aiter 0.1.21 layout +// (field-for-field identical to 0.1.20's; only the norm and rope APIs moved). // Extracted from // aiter_meta/3rdparty/composable_kernel/example/ck_tile/01_fmha/fmha_fwd.hpp. // diff --git a/include/flashinfer/rocm/attention/aiter/mha_fwd_args.h b/include/flashinfer/rocm/attention/aiter/mha_fwd_args.h index 77593ac18ea..12401bf88dd 100644 --- a/include/flashinfer/rocm/attention/aiter/mha_fwd_args.h +++ b/include/flashinfer/rocm/attention/aiter/mha_fwd_args.h @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. // SPDX-License-Identifier: Apache-2.0 // -// Vendored aiter::mha_fwd_args from AITER amd-aiter>=0.1.20. +// Vendored aiter::mha_fwd_args from AITER amd-aiter>=0.1.21. // Extracted from aiter_meta/csrc/include/mha_fwd.h. // // The floor is a hard minimum, not a recommendation: the block_scale_seqstart_* diff --git a/tests/rocm/test_aiter_version_gate.py b/tests/rocm/test_aiter_version_gate.py index 90b34a8223b..d9a64eeca0f 100644 --- a/tests/rocm/test_aiter_version_gate.py +++ b/tests/rocm/test_aiter_version_gate.py @@ -42,7 +42,10 @@ def _doc_references(): # 0.1.16 declared rmsnorm without gemma_norm, so its mangled name differs. ("0.1.16", False), ("0.1.20.dev0", False), - ("0.1.20", True), + # 0.1.21 moved rmsnorm and rope from at::Tensor to the POD aiter_tensor_t, + # so the shims no longer link against 0.1.20 at all. + ("0.1.20", False), + ("0.1.21.dev0", False), # A .dev0 *of post3* sorts above 0.1.16, which is what makes the # nightly-vs-prerelease distinction non-obvious -- still below 0.1.20. ("0.1.16.post3.dev0+g620287969.d20260725", False), From c619c28d76017c84462c1c7e9d28078d3d9363cc Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Sun, 13 Sep 2026 23:50:14 -0400 Subject: [PATCH 08/41] build(rocm): prebuild all 40 attention variants, not 36 The set was trimmed on an estimate that turned out 2-5x too slow. Measured in a real image build at --jobs 2 on 32 cores: family n measured s census est ratio mha_fwd 4 162 803 5.0x mha_varlen_fwd 16 65 140 2.2x mha_batch_prefill 16 68 206 3.0x 36 variant(s) in 23 min at --jobs 2 36/36 present in .../aiter/jit The census figures came from Step 0's S0.7, which ran two concurrent builds at MAX_JOBS=16 each on purpose. The driver sets no such cap, so ninja uses the whole box (-j 307) and every family is far cheaper than the plan assumed. That dissolves the trade-off the trim was making. Dropping mha_fwd's four LSE variants saves ~5 min of CI, once, and costs a multi-minute in-process stall to whoever first calls single_prefill_with_kv_cache(..., return_lse=True). The whole 40 lands at roughly 26 min. `--only` still narrows by family for local work. Also removes --all, which now has nothing to select. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 30 ++++++--------------- tests/rocm/test_prebuild_aiter_attention.py | 14 +++++----- 2 files changed, 14 insertions(+), 30 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 2ff22cb7fb8..4035dc5ea30 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -186,24 +186,15 @@ def reachable_variants() -> Tuple[Variant, ...]: return tuple(out) -def _is_default(v: Variant) -> bool: - """The set the image builds: everything except mha_fwd's LSE arm. - - Those four are the most expensive builds in the set (803 s each against 206 s - and 140 s), and LSE on *single* prefill serves cascade/merge rather than a - plain forward pass -- the Blaze-O1 LUM3 footprint names it unused. Skipping - them costs one ~5 min stall to anyone calling single prefill with - ``return_lse=True``; it is not an error. - """ - return not (v.family == "mha_fwd" and v.has_lse) - +def selected_variants(only: Optional[Sequence[str]] = None) -> Tuple[Variant, ...]: + """Every reachable variant. ``only`` narrows to named families. -def selected_variants( - only: Optional[Sequence[str]] = None, everything: bool = False -) -> Tuple[Variant, ...]: + Nothing is trimmed. An earlier version skipped mha_fwd's four LSE variants on + an estimate of 803 s per build; measured in the image they take 162 s, so the + whole set is 23 min at --jobs 2 and a trim saves ~5 min for a guaranteed + multi-minute stall on whoever first asks for the missing arm. + """ variants = reachable_variants() - if not everything: - variants = tuple(v for v in variants if _is_default(v)) if only: variants = tuple(v for v in variants if v.family in set(only)) return variants @@ -416,11 +407,6 @@ def main(argv: Optional[Sequence[str]] = None) -> int: p.add_argument( "--check", action="store_true", help="verify artifacts, build nothing" ) - p.add_argument( - "--all", - action="store_true", - help="build all 40 reachable variants, not just the default set", - ) p.add_argument( "--only", action="append", @@ -440,7 +426,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: build_one(_decode(a.build_one)) return 0 - variants = selected_variants(a.only, everything=a.all) + variants = selected_variants(a.only) if a.list: for v in variants: print(v.so_name) diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 510bd2dfc11..9f80871e9a0 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -88,13 +88,11 @@ def test_every_driver_filename_is_one_the_loader_asks_for(): assert {v.so_name for v in driver.reachable_variants()} == _fi_so_names() -def test_the_default_set_is_the_full_set_minus_mha_fwd_lse(): - selected = driver.selected_variants() - skipped = set(driver.reachable_variants()) - set(selected) - - assert len(selected) == 36 - assert all(v.family == "mha_fwd" and v.has_lse for v in skipped) - assert len(skipped) == 4 +def test_nothing_is_trimmed_from_the_built_set(): + """Measured in the image, the whole set is 23 min at --jobs 2, so a trim buys + ~5 min of CI for a guaranteed multi-minute stall on the missing arm.""" + assert set(driver.selected_variants()) == set(driver.reachable_variants()) + assert len(driver.selected_variants()) == 40 def test_the_blaze_lum3_variant_is_in_the_default_set(): @@ -107,7 +105,7 @@ def test_the_blaze_lum3_variant_is_in_the_default_set(): @pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) def test_only_restricts_without_renaming(family): - picked = driver.selected_variants(only=[family], everything=True) + picked = driver.selected_variants(only=[family]) assert picked assert {v.so_name for v in picked} <= _fi_so_names() assert all(v.family == family for v in picked) From fa4383ec89aebac7c681cfb6c88def3c21118f80 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 00:37:46 -0400 Subject: [PATCH 09/41] fix(rocm): re-measure the AITER soft-cap defect at 0.1.21.post2 The gate is unchanged -- gfx950 defective, gfx942 clean -- but it was recorded against 0.1.20 and the version move made that evidence stale. Re-measured on both boards, guard disabled so the kernel actually runs, non-square shapes at head_dim=128 causal, fp16, vs an fp32 reference. gfx950 (MI350X, slurm 67930310), amd-aiter 0.1.21.post2, ROCm 10.0: qo kv cap |aiter-ref| |fa2-ref| 1 1024 30 1.1930 0.0001 1 1024 50 2.6929 0.0001 17 1024 30 NaN 0.0001 37 4096 50 2.6767 0.0000 128 512 50 3.3845 0.0003 256 1024 50 3.4392 0.0002 1024 1024 50 4.1128 0.0010 12 cap>0 rows, 12 wrong (9 by magnitude, 3 returning NaN) gfx942 (MI300X), same version and sweep: 12 cap>0 rows, 0 wrong, worst 0.0003. cap=0 is clean on both boards, which is why the guard must stay narrow rather than becoming a KnownBad row over the whole op. Three things this sweep had to get right, each of which produced a wrong answer first: - The guard has to be disabled. With it on, every gfx950 capped row raises FlashInfer's own ValueError and the sweep reports "not reproduced" -- a clean result meaning "never ran". The script now exits 2 as INCONCLUSIVE when no capped row executes, so that cannot read as a pass again. - NaN output is a wrong answer, not a skipped row. Scoring it as "raised" understated the defect by three rows; nothing threw in either run. - qo_len > kv_len is dropped: causal with more queries than keys is ill-posed and *both* backends fail, which says nothing about AITER. _AITER_SOFTCAP_DEFECT_THROUGH moves to 0.1.21.post2. It is message text only; the gate is arch_caps.aiter_softcap_defect_arch, which this leaves alone. Co-Authored-By: Claude --- flashinfer/rocm/arch_caps.py | 7 ++++++- flashinfer/rocm/prefill.py | 5 +++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/flashinfer/rocm/arch_caps.py b/flashinfer/rocm/arch_caps.py index e7bb228a2f6..c1e1a281854 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -271,7 +271,12 @@ def __post_init__(self) -> None: # Not a KnownBad row: those gate a whole (op, backend, arch) on toolchain # version and would also disable the clean logits_soft_cap=0 path. What varies # by architecture is whether the kernel is affected, not at which length. -# Measured on amd-aiter 0.1.20 over qo_len x kv_len x cap, vs an fp32 reference. +# +# Re-measured on amd-aiter 0.1.21.post2, 2026-09-14, over non-square +# qo_len x kv_len x cap at head_dim=128 causal, vs an fp32 reference, with this +# guard disabled so the kernel runs. Unchanged: gfx950 wrong on all 12 capped +# rows (9 at 1.19-4.11 absolute, 3 returning NaN) where fa2 is 0.0001-0.0011 on +# the same shapes; gfx942 clean on all 12. cap=0 is clean on both. _AITER_SOFTCAP_DEFECT_ARCHS = {"gfx942": False, "gfx950": True} diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index af3b3b77cae..5bb5d670ce2 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -66,8 +66,9 @@ _AITER_NATIVE_PAGING_SINCE = "0.1.10" _AITER_LAST_VALIDATED = "0.1.21.post2" # Newest AITER carrying the mha_varlen_fwd soft-cap defect. Bump only after -# re-measuring against an fp32 reference; the wrong answer is silent. -_AITER_SOFTCAP_DEFECT_THROUGH = "0.1.21" +# re-measuring against an fp32 reference; the wrong answer is silent. Message +# text only -- the gate itself is arch_caps.aiter_softcap_defect_arch. +_AITER_SOFTCAP_DEFECT_THROUGH = "0.1.21.post2" # fp8 query dtypes that *could* be an fp8 prefill: E4M3FNUZ on gfx942, OCP # E4M3FN on gfx950. Only the arch's own encoding actually works -- the other is From 2fb6b08edd820fffd7a3c8737ea13f8d9d515ca7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 06:32:39 -0400 Subject: [PATCH 10/41] docs: the supported AITER install is a pinned source build The docs told users to pip-install a wheel and warned that a source build was untested. The image now does the opposite, so every one of those passages is not merely stale but actively wrong. Rewritten rather than renumbered, because the reasons changed too: - The wheel recipe becomes the tagged source build, with the three details that actually bite: build by tag (a lost tag makes setuptools_scm invent a .dev version that falls under the floor and silently degrades every AITER path to fa2), non-editable (develop ships no aiter_meta), and PREBUILD_KERNELS=0 (=1 needs a GPU, and its modules are no loss). - The cp312 rationale is gone. It held because every 0.1.20 wheel was cp312 only; a source build follows whichever interpreter is present, so the interpreter is now a property of the base image. - The "+rocm10.1.0a is the nearest retarget" apology is replaced by the headline correctness win: a source build compiles against the ROCm in the image rather than a retarget of another one. - CLAUDE.md's "a source build tracks master, treat as untested here" is reversed. Its neighbouring "That gap is a trap" sentence lost its antecedent with that paragraph, so it is reworded to stand alone -- the advice to read the installed tree rather than a checkout is still right. - CONTRIBUTING.md documents AITER_REF, AITER_PREBUILD_JOBS and AITER_SOURCE=0 in place of the wheel args. Left alone deliberately: backends.md:435's "measured exact on amd-aiter 0.1.20" is a record of when a measurement happened, not a pin. The notebook edit is a two-line text change; cell and output counts are unchanged (17 and 17) after the JSON round-trip. Co-Authored-By: Claude --- CLAUDE.md | 39 +++++------ CONTRIBUTING.md | 7 +- README.md | 10 +-- docs/rocm/backends.md | 73 +++++++++++---------- examples/amd_flashinfer_rocm_tutorial.ipynb | 4 +- 5 files changed, 64 insertions(+), 69 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index ad79cee4f25..3cc37b94d1e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -111,29 +111,22 @@ marking a test retires it, so do not reach for it to speed a lane up. gfx950 at `-n 1`, so 92% of a first run is compilation. Clearing `~/.cache/flashinfer` before timing anything makes the result meaningless. -**AITER is version-pinned.** The development image bundles the wheel, so no -separate install is needed there. `aiter_utils.AITER_MIN_VERSION` (0.1.20) is -the hard floor below which the vendored struct layouts stop matching, and -every 0.1.20 build is cp312 only — which is what fixes the interpreter. - -```bash -pip install amd-aiter==0.1.20+rocm10.1.0a20260819.3135022 \ - --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/ -``` - -Spell the version out in full including the local `+rocm...` segment; pip will -not select a local version from a loose specifier. `prefill.py` records -whatever is validated as `_AITER_LAST_VALIDATED`, and -[`docs/rocm/backends.md`](docs/rocm/backends.md) explains the index choice. - -A source build (`git clone --recursive https://github.com/ROCm/aiter.git && -cd aiter && python3 setup.py develop`) tracks master, which is **many releases -ahead of the pin with a different C ABI** — symbols the shim expects are -renamed, hidden rather than `extern "C"`, or absent. Nothing stops you running -one, but treat it as untested here. - -That gap is a trap when working on the C++ shim: read the **installed** tree, -never a source checkout, before designing against an AITER symbol. +**AITER is a source build at a pinned tag.** `docker/Dockerfile.rocm` clones +`v0.1.21.post2` and installs it non-editable with `PREBUILD_KERNELS=0`, so the +image needs no separate install. `aiter_utils.AITER_MIN_VERSION` (0.1.21) is +the hard floor: below it the vendored struct layouts stop matching, and 0.1.21 +moved `rmsnorm` and `rope` from `at::Tensor` to the POD `aiter_tensor_t`. + +Build it by **tag**, never a branch — a clone that loses its tag makes +`setuptools_scm` invent a `.dev` version that falls under the floor, and every +AITER path then silently degrades to `fa2`. Use a non-editable install: +`setup.py develop` ships no `aiter_meta`, which is what the C++ shims compile +against. `prefill.py` records whatever is validated as `_AITER_LAST_VALIDATED`, +and [`docs/rocm/backends.md`](docs/rocm/backends.md) has the full recipe. + +When working on the C++ shim, read the **installed** tree before designing +against an AITER symbol — never a source checkout, which is generally at a +different revision with a different C ABI. ```bash nm -D --defined-only -C /aiter/jit/module_.so | grep '(' diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 21092436714..0854e7d8b70 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -22,9 +22,10 @@ docker build -t flashinfer-dev:rocm10.0 -f docker/Dockerfile.rocm . `ROCM_VERSION`, `UBUNTU_VERSION`, `PY_VERSION`, and `TORCH_VERSION` default to 10.0, 24.04, 3.12, and 2.12.0. They select the `rocm/pytorch` base image tag, so they are not independent knobs — any override has to name a tag that exists -on Docker Hub. `AITER_VERSION` and `AITER_INDEX` pin the AITER wheel; every -0.1.20 build is cp312 only, so the interpreter and the AITER pin move together. -Do not raise `TORCH_VERSION` to 2.13 — it drops a `c10` symbol AITER's prebuilt +on Docker Hub. `AITER_REF` selects the AITER tag to build from source, and +`AITER_PREBUILD_JOBS` how many variants compile at once (~23 min at 2 on 32 +cores). `AITER_SOURCE=0` falls back to the `AITER_VERSION`/`AITER_INDEX` wheel +install. Do not raise `TORCH_VERSION` to 2.13 — it drops a `c10` symbol AITER's prefill kernels link against, and they fail to load. Pass `--build-arg USERNAME=$USER --build-arg USER_UID=$(id -u) --build-arg USER_GID=$(id -g)` to match container file ownership to your host user — diff --git a/README.md b/README.md index 37aeb4ba9cd..6dea64b8b1b 100644 --- a/README.md +++ b/README.md @@ -100,17 +100,17 @@ and this release is tested on.** | PyTorch+ROCm | 2.12.0 | | Python | 3.12 | | OS | Ubuntu 24.04 | -| `amd-aiter` | 0.1.20 | +| `amd-aiter` | 0.1.21.post2 (source build) | Nothing rejects another combination at install time, and older ROCm and torch releases have worked here before — but they are untested, uncovered by the matrix below, and not what a bug report will be reproduced against. The pins move together, which is why the supported configuration is an image -rather than a list of versions: every `amd-aiter` 0.1.20 wheel is cp312 only, -fixing the interpreter; torch must stay at 2.12, since 2.13 drops a `c10` -symbol those wheels' prebuilt prefill kernels need; and its ROCm 10.0 build -exists only in the base image. +rather than a list of versions: `amd-aiter` is built from source against the +image's own ROCm, since no wheel targets ROCm 10.0; torch must stay at 2.12, +since 2.13 drops a `c10` symbol AITER's prefill kernels need; and its ROCm 10.0 +build exists only in the base image. ## Support matrix diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index be2ccf32025..19efc87a762 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -147,49 +147,50 @@ That says nothing about their ROCm support either way. ## Installing AITER -Unless you are using the prebuilt Docker image, AITER is a separate -install. `amd-aiter` is **not** on the top-level `pypi.amd.com/simple` -index, and the ROCm-versioned channels carry no wheel at or above the -supported floor, so install from the nightlies index: +**The supported install is a source build at a pinned tag**, which is what +`docker/Dockerfile.rocm` performs: ```bash -pip install amd-aiter==0.1.20+rocm10.1.0a20260819.3135022 \ - --extra-index-url https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/ +git clone --recursive --depth 1 --shallow-submodules \ + --branch v0.1.21.post2 https://github.com/ROCm/aiter.git +cd aiter && PREBUILD_KERNELS=0 GPU_ARCHS="gfx942;gfx950" \ + pip install --no-build-isolation . ``` -Use `--extra-index-url`, not `--index-url`, so AITER's own dependencies -still resolve from PyPI, and spell the version out in full including the -local `+rocm...` segment — pip will not select a local version from a loose -specifier. Check afterwards that the install did not pull a CPU-only torch -over your ROCm one: `python -c "import torch; assert torch.version.hip"`. +Source rather than a wheel because no `amd-aiter` wheel is built against +ROCm 10.0; the published ones are retargets of the same revision. A source +build compiles against the ROCm actually present, and it is the only route to +a tagged release on this stack. -**`aiter_utils.AITER_MIN_VERSION` (0.1.20) is a hard floor**, enforced -before routing. FlashInfer links AITER's C++ symbols by mangled name -(`csrc/rocm/aiter_loader.cc`) and vendors its argument structs -(`include/flashinfer/rocm/attention/aiter/`) at the 0.1.20 layout, so an older -release shifts field offsets instead of failing to load -- and 0.1.20 renamed -the RMSNorm entry points, so an older one cannot resolve them. Below the floor -`auto` will not select AITER and an explicit `backend="aiter"` raises. - -The development image (`docker/Dockerfile.rocm`) bundles that wheel, on -CPython 3.12, and needs no separate install -- it is the one supported -configuration, so in practice this section is only for someone assembling -their own. +Three details are load-bearing. Use the **tag**, not a branch: a clone that +loses its tag makes `setuptools_scm` invent a `.dev` version that falls under +the floor below, and every AITER path silently degrades to `fa2`. Non-editable, +not `setup.py develop`: the develop branch installs no `aiter_meta`, and +`aiter_meta/csrc/include` is what every C++ shim compiles against. +`PREBUILD_KERNELS=0` because `=1` needs a live GPU — and its prebuilt modules +are no loss, since FlashInfer builds its own `lib.so` from the same +sources. -Every 0.1.20 wheel is cp312 only, which is what fixes the interpreter. None -is built against ROCm 10.0 directly; the pin above is the nearest retarget -of the same source revision (build id `3135022`). +Check afterwards that the install did not pull a CPU-only torch over your ROCm +one: `python -c "import torch; assert torch.version.hip"`. -A source build tracks master, which is many releases ahead of the pin -**with a different C ABI** — symbols the shim expects are renamed, hidden -rather than `extern "C"`, or absent: - -```bash -git clone --recursive https://github.com/ROCm/aiter.git -cd aiter && python3 setup.py develop -``` - -Nothing stops you running one, but treat it as untested here. +**`aiter_utils.AITER_MIN_VERSION` (0.1.21) is a hard floor**, enforced +before routing. FlashInfer links AITER's C++ symbols by mangled name +(`csrc/rocm/aiter_loader.cc`) and vendors its argument structs +(`include/flashinfer/rocm/attention/aiter/`), so an older release shifts field +offsets instead of failing to load. 0.1.21 also moved `rmsnorm` and the +cos/sin-cache `rope` from `at::Tensor` to the POD `aiter_tensor_t`, so those +shims cannot resolve against 0.1.20 at all. Below the floor `auto` will not +select AITER and an explicit `backend="aiter"` raises. + +The development image needs none of this: it performs the build above and +prebuilds the attention variants, so a container starts with them present. It +is the one supported configuration, so in practice this section is only for +someone assembling their own. + +The interpreter is no longer pinned by the AITER wheel's tags — a source build +follows whichever Python is present. CPython 3.12 is now a property of the base +image rather than of AITER. ### C++-level integration diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 56a09a088fb..824b1ee5da2 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -31,7 +31,7 @@ "source": [ "## Setting up amd-flashinfer\n", "\n", - "The [project README](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/README.md) describes supported ROCm and PyTorch versions and **Instinct** architectures (**gfx942**, **gfx950**). **One configuration is supported**: ROCm 10.0, Torch+ROCm 2.12.0, CPython 3.12, Ubuntu 24.04 and `amd-aiter` 0.1.20 — what `docker/Dockerfile.rocm` builds. Other combinations are untested. Align your environment with one of the following before running the cells below.\n", + "The [project README](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/README.md) describes supported ROCm and PyTorch versions and **Instinct** architectures (**gfx942**, **gfx950**). **One configuration is supported**: ROCm 10.0, Torch+ROCm 2.12.0, CPython 3.12, Ubuntu 24.04 and `amd-aiter` 0.1.21.post2 (built from source) — what `docker/Dockerfile.rocm` builds. Other combinations are untested. Align your environment with one of the following before running the cells below.\n", "\n", "**Option 1 — The development image (recommended for reproducibility)** \n", "There is no published wheel or image for this release; build the development image from the repository, which carries a matched ROCm, PyTorch, Python and AITER set:\n", @@ -164,7 +164,7 @@ "\n", "**Backend integration.** amd-flashinfer offers [AITER](https://github.com/ROCm/aiter)—AMD’s SOTA kernel library for large language model operators on ROCm—as an optional **`backend=\"aiter\"`**. When the `aiter` package is present and the requested operator is integrated with AITER in this port, work **delegates** to AITER; otherwise the default HIP path applies. The following cells exercise the AITER-backed prefill path.\n", "\n", - "**Requirements and layout.** Install the `amd-aiter` package from the nightlies index — it is not on the top-level `pypi.amd.com/simple`, and 0.1.20 is a hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, AITER natively serves page sizes `{128, 256, 1024}`. Other sizes still work, via an extra gather." + "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, AITER natively serves page sizes `{128, 256, 1024}`. Other sizes still work, via an extra gather." ] }, { From ad88c14c8fd54b01b8a7a7ee12129d8267b03f7e Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 06:47:19 -0400 Subject: [PATCH 11/41] fix(rocm): guard the stream on the migrated norm and rope shims Self-review (/code-review xhigh) before push. One finding was a regression this branch introduced; the rest are its consequences. **The stream guard.** csrc/rocm/aiter_tensor_compat.h says it outright: "The POD entry points launch on aiter::getCurrentHIPStream(), which defaults to nullptr and is otherwise set only by AITER's Python layer." All three shims that migrated at 0.1.16 install a StreamGuard; the two this branch migrated did not: activation_aiter.cu 1 fused_moe_aiter.cu 1 page_aiter.cu 1 norm_aiter.cu 0 <- added here rope_aiter.cu 0 <- added here Under `with torch.cuda.stream(s)` -- which is how vLLM and SGLang drive a rank -- the kernel went to the null stream while `at::empty_like` and the `copy_` back ran on `s`: a silent race, and an abort under graph capture. The at::Tensor entry points picked the torch stream up themselves, so nothing caught it. **AITER_SOURCE=0 installed a wheel below the floor this branch raised.** The build stayed green and every AITER path silently stopped routing -- the exact failure the surrounding assertions exist to prevent. The assertion block now rejects any resolved version under AITER_MIN_VERSION, wheel arm included. **The csrc-include fallback was unreachable.** `_aiter_importable()` still did a hard `import aiter_meta`, so on the editable layout the fallback exists to serve, routing gave up long before the resolver ran. It now probes through `_aiter_csrc_include_dir()` instead, which is the thing that actually has to work. Smaller, same review: `AITER_MIN_VERSION`'s comment still argued for 0.1.20; kAbiPinNote had lost its actionable command and now names backends.md; the benchmark-kernel skill still stated the old floor on the line that was edited for this pin; the tutorial still promised native paged page sizes that fail the probe at 0.1.21.post2; `_run_jobs` leaked children on the way out, which matters because a zombie-held AITER build lock is never treated as stale; two driver docstrings carried measurements that belong in a commit message. Dropped `test_the_blaze_lum3_variant_is_in_the_default_set`: the proper noun it cites appears nowhere in-tree, and `test_nothing_is_trimmed_from_the_built_set` already covers the assertion. 128 tests pass after the shim rebuild the StreamGuard forces. Co-Authored-By: Claude --- .claude/skills/benchmark-kernel/SKILL.md | 2 +- csrc/rocm/aiter_loader.cc | 4 +-- csrc/rocm/norm_aiter.cu | 2 ++ csrc/rocm/rope_aiter.cu | 1 + docker/Dockerfile.rocm | 9 ++++++ docker/prebuild_aiter_attention.py | 32 +++++++++++---------- examples/amd_flashinfer_rocm_tutorial.ipynb | 2 +- flashinfer/rocm/aiter_utils.py | 15 ++++++---- tests/rocm/test_prebuild_aiter_attention.py | 8 ------ 9 files changed, 43 insertions(+), 32 deletions(-) diff --git a/.claude/skills/benchmark-kernel/SKILL.md b/.claude/skills/benchmark-kernel/SKILL.md index 55c10a5ce4f..d5ac867752a 100644 --- a/.claude/skills/benchmark-kernel/SKILL.md +++ b/.claude/skills/benchmark-kernel/SKILL.md @@ -25,7 +25,7 @@ For the in-repo profiler wrapper, see [`profiler/rocm/rocm_profiler.py`](../../. - Explicit `backend="aiter"` + `kv_layout != "NHD"` → `ValueError`. Grep [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for `only supports kv_layout`; single prefill raises at the call, both batch wrappers at `plan()`. Not raised by auto-selection — that path silently falls back to `fa2`. - Explicit `backend="aiter"` on non-gfx942/gfx950 → `RuntimeError`. - `amd-aiter` not importable → `ImportError`. - - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py). The `{16, 1024}` arm there is unreachable now that `flashinfer.rocm.aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.20. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. **Measured on amd-aiter 0.1.21.post2, all three advertised sizes fail the probe** (`no matching kernel found`) on gfx942 *and* gfx950, so every paged prefill flat-gathers and `mha_batch_prefill` is unreachable — benchmark the paged path as flat-gather until that is fixed. + - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py). The `{16, 1024}` arm there is unreachable now that `flashinfer.rocm.aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.21. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. **Measured on amd-aiter 0.1.21.post2, all three advertised sizes fail the probe** (`no matching kernel found`) on gfx942 *and* gfx950, so every paged prefill flat-gathers and `mha_batch_prefill` is unreachable — benchmark the paged path as flat-gather until that is fixed. - Auto-selection (no explicit `backend=`) silently falls back to `fa2` for any of: `kv_layout != "NHD"`, custom mask, dtype not in `{fp16, bf16}`, `dtype_q != dtype_kv`, `head_dim_qk != head_dim_vo`, `pos_encoding_mode != "NONE"`, or `amd-aiter` not importable. See `_auto_select_prefill_backend()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for the authoritative list; it returns `(backend, reason)`, and the reason names the constraint that forced the fallback. - **Always verify numerical parity before trusting perf numbers.** Compare default-HIP vs AITER outputs with `torch.testing.assert_close(rtol=1e-2, atol=1e-2)` for BF16/FP16 first. - **`gcnArchName` is the unambiguous arch marker.** Device strings show `cuda:0` on AMD too. Record `torch.cuda.get_device_properties(0).gcnArchName` and `torch.version.hip` alongside every number — a `gfx942` / ROCm 7.2 result is not comparable to a `gfx950` / ROCm 7.0.2 result. diff --git a/csrc/rocm/aiter_loader.cc b/csrc/rocm/aiter_loader.cc index 226ec578ba4..50a78933f15 100644 --- a/csrc/rocm/aiter_loader.cc +++ b/csrc/rocm/aiter_loader.cc @@ -26,8 +26,8 @@ constexpr const char* kAbiPinNote = "\n These symbols and .so names are pinned to amd-aiter 0.1.21.post2, which has no" "\n stable C++ ABI -- 0.1.21 moved rmsnorm and rope from at::Tensor to the POD" "\n aiter_tensor_t. If AITER was changed, that is the likely cause. The supported" - "\n install is the source build docker/Dockerfile.rocm performs at tag v0.1.21.post2," - "\n or re-pin the symbols in csrc/rocm/aiter_loader.cc."; + "\n install is a source build at tag v0.1.21.post2; docs/rocm/backends.md has the" + "\n exact clone and pip command, or re-pin the symbols in csrc/rocm/aiter_loader.cc."; // Directories to try, in order, when resolving a variant .so. // diff --git a/csrc/rocm/norm_aiter.cu b/csrc/rocm/norm_aiter.cu index e30edce6f3f..2ee8bd49c16 100644 --- a/csrc/rocm/norm_aiter.cu +++ b/csrc/rocm/norm_aiter.cu @@ -52,6 +52,7 @@ void fused_add_rmsnorm_aiter(at::Tensor input, at::Tensor residual, at::Tensor w auto a_residual = flashinfer::aiter_compat::to_aiter(residual); auto a_residual_out = flashinfer::aiter_compat::to_aiter(residual_out); auto a_weight = flashinfer::aiter_compat::to_aiter(weight2d); + const flashinfer::aiter_compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); aiter::add_rmsnorm(a_out, a_input, a_residual, a_residual_out, a_weight, eps, /*gemma_norm=*/false); input.copy_(out); @@ -67,6 +68,7 @@ void rmsnorm_aiter(at::Tensor out, at::Tensor input, at::Tensor weight, double e // same way; stage it only when the caller actually aliased. auto a_input = flashinfer::aiter_compat::to_aiter(input); auto a_weight = flashinfer::aiter_compat::to_aiter(weight2d); + const flashinfer::aiter_compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); if (out.data_ptr() == input.data_ptr()) { at::Tensor staged = at::empty_like(out); auto a_staged = flashinfer::aiter_compat::to_aiter(staged); diff --git a/csrc/rocm/rope_aiter.cu b/csrc/rocm/rope_aiter.cu index 4ff842fe0df..dd50bad342f 100644 --- a/csrc/rocm/rope_aiter.cu +++ b/csrc/rocm/rope_aiter.cu @@ -85,6 +85,7 @@ void apply_rope_pos_ids_cos_sin_cache_aiter(at::Tensor query, at::Tensor key, at namespace compat = flashinfer::aiter_compat; auto a_q_out = compat::to_aiter(q_out_rot); auto a_k_out = compat::to_aiter(k_out_rot); + const compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); rope_cached_positions_2c_fwd_impl(a_q_out, a_k_out, compat::to_aiter(q_in_rot), compat::to_aiter(k_in_rot), compat::to_aiter(cos), compat::to_aiter(sin), compat::to_aiter(pos), diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7feadd13696..351844125e1 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -210,6 +210,15 @@ want = want if os.environ["AITER_SOURCE"] == "1" else version if got != want: sys.exit(f"ERROR: amd-aiter resolved to {got!r}, expected {want!r}") +# The AITER_SOURCE=0 wheel is pinned separately and can fall under the floor the +# shims need, which no runtime check would catch loudly: every AITER path just +# stops routing and the image is quietly several times slower on prefill. +FLOOR = "0.1.21" +from packaging.version import Version # noqa: E402 + +if Version(got).base_version < Version(FLOOR).base_version: + sys.exit(f"ERROR: amd-aiter {got} is below AITER_MIN_VERSION {FLOOR}") + inc = Path(md.distribution("amd-aiter").locate_file("aiter_meta/csrc/include")) if not inc.is_dir(): sys.exit(f"ERROR: {inc} is missing; the shims cannot compile against this install") diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 4035dc5ea30..99d7bf62378 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -187,12 +187,9 @@ def reachable_variants() -> Tuple[Variant, ...]: def selected_variants(only: Optional[Sequence[str]] = None) -> Tuple[Variant, ...]: - """Every reachable variant. ``only`` narrows to named families. + """Every reachable variant; ``only`` narrows to named families. - Nothing is trimmed. An earlier version skipped mha_fwd's four LSE variants on - an estimate of 803 s per build; measured in the image they take 162 s, so the - whole set is 23 min at --jobs 2 and a trim saves ~5 min for a guaranteed - multi-minute stall on whoever first asks for the missing arm. + Nothing is trimmed: the whole set costs minutes, a missing arm costs a stall. """ variants = reachable_variants() if only: @@ -267,12 +264,10 @@ def jit_dir(core) -> Path: def _args_of_build(core, family: str) -> Dict[str, object]: - """AITER's compile flags for a family, keyed as `compile_ops` keys them. + """AITER's compile flags for a family, keyed as ``compile_ops`` keys them. - The config key carries a ``module_`` prefix; without it ``get_args_of_build`` - *warns and returns empty args* rather than raising, which drops every include - path and fails as a missing header several minutes into the build. Assert - rather than trust, so a renamed key is loud. + The ``module_`` prefix matters: without it ``get_args_of_build`` warns and + returns empty args rather than raising, so assert rather than trust. """ args = core.get_args_of_build(f"module_{family}") if not args.get("srcs") or not args.get("extra_include"): @@ -321,9 +316,20 @@ def _run_jobs(variants: Sequence[Variant], jobs: int) -> int: # would spin forever. Only read on failure, and only the tail. running: List[Tuple[Variant, subprocess.Popen, float, str]] = [] failed: List[str] = [] - done = 0 total = len(pending) + try: + _drain(pending, running, failed, jobs, total) + finally: + for _, proc, _, _ in running: + proc.kill() + if failed: + print(f"\n{len(failed)} build(s) failed:\n " + "\n ".join(failed)) + return 1 if failed else 0 + + +def _drain(pending, running, failed, jobs, total) -> None: + done = 0 while pending or running: while pending and len(running) < jobs: v = pending.pop(0) @@ -361,10 +367,6 @@ def _run_jobs(variants: Sequence[Variant], jobs: int) -> int: flush=True, ) - if failed: - print(f"\n{len(failed)} build(s) failed:\n " + "\n ".join(failed)) - return 1 if failed else 0 - def _encode(v: Variant) -> str: return "|".join( diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 824b1ee5da2..e344a060644 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -164,7 +164,7 @@ "\n", "**Backend integration.** amd-flashinfer offers [AITER](https://github.com/ROCm/aiter)—AMD’s SOTA kernel library for large language model operators on ROCm—as an optional **`backend=\"aiter\"`**. When the `aiter` package is present and the requested operator is integrated with AITER in this port, work **delegates** to AITER; otherwise the default HIP path applies. The following cells exercise the AITER-backed prefill path.\n", "\n", - "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, AITER natively serves page sizes `{128, 256, 1024}`. Other sizes still work, via an extra gather." + "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the probe fails for every advertised size on both architectures, so this path always gathers." ] }, { diff --git a/flashinfer/rocm/aiter_utils.py b/flashinfer/rocm/aiter_utils.py index 4ccf77827a0..76572af0c35 100644 --- a/flashinfer/rocm/aiter_utils.py +++ b/flashinfer/rocm/aiter_utils.py @@ -50,10 +50,9 @@ def _ensure_aiter_gpu_archs() -> None: os.environ["GPU_ARCHS"] = arch -# 0.1.20 renamed the RMSNorm entry points and added a gemma_norm parameter, so -# an older AITER cannot resolve the shim's symbols at all. The vendored prefill -# structs travel by value through dlsym'd pointers, where a mismatch corrupts -# silently -- hence a hard floor rather than a warning. +# 0.1.21 moved rmsnorm and the cos/sin-cache rope from at::Tensor to the POD +# aiter_tensor_t, so csrc/rocm/{norm,rope}_aiter.cu cannot resolve their mangled +# symbols against anything older, and the vendored arg structs assume this layout. AITER_MIN_VERSION = "0.1.21" @@ -98,8 +97,14 @@ def _aiter_importable() -> bool: try: _ensure_aiter_gpu_archs() import aiter # noqa: F401 - import aiter_meta # noqa: F401 from aiter.jit import core as _core # noqa: F401 + + # Not `import aiter_meta`: an editable AITER ships none, and + # jit.rocm.aiter_source._aiter_csrc_include_dir now falls back to + # AITER_CSRC_DIR. Requiring it here made that fallback unreachable. + from ..jit.rocm.aiter_source import _aiter_csrc_include_dir + + _aiter_csrc_include_dir() except Exception: return False return _aiter_version_supported() diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 9f80871e9a0..353bfc90c0b 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -95,14 +95,6 @@ def test_nothing_is_trimmed_from_the_built_set(): assert len(driver.selected_variants()) == 40 -def test_the_blaze_lum3_variant_is_in_the_default_set(): - """Blaze-O1's LUM3 pins fp16 / batch mode / causal / no soft-cap / no LSE; - batch mode is mha_fwd, so this exact file is the one it dlopens.""" - assert "mha_fwd_fp16_nbias_mask_nlse_ndropout_nqscale.so" in { - v.so_name for v in driver.selected_variants() - } - - @pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) def test_only_restricts_without_renaming(family): picked = driver.selected_variants(only=[family]) From 2b8527c2899ad5c98a7aa8cb9b968d67fffe5fbe Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 07:09:11 -0400 Subject: [PATCH 12/41] test(rocm): cover AITER shim stream affinity The StreamGuard regression fixed in the previous commit passed all 69 norm and rope AITER tests while it was live. Every one of them runs on the default stream, where the null stream AITER falls back to and the caller's stream are the same stream, so the race cannot appear. A reviewer found it by reading. These tests run the same ops inside `with torch.cuda.stream(s)` with real work queued ahead, which is how vLLM and SGLang drive a rank. A/B, guards stripped from both shims and the shim cache cleared: ..F.F 2 failed, 3 passed FAILED test_rmsnorm_aiter_is_correct_on_a_side_stream[4096] FAILED test_rope_aiter_is_correct_on_a_side_stream Restored, 5 passed. Only the larger shapes lose the race -- hence hidden=4096 in the parametrize and the matmul ballast, without which nothing fails. A race test catches a subset by nature; two of five is enough to stop this returning silently, which is the whole problem with it. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 5 ++ tests/rocm/test_aiter_stream_affinity.py | 97 ++++++++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 tests/rocm/test_aiter_stream_affinity.py diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 99d7bf62378..a4c358d458b 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -36,6 +36,7 @@ import argparse import os +import shutil import subprocess import sys import tempfile @@ -300,6 +301,10 @@ def build_one(v: Variant) -> Path: out = jit_dir(core) / v.so_name if not out.is_file(): raise RuntimeError(f"build of {v.md_name} produced no {out}") + # Drop the CK blob and object tree: ~15 MB per variant, 603 MB over the set, + # and nothing reads it again -- AITER decides "already built" from the .so. + # Kept on failure, where it is the only diagnostic. + shutil.rmtree(jit_dir(core) / "build" / v.md_name, ignore_errors=True) return out diff --git a/tests/rocm/test_aiter_stream_affinity.py b/tests/rocm/test_aiter_stream_affinity.py new file mode 100644 index 00000000000..53d975fb8e0 --- /dev/null +++ b/tests/rocm/test_aiter_stream_affinity.py @@ -0,0 +1,97 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +"""AITER POD-API shims must launch on the caller's HIP stream. + +The POD entry points read `aiter::getCurrentHIPStream()`, a thread-local that +defaults to nullptr and is set only by AITER's Python layer. A shim that omits +`aiter_compat::StreamGuard` enqueues its kernel on the null stream while the +surrounding `at::empty_like` and `copy_` run on the caller's -- a silent race +that every existing test misses, because they all run on the default stream +where the two happen to coincide. + +vLLM and SGLang both drive a rank inside `with torch.cuda.stream(s)`, so this is +the production configuration, not an exotic one. +""" + +import pytest +import torch + +import flashinfer +from flashinfer.rocm.aiter_utils import is_aiter_supported + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_aiter_supported(torch.device("cuda:0")), + reason="needs a GPU with a supported AITER", +) + +DT = torch.float16 + + +def _on_side_stream(fn): + """Run `fn` on a non-default stream with real work queued ahead of it. + + The preceding work matters: it is what makes a null-stream launch observably + early rather than merely differently ordered. + """ + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + ballast = torch.randn(4096, 4096, device="cuda", dtype=DT) + for _ in range(8): + ballast = ballast @ ballast.T / 64.0 + out = fn() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + return out + + +@pytest.mark.parametrize("hidden", [128, 1024, 4096]) +def test_rmsnorm_aiter_is_correct_on_a_side_stream(hidden): + x = torch.randn(256, hidden, device="cuda", dtype=DT) + w = torch.randn(hidden, device="cuda", dtype=DT) + expected = flashinfer.rmsnorm(x, w, backend="aiter") + + got = _on_side_stream(lambda: flashinfer.rmsnorm(x, w, backend="aiter")) + + torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) + + +def test_fused_add_rmsnorm_aiter_is_correct_on_a_side_stream(): + """The in-place path is the exposed one: it stages into a fresh buffer and + copies back, so a kernel on the wrong stream can lose the race to the copy.""" + hidden = 1024 + x = torch.randn(256, hidden, device="cuda", dtype=DT) + res = torch.randn(256, hidden, device="cuda", dtype=DT) + w = torch.randn(hidden, device="cuda", dtype=DT) + + x_ref, res_ref = x.clone(), res.clone() + flashinfer.fused_add_rmsnorm(x_ref, res_ref, w, backend="aiter") + + x_got, res_got = x.clone(), res.clone() + _on_side_stream( + lambda: flashinfer.fused_add_rmsnorm(x_got, res_got, w, backend="aiter") + ) + + torch.testing.assert_close(x_got, x_ref, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(res_got, res_ref, rtol=1e-2, atol=1e-2) + + +def test_rope_aiter_is_correct_on_a_side_stream(): + nnz, heads, head_dim = 512, 8, 128 + q = torch.randn(nnz, heads * head_dim, device="cuda", dtype=DT) + k = torch.randn(nnz, heads * head_dim, device="cuda", dtype=DT) + pos = torch.arange(nnz, device="cuda", dtype=torch.int32) + cache = torch.randn(nnz, head_dim, device="cuda", dtype=torch.float32) + + def call(): + return flashinfer.apply_rope_with_cos_sin_cache( + pos, q.clone(), k.clone(), head_dim, cache, is_neox=True, backend="aiter" + ) + + expected = call() + got = _on_side_stream(call) + + for a, b in zip(got, expected, strict=True): + torch.testing.assert_close(a, b, rtol=1e-2, atol=1e-2) From bb3e6e550c11e26c17d611cf30488c1599546835 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 07:46:13 -0400 Subject: [PATCH 13/41] fix(rocm): drop the variant build trees, and make the floor check real Three defects found while verifying the image, each by running it rather than reading it. **603 MB of dead weight shipped.** The 40 variants are 224 MB; AITER's jit/build/ staging -- CK blob sources and objects, ~15 MB per variant -- was another 603 MB, three quarters of the directory. Proven removable rather than assumed: with build/ deleted, a paged prefill plan() built 0 artifacts, because AITER decides "already built" from the .so. The driver now removes each tree on success and keeps it on failure, where it is the only diagnostic. **The floor check compared strings.** `Version(got).base_version < Version(FLOOR).base_version` is lexicographic, so it would have admitted 0.1.9 against a 0.1.10 floor. Comparing Version objects is both correct and simpler; PEP 440 already puts 0.1.20+rocm10.1.0a below 0.1.21 and 0.1.21.dev0 below 0.1.21, which is what the floor wants. **The wheel fallback could not build at all.** With the floor enforced, AITER_SOURCE=0 correctly refused its 0.1.20 pin -- leaving a documented escape hatch that always fails. The index does publish 0.1.21.post2, but only +rocm7.14.0 and +rocm7.2.3: no ROCm 10 build exists at this revision, which is the clearest statement yet of why the default path builds from source. The arm now pins the 7.14 retarget and CONTRIBUTING says plainly that it is an unblock. Also regenerates tmp/Dockerfile.c5probe from docker/Dockerfile.rocm instead of hand-copying it. The hand-copied version silently lacked the floor assertion, so the first AITER_SOURCE=0 run "passed" a check the real file would have failed -- a test that proved nothing, which is worse than no test. Cold-start A/B on the built image, fresh container, same FlashInfer JIT cost in both arms so the delta is AITER's: family prebuilt cold saved mha_fwd (single prefill) 52.5s 485.9s 433s mha_varlen_fwd (soft cap) 41.6s 156.7s 115s mha_varlen_fwd (ragged) 39.7s 177.5s 138s mha_batch_prefill (paged) 0.4s 129.3s 129s total 134.2s 949.4s 815s (7.1x) Co-Authored-By: Claude --- CONTRIBUTING.md | 5 +++-- docker/Dockerfile.rocm | 14 +++++++------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 0854e7d8b70..cdf41dc9fc3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -24,8 +24,9 @@ docker build -t flashinfer-dev:rocm10.0 -f docker/Dockerfile.rocm . so they are not independent knobs — any override has to name a tag that exists on Docker Hub. `AITER_REF` selects the AITER tag to build from source, and `AITER_PREBUILD_JOBS` how many variants compile at once (~23 min at 2 on 32 -cores). `AITER_SOURCE=0` falls back to the `AITER_VERSION`/`AITER_INDEX` wheel -install. Do not raise `TORCH_VERSION` to 2.13 — it drops a `c10` symbol AITER's +cores). `AITER_SOURCE=0` falls back to the `AITER_VERSION`/`AITER_INDEX` wheel install, +which is an unblock only: no `amd-aiter` wheel targets ROCm 10, so it installs a +retarget of another ROCm. Do not raise `TORCH_VERSION` to 2.13 — it drops a `c10` symbol AITER's prefill kernels link against, and they fail to load. Pass `--build-arg USERNAME=$USER --build-arg USER_UID=$(id -u) --build-arg USER_GID=$(id -g)` to match container file ownership to your host user — diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 351844125e1..43452a1ec75 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -47,12 +47,12 @@ ARG AITER_REPO=https://github.com/ROCm/aiter.git ARG AITER_GPU_ARCHS="gfx942;gfx950" ARG AITER_PREBUILD_JOBS=2 -# The wheel path (AITER_SOURCE=0). Pinned to an exact build including the local -# version segment, because pip will not select a local version from a loose -# specifier. No +rocm10.0 build of amd-aiter exists; every 0.1.20 wheel shares -# source build id 3135022, so the +rocm10.1.0a one is the nearest retarget of the -# same revision. -ARG AITER_VERSION=0.1.20+rocm10.1.0a20260819.3135022 +# The wheel path (AITER_SOURCE=0), an unblock only. Spelled in full including the +# local segment: pip will not select a local version from a loose specifier. It +# must clear AITER_MIN_VERSION or the assertion below rejects it -- and at +# 0.1.21.post2 the index publishes only +rocm7.14.0 and +rocm7.2.3, no ROCm 10 +# build at all, which is the whole reason the default path builds from source. +ARG AITER_VERSION=0.1.21.post2+rocm7.14.0.c2a6601 ARG AITER_INDEX=https://rocm.frameworks-nightlies.amd.com/whl-multi-arch/ # Fail fast if the base image ever stops shipping a HIP-enabled torch. @@ -216,7 +216,7 @@ if got != want: FLOOR = "0.1.21" from packaging.version import Version # noqa: E402 -if Version(got).base_version < Version(FLOOR).base_version: +if Version(got) < Version(FLOOR): sys.exit(f"ERROR: amd-aiter {got} is below AITER_MIN_VERSION {FLOOR}") inc = Path(md.distribution("amd-aiter").locate_file("aiter_meta/csrc/include")) From f07b293a6b1705c71e212011fd5fc83a65dd1db7 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 08:05:33 -0400 Subject: [PATCH 14/41] fix(rocm): prebuild the asm module the loader now dlopens Rebasing onto amd-integration picked up #365, which routes single prefill to AITER's asm arm on CDNA4 and dlopens module_fmha_v3_fwd.so. Its own failure hint says why that is safe: module_fmha_v3_fwd.so ships prebuilt in the amd-aiter wheel, so unlike the mha_fwd variants there is no JIT build to trigger. Its absence means the wheel is incomplete or the JIT dir points elsewhere. This branch replaces the wheel with a PREBUILD_KERNELS=0 source install, which ships no prebuilt modules at all. Measured on the image built from this branch: module_fmha_v3_fwd.so present: NO module_* present: 0 So the two changes are individually correct and jointly broken: #365's default CDNA4 prefill route would throw at dlopen in the new image, with a message blaming an incomplete wheel that is not being used. The driver now builds whole modules as well as variants. module_fmha_v3_fwd takes 96 s GPU-free at GPU_ARCHS="gfx942;gfx950" and needs no filter, because AITER resolves the asm kernel from its own config table per call rather than from filename axes. The durable part is the parity test: it greps aiter_loader.cc for every "module_*.so" it opens and fails when the driver's set does not cover them. A/B'd by emptying LOADER_MODULES -- caught. Without it the next module added to the loader repeats this silently, since nothing else connects the two files. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 63 +++++++++++++++++---- tests/rocm/test_prebuild_aiter_attention.py | 20 +++++++ 2 files changed, 72 insertions(+), 11 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index a4c358d458b..a5e2a3cb6f6 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -55,6 +55,12 @@ GENERATE_PY = "example/ck_tile/01_fmha/generate.py" +# Whole AITER modules aiter_loader.cc dlopens by name, with no variant axes. The +# wheel shipped these prebuilt; a PREBUILD_KERNELS=0 source install ships none, +# so the image has to build them or the load throws. +# tests/rocm/test_prebuild_aiter_attention.py checks this against the loader. +LOADER_MODULES = ("module_fmha_v3_fwd",) + class Variant(NamedTuple): """One ``.so``. Mirrors ``VariantKey`` in flashinfer/jit/rocm/aiter_variants.py. @@ -279,6 +285,33 @@ def _args_of_build(core, family: str) -> Dict[str, object]: return args +def build_module_by_name(md_name: str) -> Path: + """Build one whole AITER module, e.g. the asm prefill arm.""" + core = _aiter_jit_core() + args = core.get_args_of_build(md_name) + if not args.get("srcs"): + raise RuntimeError(f"aiter's config gave no srcs for {md_name}") + core.build_module( + md_name, + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + args["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args.get("torch_exclude", False), + args.get("third_party", []), + ) + out = jit_dir(core) / f"{md_name}.so" + if not out.is_file(): + raise RuntimeError(f"build of {md_name} produced no {out}") + shutil.rmtree(jit_dir(core) / "build" / md_name, ignore_errors=True) + return out + + def build_one(v: Variant) -> Path: """Build a single variant. Returns the artifact path; raises if absent after.""" core = _aiter_jit_core() @@ -308,14 +341,14 @@ def build_one(v: Variant) -> Path: return out -def _run_jobs(variants: Sequence[Variant], jobs: int) -> int: +def _run_jobs(variants: Sequence[Variant], jobs: int, extra: Sequence[str] = ()) -> int: """Build in parallel across *processes*. Never threads: AITER's build mutates process-global environment and takes a process-global lock, so two concurrent builds in one process have the first to finish restore the environment under the second. """ - pending = list(variants) + pending = list(variants) + list(extra) # Child output goes to a file, never a pipe: hipcc is verbose enough to fill # a pipe buffer, and a child blocked on a full pipe never exits, so poll() # would spin forever. Only read on failure, and only the tail. @@ -338,12 +371,16 @@ def _drain(pending, running, failed, jobs, total) -> None: while pending or running: while pending and len(running) < jobs: v = pending.pop(0) - fd, log_path = tempfile.mkstemp( - prefix=f"prebuild-{v.md_name}-", suffix=".log" - ) + name = v if isinstance(v, str) else v.md_name + fd, log_path = tempfile.mkstemp(prefix=f"prebuild-{name}-", suffix=".log") try: proc = subprocess.Popen( - [sys.executable, __file__, "--build-one", _encode(v)], + [ + sys.executable, + __file__, + "--build-one", + v if isinstance(v, str) else _encode(v), + ], stdout=fd, stderr=subprocess.STDOUT, ) @@ -357,17 +394,18 @@ def _drain(pending, running, failed, jobs, total) -> None: continue running.remove(entry) done += 1 + label = f"{v}.so" if isinstance(v, str) else v.so_name if proc.returncode == 0: print( - f"[{done}/{total}] ok {v.so_name} {time.time() - started:.0f}s", + f"[{done}/{total}] ok {label} {time.time() - started:.0f}s", flush=True, ) os.unlink(log_path) else: - failed.append(v.so_name) + failed.append(label) tail = Path(log_path).read_text(errors="replace").splitlines()[-15:] print( - f"[{done}/{total}] FAIL {v.so_name} {time.time() - started:.0f}s" + f"[{done}/{total}] FAIL {label} {time.time() - started:.0f}s" f" (full log: {log_path})\n" + "\n".join(tail), flush=True, ) @@ -430,7 +468,10 @@ def main(argv: Optional[Sequence[str]] = None) -> int: a = p.parse_args(argv) if a.build_one: - build_one(_decode(a.build_one)) + if "|" in a.build_one: + build_one(_decode(a.build_one)) + else: + build_module_by_name(a.build_one) return 0 variants = selected_variants(a.only) @@ -445,7 +486,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: if a.jobs < 1: p.error("--jobs must be >= 1") t0 = time.time() - rc = _run_jobs(variants, a.jobs) + rc = _run_jobs(variants, a.jobs, extra=list(LOADER_MODULES) if not a.only else []) print( f"\n{len(variants)} variant(s) in {(time.time() - t0) / 60:.0f} min " f"at --jobs {a.jobs}" diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 353bfc90c0b..53716d153ab 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -207,3 +207,23 @@ def test_varlen_goes_through_aiters_own_helper(): src = _DRIVER_PATH.read_text() assert "get_mha_varlen_prebuild_variants_by_names" in src assert "_md_varlen" in src # the mirror exists only to name the file we ask for + + +# --------------------------------------------------------------------------- +# 3. whole modules the loader dlopens by name +# --------------------------------------------------------------------------- + +_LOADER_CC = Path(__file__).resolve().parents[2] / "csrc" / "rocm" / "aiter_loader.cc" + + +def test_the_driver_builds_every_module_the_loader_dlopens(): + """A PREBUILD_KERNELS=0 source install ships no module_*.so at all, so any + name the loader opens has to be in the driver's set or the load throws -- + which is what happened to module_fmha_v3_fwd when the asm arm landed.""" + opened = set(re.findall(r'"(module_[a-z0-9_]+)\.so"', _LOADER_CC.read_text())) + + assert opened, "no module_*.so names found in aiter_loader.cc; did the form change?" + assert opened <= set(driver.LOADER_MODULES), ( + f"aiter_loader.cc dlopens {sorted(opened - set(driver.LOADER_MODULES))}, which " + "docker/prebuild_aiter_attention.py does not build" + ) From 13f8053d147f842c8d127882635f7f4d4c385299 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 08:40:16 -0400 Subject: [PATCH 15/41] fix(rocm): make --check cover the loader modules it builds The previous commit taught the driver to build module_fmha_v3_fwd but its own --check still enumerated only the 40 variants, so the image-build gate would pass with the asm module missing -- the exact hole that commit set out to close. Visible in that build's own output: the job runner counted [41/41] while the check printed 40/40. The edit that was meant to do this silently matched nothing, because ruff-format had rewrapped the function between writing and applying it, and unlike the other edits in that batch this one asserted no anchor. Verified against the image built from the previous commit: 520896 bytes .../aiter/jit/module_fmha_v3_fwd.so mha_*.so: 40 amdgcn-amd-amdhsa--gfx942 amdgcn-amd-amdhsa--gfx950 The artifact is there and carries both code objects; only the count was wrong. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index a5e2a3cb6f6..49ad0569feb 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -432,12 +432,9 @@ def _check(variants: Sequence[Variant]) -> int: """Assert every selected artifact is present and non-empty.""" core = _aiter_jit_core() d = jit_dir(core) - missing = [ - v.so_name - for v in variants - if not (d / v.so_name).is_file() or (d / v.so_name).stat().st_size == 0 - ] - print(f"{len(variants) - len(missing)}/{len(variants)} present in {d}") + names = [v.so_name for v in variants] + [f"{m}.so" for m in LOADER_MODULES] + missing = [n for n in names if not (d / n).is_file() or (d / n).stat().st_size == 0] + print(f"{len(names) - len(missing)}/{len(names)} present in {d}") if missing: print("missing or empty:\n " + "\n ".join(missing)) return 1 From 9f99d53879726716b67d3c35bd840da6e48e2460 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 11:57:50 -0400 Subject: [PATCH 16/41] test(rocm): pin the AITER state these two tests depend on The full suite on a box whose AITER is 0.1.20 -- below the floor this branch raises -- failed six tests that pass on the image it ships. Both causes are the tests trusting whatever AITER the host happens to have. **The new stream tests guarded on `is_aiter_supported`**, which answers for the architecture alone. On a sub-floor box they ran anyway and died inside `require_aiter` rather than skipping. They now use `requires_aiter`, the helper every other AITER test in tests/rocm uses, which is why those skipped cleanly on the same run. **`test_a_missing_aiter_package_is_its_own_reason` stubbed importability but not presence.** `_auto_select_prefill_backend` reports the ABI-floor reason when a sub-floor AITER *is* installed and the not-installed reason otherwise, so with 0.1.20 present the test asserted on the wrong branch: assert 'aiter package not installed' in 'amd-aiter 0.1.20+rocm10.1.0a... is below the 0.1.21 ABI floor (the vendored struct layouts do not match older releases)' It now pins absence as well, which is what its name claims. `TestRequireAiterRuntime` directly below already does this, for the stated reason that which branch raises "does not depend on the amd-aiter the host happens to have installed". Verified both ways: aiter 0.1.21.post2 43 passed aiter 0.1.20 sssss...........s.......ss................. exit 0 Co-Authored-By: Claude --- tests/rocm/test_aiter_stream_affinity.py | 10 +++++----- tests/rocm/test_prefill_decode_dispatch.py | 7 +++++++ 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/tests/rocm/test_aiter_stream_affinity.py b/tests/rocm/test_aiter_stream_affinity.py index 53d975fb8e0..b9a209d7baf 100644 --- a/tests/rocm/test_aiter_stream_affinity.py +++ b/tests/rocm/test_aiter_stream_affinity.py @@ -19,12 +19,12 @@ import torch import flashinfer -from flashinfer.rocm.aiter_utils import is_aiter_supported +from tests.test_helpers.test_helpers import requires_aiter -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available() or not is_aiter_supported(torch.device("cuda:0")), - reason="needs a GPU with a supported AITER", -) +# requires_aiter, not is_aiter_supported: the latter answers for the +# architecture only, so on a box whose AITER is below the ABI floor these ran +# and failed inside require_aiter instead of skipping. +pytestmark = requires_aiter DT = torch.float16 diff --git a/tests/rocm/test_prefill_decode_dispatch.py b/tests/rocm/test_prefill_decode_dispatch.py index 2f172bb8968..561ec4e088a 100644 --- a/tests/rocm/test_prefill_decode_dispatch.py +++ b/tests/rocm/test_prefill_decode_dispatch.py @@ -279,8 +279,15 @@ def test_the_reason_is_returned_on_every_call_not_only_the_first(self, device): assert first == second is not None def test_a_missing_aiter_package_is_its_own_reason(self, device, monkeypatch): + from flashinfer.rocm import aiter_utils + prefill_rocm._aiter_ops_importable.cache_clear() monkeypatch.setattr(prefill_rocm, "_aiter_ops_importable", lambda: False) + # Pin absence, not just unimportability: with an installed AITER below + # the ABI floor the other branch answers first, so this asserted on + # whatever the host happened to have -- like the class below, which pins + # both branches for exactly this reason. + monkeypatch.setattr(aiter_utils, "_aiter_installed_version", lambda: None) backend, reason = _auto(device) From de09d06526bac124ea2f92aeea706c1d4bfba591 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 14:26:19 -0400 Subject: [PATCH 17/41] fix(rocm): prebuild the fp8 batch-prefill arm the loader can ask for The driver enumerated bf16 and fp16 only, so an fp8 deployment got nothing from this branch: every fp8 paged prefill still compiled AITER inside the serving process, which is the stall the prebuild exists to remove. fp8 is reachable and batch-prefill only. prefill.py routes every fp8 query through the native paged kernel because the flat-gather route runs mha_varlen_fwd, which has no fp8 arm, and the in-tree fa2 kernel rejects 8-bit types at compile time. All three AITER mha families carry an fp8 arm upstream, but kFp8Bf16 is constructed at exactly one site in this tree (batch_prefill_paged_aiter.cu:110), so the reachable set is 8 names, not 24: mha_batch_prefill_fp8bf16_{logits,nlogits}_nbias_{mask,nmask}_{lse,nlse} _ndropout_pertensor_nsink.so. All 8 build clean, ~66 s each. The old parity test could not have caught this: it compared the driver's table against flashinfer's, and both omit fp8. The replacement reads the token literals out of aiter_loader.cc instead. A/B'd by reverting the dtype set -- it fails with "aiter_loader.cc can emit ['fp8bf16', 'pertensor'], never built". flashinfer's own table stays at 40. It serves the store route, which has no fp8 arm; its docstrings claimed to enumerate everything the loader can ask for, which was already false and is now corrected. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 21 ++++++-- flashinfer/jit/rocm/aiter_variants.py | 8 +-- tests/rocm/test_prebuild_aiter_attention.py | 59 ++++++++++++++++----- 3 files changed, 67 insertions(+), 21 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 49ad0569feb..1c4ee8c8510 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -46,6 +46,14 @@ DTYPES = ("bf16", "fp16") +# fp8 is batch-prefill only: prefill.py sends every fp8 query down the native +# paged route, because the flat-gather route runs varlen, which has no fp8 arm. +_FAMILY_DTYPES = { + "mha_fwd": DTYPES, + "mha_varlen_fwd": DTYPES, + "mha_batch_prefill": DTYPES + ("fp8bf16",), +} + # Receipts select which CK-tile instance set the codegen emits. They differ per # family and a wrong one ships a *different kernel* under the right filename, # with nothing downstream to detect it. @@ -83,6 +91,11 @@ def md_name(self) -> str: def so_name(self) -> str: return f"{self.md_name}.so" + @property + def qscale_token(self) -> str: + """fp8 always arrives with per-tensor descales; bf16/fp16 never do.""" + return "pertensor" if self.dtype == "fp8bf16" else "nqscale" + def _md_mha_fwd(v: Variant) -> str: """Mirrors aiter/ops/mha.py::mha_fwd. No logits or skip token in this family.""" @@ -122,7 +135,7 @@ def _md_batch_prefill(v: Variant) -> str: + ("_mask" if v.needs_mask else "_nmask") + ("_lse" if v.has_lse else "_nlse") + "_ndropout" - + "_nqscale" + + f"_{v.qscale_token}" + "_nsink" ) @@ -173,17 +186,17 @@ def _filter_batch_prefill(v: Variant) -> str: + ("_mask*" if v.needs_mask else "_nmask*") + ("_lse*" if v.has_lse else "_nlse*") + "_ndropout*" - + "_nqscale*" + + f"_{v.qscale_token}*" + "_nsink*" ) def reachable_variants() -> Tuple[Variant, ...]: - """Every ``.so`` the loader can ask for: 8 + 16 + 16 = 40 per architecture.""" + """Every ``.so`` the loader can ask for: 8 + 16 + 24 = 48 per architecture.""" out: List[Variant] = [] for family, has_logits_axis in _HAS_LOGITS_AXIS.items(): logits_values = (False, True) if has_logits_axis else (False,) - for dtype in DTYPES: + for dtype in _FAMILY_DTYPES[family]: for has_logits_cap in logits_values: for needs_mask in (False, True): for has_lse in (False, True): diff --git a/flashinfer/jit/rocm/aiter_variants.py b/flashinfer/jit/rocm/aiter_variants.py index 32b74e1362b..9e78ddb5f70 100644 --- a/flashinfer/jit/rocm/aiter_variants.py +++ b/flashinfer/jit/rocm/aiter_variants.py @@ -90,9 +90,9 @@ class VariantKey: def so_name(key: VariantKey) -> str: """The filename ``aiter_loader.cc`` will ask ``dlopen`` for. - The quantisation token is always ``nqscale``: the ``pertensor`` arm is fp8, - which only the batch-prefill family serves, and that family is not - ``servable_from_store`` -- see ``Family.servable_from_store``. + The quantisation token is always ``nqscale`` because this table covers only + the store route, which has no fp8 arm. The loader can also ask for the fp8 + ``pertensor`` names; ``docker/prebuild_aiter_attention.py`` builds those. """ name = key.family.prefix + key.dtype if key.family.include_logits: @@ -107,7 +107,7 @@ def so_name(key: VariantKey) -> str: def reachable_variants() -> Tuple[VariantKey, ...]: - """Every ``.so`` the loader can ask for: 8 + 16 + 16 = 40 per architecture. + """Every non-fp8 ``.so`` the loader can ask for: 8 + 16 + 16 = 40 per arch. ``mha_fwd`` contributes only 8 because ``get_aiter_mha_fwd_handle`` refuses ``has_logits_cap=true`` -- that template has no ``_logits`` arm, and a diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 53716d153ab..4c365884be1 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -77,32 +77,65 @@ def _fi_so_names() -> set: return {fi_variants.so_name(k) for k in fi_variants.reachable_variants()} -def test_the_driver_enumerates_the_same_40_variants_as_flashinfer(): - assert ( - len(driver.reachable_variants()) == len(fi_variants.reachable_variants()) == 40 - ) - - -def test_every_driver_filename_is_one_the_loader_asks_for(): - """A name the loader never asks for is built and then never opened.""" - assert {v.so_name for v in driver.reachable_variants()} == _fi_so_names() +def test_the_driver_covers_flashinfers_table_and_the_fp8_arm_it_omits(): + """flashinfer's table is the store's domain, which has no fp8 arm; the + driver writes where the loader looks, so it must carry fp8 as well.""" + assert len(driver.reachable_variants()) == 48 + assert _fi_so_names() < {v.so_name for v in driver.reachable_variants()} def test_nothing_is_trimmed_from_the_built_set(): - """Measured in the image, the whole set is 23 min at --jobs 2, so a trim buys - ~5 min of CI for a guaranteed multi-minute stall on the missing arm.""" + """A trim buys a few minutes of image build for a multi-minute stall on the + arm that turns out to be missing.""" assert set(driver.selected_variants()) == set(driver.reachable_variants()) - assert len(driver.selected_variants()) == 40 + assert len(driver.selected_variants()) == 48 @pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) def test_only_restricts_without_renaming(family): picked = driver.selected_variants(only=[family]) assert picked - assert {v.so_name for v in picked} <= _fi_so_names() + assert {v.so_name for v in picked} <= { + v.so_name for v in driver.reachable_variants() + } assert all(v.family == family for v in picked) +# --------------------------------------------------------------------------- +# 1b. every token the loader can compose is one the driver emits +# --------------------------------------------------------------------------- + +# alibi is the one arm the loader can spell but no call site reaches: all three +# construction sites hard-code has_alibi false. +_UNREACHABLE_LOADER_TOKENS = {"_alibi"} + + +def _loader_name_tokens() -> set: + """String literals `build_so_name`/`dtype_token` can put in a filename.""" + text = _LOADER_CC.read_text() + start = text.find("const char* dtype_token(") + end = text.find("\n}", text.find("std::string build_so_name(")) + assert start != -1 and end != -1, "aiter_loader.cc: name composition not found" + return set(re.findall(r'"([a-z0-9_]+)"', text[start:end])) + + +def test_every_name_token_the_loader_can_emit_is_one_the_driver_builds(): + """Table-vs-table parity cannot see a token both tables lack -- which is how + the fp8/pertensor batch-prefill arm went unbuilt.""" + segments = set() + for v in driver.reachable_variants(): + segments.update(v.md_name.split("_")) + + tokens = _loader_name_tokens() + assert "fp8bf16" in tokens, "loader no longer spells fp8; is this test stale?" + missing = { + t + for t in tokens + if t not in _UNREACHABLE_LOADER_TOKENS and t.lstrip("_") not in segments + } + assert not missing, f"aiter_loader.cc can emit {sorted(missing)}, never built" + + # --------------------------------------------------------------------------- # 2. receipts and token spellings match the installed aiter # --------------------------------------------------------------------------- From f640a1b49c9ebab6061e660c9f717f4a7fa06a9a Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 14:27:27 -0400 Subject: [PATCH 18/41] fix(rocm): scope --only's check, honour AITER's clang, keep artifacts readable Three defects found by review, each verified by reproducing it. --only built the family it was asked for and then failed its own check on the whole-module set it deliberately skips: $ prebuild_aiter_attention.py --only mha_fwd 8/9 present in .../aiter/jit missing or empty: module_fmha_v3_fwd.so # exit 1 on a healthy tree _check now takes the same `extra` set the invocation builds. build_module was called directly, which skips the HIP_CLANG_PATH export that AITER's own compile_ops wraps the build in. optCompilerConfig.json sets it from MHA_HIP_CLANG_PATH for module_mha_fwd and module_mha_varlen_fwd, so with that variable set the prebuilt .so would come from a different clang than the lazy build -- falsifying the byte-identical claim. Unset in this image, so no artifact shipped so far is affected. The chmod is the one with teeth. AITER's get_user_jit_dir() copies the whole jit tree into $HOME when its own directory is unwritable, and the tree is owned by uid 1003. Any container run with --user of another uid therefore sees none of the prebuilt artifacts -- and dies outright where HOME is not writable, which is not a fallback but a crash: PermissionError: [Errno 13] Permission denied: '/nonexistent' in get_user_jit_dir() -> shutil.copytree(this_dir, home_jit_dir) Verified fixed: the same foreign uid now resolves to the installed jit dir. --user is the documented run mode, so this would have hit normal use. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 12 ++- docker/prebuild_aiter_attention.py | 101 +++++++++++++------- tests/rocm/test_prebuild_aiter_attention.py | 14 +++ 3 files changed, 90 insertions(+), 37 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 43452a1ec75..b1df3f41584 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -43,7 +43,9 @@ ARG AITER_SOURCE=1 ARG AITER_REF=v0.1.21.post2 ARG AITER_REPO=https://github.com/ROCm/aiter.git # Both architectures in one fat binary. Must be explicit: there is no device here -# for AITER's arch detection to find. +# for AITER's arch detection to find. Note AITER's get_gfx() reads only the LAST +# entry, which gates -D__Float4_e2m1fn_x2; no attention source references it, so +# extending this prebuild to the quant or MoE ops needs that checked first. ARG AITER_GPU_ARCHS="gfx942;gfx950" ARG AITER_PREBUILD_JOBS=2 @@ -236,9 +238,15 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # Compile the attention variants that the first request would otherwise compile # in-process, under a process-global lock, minutes at a time. The driver ends in # its own --check, so a missing or empty artifact fails the build here. +# +# The chmod is load-bearing: AITER's get_user_jit_dir() falls back to copying the +# whole jit tree into $HOME when its own directory is unwritable, so a container +# run with --user of any other uid would miss every artifact built here -- or die +# outright when HOME is not writable either. RUN set -eu; \ GPU_ARCHS="$AITER_GPU_ARCHS" \ - python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" + python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ + chmod -R a+w "$(python3 -c 'import importlib.util as u, os; print(os.path.join(list(u.find_spec("aiter").submodule_search_locations)[0], "jit"))')" WORKDIR /home/$USERNAME diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 1c4ee8c8510..06be466a377 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -35,6 +35,7 @@ from __future__ import annotations import argparse +import contextlib import os import shutil import subprocess @@ -42,7 +43,7 @@ import tempfile import time from pathlib import Path -from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple +from typing import Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple DTYPES = ("bf16", "fp16") @@ -298,26 +299,49 @@ def _args_of_build(core, family: str) -> Dict[str, object]: return args +@contextlib.contextmanager +def _hip_clang_path(args: Dict) -> Iterator[None]: + """Export ``HIP_CLANG_PATH`` the way AITER's ``compile_ops`` does. + + Only ``compile_ops`` sets it; calling ``build_module`` directly would leave + the mha families on a different clang from their lazy builds. + """ + path = args.get("hip_clang_path") + if not path: + yield + return + prev = os.environ.get("HIP_CLANG_PATH") + os.environ["HIP_CLANG_PATH"] = path + try: + yield + finally: + if prev is None: + os.environ.pop("HIP_CLANG_PATH", None) + else: + os.environ["HIP_CLANG_PATH"] = prev + + def build_module_by_name(md_name: str) -> Path: """Build one whole AITER module, e.g. the asm prefill arm.""" core = _aiter_jit_core() args = core.get_args_of_build(md_name) if not args.get("srcs"): raise RuntimeError(f"aiter's config gave no srcs for {md_name}") - core.build_module( - md_name, - args["srcs"], - args["flags_extra_cc"], - args["flags_extra_hip"], - args["blob_gen_cmd"], - args["extra_include"], - args["extra_ldflags"], - args["verbose"], - args["is_python_module"], - args["is_standalone"], - args.get("torch_exclude", False), - args.get("third_party", []), - ) + with _hip_clang_path(args): + core.build_module( + md_name, + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + args["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args.get("torch_exclude", False), + args.get("third_party", []), + ) out = jit_dir(core) / f"{md_name}.so" if not out.is_file(): raise RuntimeError(f"build of {md_name} produced no {out}") @@ -330,20 +354,21 @@ def build_one(v: Variant) -> Path: core = _aiter_jit_core() recipe = build_recipe(v, core) args = _args_of_build(core, v.family) - core.build_module( - recipe["md_name"], - args["srcs"], - args["flags_extra_cc"], - args["flags_extra_hip"], - recipe["blob_gen_cmd"], - args["extra_include"], - args["extra_ldflags"], - args["verbose"], - args["is_python_module"], - args["is_standalone"], - args.get("torch_exclude", False), - args.get("third_party", []), - ) + with _hip_clang_path(args): + core.build_module( + recipe["md_name"], + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + recipe["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args.get("torch_exclude", False), + args.get("third_party", []), + ) out = jit_dir(core) / v.so_name if not out.is_file(): raise RuntimeError(f"build of {v.md_name} produced no {out}") @@ -441,11 +466,15 @@ def _decode(s: str) -> Variant: return Variant(family, dtype, bool(int(lc)), bool(int(mask)), bool(int(lse))) -def _check(variants: Sequence[Variant]) -> int: - """Assert every selected artifact is present and non-empty.""" +def _check(variants: Sequence[Variant], extra: Sequence[str]) -> int: + """Assert every selected artifact is present and non-empty. + + ``extra`` must be the module set the same invocation builds, or a ``--only`` + run fails its own check on artifacts it was never asked to produce. + """ core = _aiter_jit_core() d = jit_dir(core) - names = [v.so_name for v in variants] + [f"{m}.so" for m in LOADER_MODULES] + names = [v.so_name for v in variants] + [f"{m}.so" for m in extra] missing = [n for n in names if not (d / n).is_file() or (d / n).stat().st_size == 0] print(f"{len(names) - len(missing)}/{len(names)} present in {d}") if missing: @@ -485,23 +514,25 @@ def main(argv: Optional[Sequence[str]] = None) -> int: return 0 variants = selected_variants(a.only) + # The whole-module set has no family axis, so --only excludes it. + extra = [] if a.only else list(LOADER_MODULES) if a.list: for v in variants: print(v.so_name) print(f"\n{len(variants)} variant(s)") return 0 if a.check: - return _check(variants) + return _check(variants, extra) if a.jobs < 1: p.error("--jobs must be >= 1") t0 = time.time() - rc = _run_jobs(variants, a.jobs, extra=list(LOADER_MODULES) if not a.only else []) + rc = _run_jobs(variants, a.jobs, extra=extra) print( f"\n{len(variants)} variant(s) in {(time.time() - t0) / 60:.0f} min " f"at --jobs {a.jobs}" ) - return rc or _check(variants) + return rc or _check(variants, extra) if __name__ == "__main__": diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 4c365884be1..7850dca361b 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -101,6 +101,20 @@ def test_only_restricts_without_renaming(family): assert all(v.family == family for v in picked) +def test_check_does_not_demand_what_only_never_builds(tmp_path, monkeypatch): + """`--only` skips the whole-module set, so `_check` must skip it too or a + healthy restricted build exits 1.""" + monkeypatch.setattr(driver, "_aiter_jit_core", lambda: object()) + monkeypatch.setattr(driver, "jit_dir", lambda _core: tmp_path) + picked = driver.selected_variants(only=["mha_fwd"]) + for v in picked: + (tmp_path / v.so_name).write_bytes(b"x") + + assert driver.main(["--check", "--only", "mha_fwd"]) == 0 + # Unrestricted, the same directory is incomplete: the module is absent. + assert driver.main(["--check"]) == 1 + + # --------------------------------------------------------------------------- # 1b. every token the loader can compose is one the driver emits # --------------------------------------------------------------------------- From 0e28292204f70899b1e06acc7de7fe01d4b4800d Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 14:28:25 -0400 Subject: [PATCH 19/41] refactor(rocm): one composer for the AITER cache tag, and trim to the comment cap Moving _rocm_version into aiter_source let _aiter_cache_tag gain the ROCm component that variant_store_dir already had, which left the tag shape and its path-safety guard written out twice in two modules. The next change to either could land in one and not the other, and the two cache directories would then disagree about staleness. Both now call compose_cache_tag(). The rest is the repo's own comment cap, which this branch had been ignoring: the driver's module docstring ran 27 lines and carried a measurement, and two more docstrings ran long. Measurements belong in these commit messages. Also drops a stale claim in test_aiter_version_gate.py, which still said the vendored structs follow the 0.1.20 layout after the floor moved to 0.1.21. No behaviour change: the driver diff is docstrings only, and the tag helper composes the same string from the same inputs. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 38 ++++++------------- flashinfer/jit/rocm/aiter_source.py | 41 +++++++++------------ flashinfer/jit/rocm/aiter_variants.py | 15 +------- tests/rocm/test_aiter_version_gate.py | 4 +- tests/rocm/test_prebuild_aiter_attention.py | 11 +++--- 5 files changed, 39 insertions(+), 70 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 06be466a377..253bc2e81f4 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -3,33 +3,17 @@ # SPDX-License-Identifier: Apache-2.0 """Build AITER's attention variant ``.so`` at image-build time, with no GPU. -FlashInfer ``dlopen``s these by composed filename. The wheel prebuilds almost -none of them, so AITER's lazy JIT compiles each one *inside the serving process* -on the first request that needs it -- minutes of stall, possibly hours into the -process's life, under a process-global build lock. - -Runs during ``docker build``, where there is no device. Everything here is -therefore stdlib-only and reaches AITER through the flat ``sys.path`` import its -own ``setup.py`` uses; ``import aiter`` would run arch detection and a triton -import that both need a GPU. The companion GPU-attached driver -(``flashinfer.rocm.prebuild_aiter_variants``) produces variants by *calling the -op*, which is why that one cannot run here. - -Artifacts land in AITER's own ``jit/`` directory, so AITER itself finds them -already built and the bootstraps return without compiling. That is what makes -``mha_batch_prefill`` worth building by this route and not by the store route: -its bootstrap doubles as the page-size probe and runs either way, so only a hit -AITER can see spares the compile. Measured: 101.4 s cold against 0.2 s prebuilt. - - python3 docker/prebuild_aiter_attention.py --list - python3 docker/prebuild_aiter_attention.py --jobs 2 - python3 docker/prebuild_aiter_attention.py --check - -This file deliberately carries its own copy of the variant table: importing -``flashinfer.jit.rocm.aiter_variants`` would pull in torch. -``tests/rocm/test_prebuild_aiter_attention.py`` asserts the two agree, and -asserts the receipts and token spellings against the installed ``aiter/ops/mha.py`` --- the only detector for an upstream recipe change. +Without this, AITER's lazy JIT compiles each variant inside the serving process +on the first request that needs it. Artifacts land in AITER's own ``jit/`` +directory so its bootstraps find them built -- the store route cannot serve +``mha_batch_prefill``, whose bootstrap doubles as the page-size probe. + +Stdlib-only, and reaches AITER by the flat ``sys.path`` import its own +``setup.py`` uses: ``import aiter`` needs a device. The table here is a +deliberate copy of ``flashinfer.jit.rocm.aiter_variants``, which pulls torch; +``tests/rocm/test_prebuild_aiter_attention.py`` holds the two in step. + + python3 docker/prebuild_aiter_attention.py --list|--check|--jobs 2 """ from __future__ import annotations diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 4773158a148..81d02834b9f 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -211,29 +211,17 @@ def _rocm_version() -> str: return "unknown" -def _aiter_cache_tag() -> str: - """A filesystem-safe tag keying the lib cache by arch, AITER version and ROCm. - - Without the arch component, a lib built for one arch would be silently reused - on a machine with a different arch. Without the version component, an AITER - upgrade (which can change the C++ ABI the FlashInfer shim links against) would - silently reuse the stale .so. The FlashInfer JIT dir already keys by its own - version+arch, but this cache sits outside it. - - ROCm is a component because a source-built AITER's version is bare - ("0.1.21.post2"), where a wheel's carried "+rocm10.1.0a..." and keyed ROCm by - accident. Without it, two ROCm toolchains share one tag. - - Keyed on the *resolved* architecture -- the one actually compiled for -- so - the tag cannot disagree with the contents of the directory it names.""" - arch = resolve_aiter_build_arch() - # The tag is joined onto the cache root to create a directory, so - # "filesystem-safe" above has to be enforced, not just asserted in prose. - # _env_arch_list already rejects anything that is not an architecture name; - # this keeps the guarantee true for the other two sources of `arch` (the - # device probe and the default) and for any future caller. Loud rather than - # silently sanitized: an arch that needs rewriting means the resolver is - # wrong, and a quietly renamed cache directory would hide that. +def compose_cache_tag(arch: Optional[str] = None) -> str: + """``__aiter-__rocm-``, the key both AITER caches use. + + Staleness is structural rather than detected: bump the arch, AITER's ABI or + the toolchain and the tag names a different directory, so a mismatched + artifact is never found. ``arch`` defaults to the *resolved* build arch. + """ + arch = arch or resolve_aiter_build_arch() + # The tag becomes a directory name, so enforce that rather than assume it. + # Loud rather than silently sanitized: an arch that needs rewriting means the + # resolver is wrong, and a quietly renamed cache directory would hide that. if not arch or arch != Path(arch).name or arch.startswith("."): raise ValueError( f"refusing to build a cache directory name from architecture " @@ -245,6 +233,9 @@ def _aiter_cache_tag() -> str: version = _md.version("amd-aiter") except Exception: version = "unknown" + # ROCm is a component because a source-built AITER's version is bare + # ("0.1.21.post2") where a wheel's carried "+rocm..." and keyed ROCm by + # accident; without it, two toolchains would share one tag. tag = f"{arch}__aiter-{version}__rocm-{_rocm_version()}" # The version strings come from package metadata and torch, neither of which # this module controls, so the composed tag gets the same check the arch did. @@ -256,6 +247,10 @@ def _aiter_cache_tag() -> str: return tag +def _aiter_cache_tag() -> str: + return compose_cache_tag() + + @functools.lru_cache(maxsize=1) def _aiter_libs_dir() -> Path: # Keyed by arch + AITER version + ROCm so a cached lib is never reused across diff --git a/flashinfer/jit/rocm/aiter_variants.py b/flashinfer/jit/rocm/aiter_variants.py index 9e78ddb5f70..eed7cf1a7c2 100644 --- a/flashinfer/jit/rocm/aiter_variants.py +++ b/flashinfer/jit/rocm/aiter_variants.py @@ -21,7 +21,7 @@ from typing import List, Optional, Tuple from .. import env as jit_env -from .aiter_source import _rocm_version, resolve_aiter_build_arch +from .aiter_source import compose_cache_tag __all__ = [ "Family", @@ -209,18 +209,7 @@ def variant_store_dir(arch: Optional[str] = None) -> Path: old contents are simply never found and the lookup misses into a rebuild. A mismatched artifact is never loaded. """ - try: - import importlib.metadata as _md - - aiter_version = _md.version("amd-aiter") - except Exception: - aiter_version = "unknown" - tag = f"{arch or resolve_aiter_build_arch()}__aiter-{aiter_version}__rocm-{_rocm_version()}" - # The tag becomes a directory name; refuse anything that is not one - # component, the same guard _aiter_cache_tag applies. - if not tag or tag != Path(tag).name or tag.startswith("."): - raise ValueError(f"refusing to build a cache directory name from {tag!r}") - return jit_env.FLASHINFER_CACHE_DIR / "aiter_variants" / tag + return jit_env.FLASHINFER_CACHE_DIR / "aiter_variants" / compose_cache_tag(arch) def _explicit_stores() -> List[Path]: diff --git a/tests/rocm/test_aiter_version_gate.py b/tests/rocm/test_aiter_version_gate.py index d9a64eeca0f..4ec4ae2462d 100644 --- a/tests/rocm/test_aiter_version_gate.py +++ b/tests/rocm/test_aiter_version_gate.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 """The amd-aiter ABI floor. -The vendored structs under include/flashinfer/rocm/attention/aiter/ follow the 0.1.20 +The vendored structs under include/flashinfer/rocm/attention/aiter/ follow the 0.1.21 layout and travel by value through dlsym'd pointers, so an older AITER shifts field offsets instead of failing to load. Nothing downstream can detect that, which is why the floor is enforced before routing rather than at the call. @@ -47,7 +47,7 @@ def _doc_references(): ("0.1.20", False), ("0.1.21.dev0", False), # A .dev0 *of post3* sorts above 0.1.16, which is what makes the - # nightly-vs-prerelease distinction non-obvious -- still below 0.1.20. + # nightly-vs-prerelease distinction non-obvious -- still below the floor. ("0.1.16.post3.dev0+g620287969.d20260725", False), ("0.1.21", True), ("0.2.0", True), diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 7850dca361b..b40a78ab0dc 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -40,6 +40,10 @@ def _load_driver(): return mod +if not _DRIVER_PATH.is_file(): + # An installed wheel has no docker/ tree; skip rather than fail collection. + pytest.skip(f"{_DRIVER_PATH} not present", allow_module_level=True) + driver = _load_driver() @@ -216,11 +220,8 @@ def _mha_recipes_source() -> str: def _md_tokens(family: str) -> set: """Every name token AITER can emit for this family, dtype excluded. - mha_fwd and batch_prefill append string literals to md_name directly. varlen - delegates its whole suffix to `compose_mha_fwd_variant_suffix_and_filter`, so - its literals live in mha_recipes.py instead. Both spell the dtype with an - f-string, which leaves no literal to find -- hence dtype is stripped by the - caller rather than matched here. + varlen delegates its suffix to `compose_mha_fwd_variant_suffix_and_filter`, + so its literals live in mha_recipes.py rather than the composer body. """ source = ( _mha_recipes_source() if family == "mha_varlen_fwd" else _composer_body(family) From 6d80af63d987ac7a281c40f3f02897e64010ce30 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 15:43:14 -0400 Subject: [PATCH 20/41] docs(rocm): record which stream-affinity cases actually catch a missing guard A/B on gfx950, guards stripped from norm_aiter.cu and rope_aiter.cu and the shim rebuilt from a cold cache: guards in source: 0 -> 2 failed, 3 passed FAILED test_rmsnorm_aiter_is_correct_on_a_side_stream[4096] NaN, 1048576 / 1048576 mismatched FAILED test_rope_aiter_is_correct_on_a_side_stream 518956 / 524288 mismatched, 13.5 absolute guards in source: 3 -> 5 passed So the suite does catch the regression, but not through the case whose docstring claimed it: fused_add_rmsnorm passes either way, because its buffers are cloned on the default stream before the call and so never meet the stream-ordered allocator hazard the other two hit. Left as coverage of the in-place path with the claim corrected rather than strengthened -- a race that reproduces sometimes is worse than two that reproduce every run. Co-Authored-By: Claude --- tests/rocm/test_aiter_stream_affinity.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/rocm/test_aiter_stream_affinity.py b/tests/rocm/test_aiter_stream_affinity.py index b9a209d7baf..2a196bb6e3d 100644 --- a/tests/rocm/test_aiter_stream_affinity.py +++ b/tests/rocm/test_aiter_stream_affinity.py @@ -59,8 +59,11 @@ def test_rmsnorm_aiter_is_correct_on_a_side_stream(hidden): def test_fused_add_rmsnorm_aiter_is_correct_on_a_side_stream(): - """The in-place path is the exposed one: it stages into a fresh buffer and - copies back, so a kernel on the wrong stream can lose the race to the copy.""" + """Covers the in-place path, which stages into a fresh buffer and copies back. + + Measured: this case still passes with the guard removed. rmsnorm at 4096 and + the rope case are the two that actually detect a missing guard. + """ hidden = 1024 x = torch.randn(256, hidden, device="cuda", dtype=DT) res = torch.randn(256, hidden, device="cuda", dtype=DT) From 77334430c9132169dd78c8f989a5a67eebe0ab84 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 17:45:15 -0400 Subject: [PATCH 21/41] test(rocm): cover fp8 paged prefill, the only route that serves fp8 Nothing in tests/rocm exercised fp8 attention, which left the eight fp8 variants the previous commit added with no end-to-end check at all. That matters more than the usual coverage argument: the driver's `_check` only proves the .so exists and is non-empty, so a CK --filter that matched zero instances would still link, still pass the build, and fail only at dispatch. Measured on gfx950 with the driver-built artifacts in place -- AITER loads the file our composed name asks for and the kernel runs: [aiter] import [mha_batch_prefill_fp8bf16_nlogits_nbias_mask_nlse _ndropout_pertensor_nsink] under .../aiter/jit/ fp8 aiter ok, out=(128, 8, 128) torch.bfloat16 mean abs err 0.0031 rel 0.0505 max 0.0293 The tolerance is deliberately loose. This asserts the kernel dispatched and did not return garbage or NaN, not that fp8 quantisation is accurate; a tight bound here would fail on the quantisation error rather than on a defect. The second case pins the refusal: fa2 has no fp8 kernel, and without an explicit raise the caller gets a ninja compiler log from a static_assert. Co-Authored-By: Claude --- tests/rocm/test_fp8_paged_prefill_aiter.py | 133 +++++++++++++++++++++ 1 file changed, 133 insertions(+) create mode 100644 tests/rocm/test_fp8_paged_prefill_aiter.py diff --git a/tests/rocm/test_fp8_paged_prefill_aiter.py b/tests/rocm/test_fp8_paged_prefill_aiter.py new file mode 100644 index 00000000000..710cf86049c --- /dev/null +++ b/tests/rocm/test_fp8_paged_prefill_aiter.py @@ -0,0 +1,133 @@ +# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. +# +# SPDX-License-Identifier: Apache-2.0 + +"""fp8 paged prefill, the only route that serves fp8 at all. + +The in-tree fa2 kernel rejects 8-bit types at compile time and the flat-gather +route runs mha_varlen_fwd, which has no fp8 arm -- so AITER's native paged +kernel is the whole of fp8 prefill support, and nothing else here exercised it. + +It is also the only end-to-end check on the eight fp8 variants +docker/prebuild_aiter_attention.py composes: a CK ``--filter`` that matched no +instances still links, still passes the driver's own --check, and fails only +here, at dispatch. +""" + +import pytest +import torch + +import flashinfer +from flashinfer.rocm.prefill import ( + FP8_PREFILL_OUT_DTYPE, + _aiter_paged_route_page_sizes, +) +from tests.test_helpers.test_helpers import requires_aiter + +pytestmark = requires_aiter + +NUM_QO_HEADS, NUM_KV_HEADS, HEAD_DIM = 8, 2, 128 +PAGE, KV_LEN, Q_LEN, BATCH = 16, 512, 64, 2 +FP8_MAX = 240.0 + + +def _native_fp8_dtype(): + """This GPU's fp8 encoding. The two are indistinguishable by shape or by the + .so name, and the wrong one is read under the wrong bias and returns NaN.""" + aiter = pytest.importorskip("aiter") + try: + return aiter.dtypes.fp8 + except AttributeError: + pytest.skip("aiter.dtypes.fp8 is unreadable, so the encoding is unknown") + + +@pytest.fixture(scope="module") +def workspace(): + return torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0") + + +def _plan_and_run(workspace, q, kv, backend, **kw): + pages_per_seq = KV_LEN // PAGE + dev = q.device + w = flashinfer.BatchPrefillWithPagedKVCacheWrapper( + workspace, "NHD", backend=backend + ) + w.plan( + torch.arange(0, (BATCH + 1) * Q_LEN, Q_LEN, device=dev, dtype=torch.int32), + torch.arange( + 0, (BATCH + 1) * pages_per_seq, pages_per_seq, device=dev, dtype=torch.int32 + ), + torch.arange(BATCH * pages_per_seq, device=dev, dtype=torch.int32), + torch.full((BATCH,), PAGE, device=dev, dtype=torch.int32), + NUM_QO_HEADS, + NUM_KV_HEADS, + HEAD_DIM, + PAGE, + causal=True, + q_data_type=q.dtype, + kv_data_type=kv.dtype, + ) + return w.run(q, kv, **kw) + + +def test_fp8_paged_prefill_matches_a_bf16_reference(workspace): + fp8 = _native_fp8_dtype() + if PAGE not in _aiter_paged_route_page_sizes(fp8): + pytest.skip(f"page_size={PAGE} does not route to the native paged kernel") + + torch.manual_seed(0) + dev = "cuda:0" + q = torch.randn( + BATCH * Q_LEN, NUM_QO_HEADS, HEAD_DIM, device=dev, dtype=torch.bfloat16 + ) + kv = torch.randn( + BATCH * (KV_LEN // PAGE), + 2, + PAGE, + NUM_KV_HEADS, + HEAD_DIM, + device=dev, + dtype=torch.bfloat16, + ) + + ref = _plan_and_run(workspace, q, kv, "fa2") + + sq = q.abs().amax().float() / FP8_MAX + sk = kv.abs().amax().float() / FP8_MAX + got = _plan_and_run( + workspace, + (q.float() / sq).to(fp8), + (kv.float() / sk).to(fp8), + "aiter", + scale_q=sq.reshape(1), + scale_k=sk.reshape(1), + scale_v=sk.reshape(1), + ) + + assert got.dtype == FP8_PREFILL_OUT_DTYPE + assert not torch.isnan(got).any(), "fp8 prefill returned NaN" + rel = (got.float() - ref.float()).abs().mean() / ref.float().abs().mean() + # Loose on purpose: this asserts the kernel ran and is not garbage, not that + # fp8 quantisation is accurate. + assert rel < 0.15, f"fp8 output differs from the bf16 reference by rel={rel:.4f}" + + +def test_fp8_prefill_is_refused_on_fa2(workspace): + """fa2 has no fp8 kernel, so the refusal must be an error, not a fallback.""" + fp8 = _native_fp8_dtype() + dev = "cuda:0" + q = torch.zeros( + BATCH * Q_LEN, NUM_QO_HEADS, HEAD_DIM, device=dev, dtype=torch.bfloat16 + ) + kv = torch.zeros( + BATCH * (KV_LEN // PAGE), + 2, + PAGE, + NUM_KV_HEADS, + HEAD_DIM, + device=dev, + dtype=torch.bfloat16, + ) + + with pytest.raises(NotImplementedError, match="fp8"): + _plan_and_run(workspace, q.to(fp8), kv.to(fp8), "fa2") From 18c224ae9b5c34847366fa7064a6964fc5b8ed69 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 17:45:35 -0400 Subject: [PATCH 22/41] fix(rocm): reap killed builds, chmod what the driver wrote, run the long job first Three from review, all in the prebuild path. _run_jobs killed its in-flight children on the way out but never reaped them. AITER treats a build baton as stale only when kill(pid, 0) fails, which a zombie child of a living parent answers successfully -- so an interrupted run left the next build for that md_name blocking forever, indistinguishable from a slow compile. It also leaked the mkstemp log of every killed job. The chmod recomputed AITER's jit directory from the package path while the driver wrote to core.get_user_jit_dir(), which honours AITER_JIT_DIR and falls back to $HOME. The two agree in this image and could silently stop agreeing, chmod-ing an empty directory while reporting success -- reintroducing exactly the foreign-uid failure the chmod exists to prevent. The driver now prints the directory it resolved (--print-jit-dir) and the Dockerfile chmods that. a+rwX rather than a+w: the intent is that any uid can write into the tree, and X grants directory traversal without marking regular files executable. module_fmha_v3_fwd is the longest single job -- it compiles the whole asm FAV3 source set rather than one filtered CK instance group -- and was queued last, so at --jobs 2 it ran alone on the tail with a core idle. Queued first it overlaps the variants. --list and the closing summary now count the whole-module builds too; before, progress said [49/49] while the summary said 48 and --list named neither. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 2 +- docker/prebuild_aiter_attention.py | 28 +++++++++++++++++++++------- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b1df3f41584..172a8df6ecb 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -246,7 +246,7 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil RUN set -eu; \ GPU_ARCHS="$AITER_GPU_ARCHS" \ python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ - chmod -R a+w "$(python3 -c 'import importlib.util as u, os; print(os.path.join(list(u.find_spec("aiter").submodule_search_locations)[0], "jit"))')" + chmod -R a+rwX "$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)" WORKDIR /home/$USERNAME diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 253bc2e81f4..6ac2990d4d8 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -370,7 +370,7 @@ def _run_jobs(variants: Sequence[Variant], jobs: int, extra: Sequence[str] = ()) process-global lock, so two concurrent builds in one process have the first to finish restore the environment under the second. """ - pending = list(variants) + list(extra) + pending = list(extra) + list(variants) # Child output goes to a file, never a pipe: hipcc is verbose enough to fill # a pipe buffer, and a child blocked on a full pipe never exits, so poll() # would spin forever. Only read on failure, and only the tail. @@ -381,8 +381,13 @@ def _run_jobs(variants: Sequence[Variant], jobs: int, extra: Sequence[str] = ()) try: _drain(pending, running, failed, jobs, total) finally: - for _, proc, _, _ in running: + for _, proc, _, log_path in running: proc.kill() + # wait(), not just kill(): AITER's baton is only stale when + # kill(pid, 0) fails, which a zombie answers successfully, so an + # unreaped child makes the next build for that name hang. + proc.wait() + os.unlink(log_path) if failed: print(f"\n{len(failed)} build(s) failed:\n " + "\n ".join(failed)) return 1 if failed else 0 @@ -487,9 +492,18 @@ def main(argv: Optional[Sequence[str]] = None) -> int: default=int(os.environ.get("AITER_PREBUILD_JOBS", "1")), help="concurrent build processes", ) + p.add_argument( + "--print-jit-dir", + action="store_true", + help="print the directory artifacts are written to, and exit", + ) p.add_argument("--build-one", help=argparse.SUPPRESS) a = p.parse_args(argv) + if a.print_jit_dir: + print(jit_dir(_aiter_jit_core())) + return 0 + if a.build_one: if "|" in a.build_one: build_one(_decode(a.build_one)) @@ -501,9 +515,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: # The whole-module set has no family axis, so --only excludes it. extra = [] if a.only else list(LOADER_MODULES) if a.list: - for v in variants: - print(v.so_name) - print(f"\n{len(variants)} variant(s)") + for name in [v.so_name for v in variants] + [f"{m}.so" for m in extra]: + print(name) + print(f"\n{len(variants)} variant(s) + {len(extra)} module(s)") return 0 if a.check: return _check(variants, extra) @@ -513,8 +527,8 @@ def main(argv: Optional[Sequence[str]] = None) -> int: t0 = time.time() rc = _run_jobs(variants, a.jobs, extra=extra) print( - f"\n{len(variants)} variant(s) in {(time.time() - t0) / 60:.0f} min " - f"at --jobs {a.jobs}" + f"\n{len(variants)} variant(s) + {len(extra)} module(s) in " + f"{(time.time() - t0) / 60:.0f} min at --jobs {a.jobs}" ) return rc or _check(variants, extra) From ea1489620f5d792e41b7270ec698e817ae0aa529 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 17:45:35 -0400 Subject: [PATCH 23/41] docs(skill): correct the native page-size guidance, and guard two probes The benchmark skill's page-size bullet stated {128, 256, 1024} and called {1, 16, 1024} "wrong". It has it backwards: _aiter_native_page_sizes() returns frozenset({1, 16, 1024}) for any AITER >= 0.1.10. The bullet also called the {16, 1024} arm unreachable at the 0.1.21 floor, when it is the PackageNotFoundError fallback and reachable whenever metadata is missing. Its measured claim was stale too -- "all three advertised sizes fail the probe ... mha_batch_prefill is unreachable". On 0.1.21.post2 the native paged kernel runs: 62 TFLOPS at page_size=1024, kv=1024 on gfx950. What is true, and now recorded, is that the sweep has to stop there; beyond kv=1024 it faults inside FmhaBatchPrefillWithPagedKVCacheKernel on 0.1.20 and 0.1.21.post2 alike. Also adds importorskip to the two aiter_meta probes. They were the only aiter-dependent tests in the tree without a guard, and would error rather than skip on precisely the layouts the fallback they cover was written for. Co-Authored-By: Claude --- .claude/skills/benchmark-kernel/SKILL.md | 3 ++- tests/rocm/test_aiter_source_probes.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.claude/skills/benchmark-kernel/SKILL.md b/.claude/skills/benchmark-kernel/SKILL.md index d5ac867752a..69d54db3e25 100644 --- a/.claude/skills/benchmark-kernel/SKILL.md +++ b/.claude/skills/benchmark-kernel/SKILL.md @@ -25,7 +25,8 @@ For the in-repo profiler wrapper, see [`profiler/rocm/rocm_profiler.py`](../../. - Explicit `backend="aiter"` + `kv_layout != "NHD"` → `ValueError`. Grep [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for `only supports kv_layout`; single prefill raises at the call, both batch wrappers at `plan()`. Not raised by auto-selection — that path silently falls back to `fa2`. - Explicit `backend="aiter"` on non-gfx942/gfx950 → `RuntimeError`. - `amd-aiter` not importable → `ImportError`. - - **"Native" page sizes** (no flat-gather): `{128, 256, 1024}` — see `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py). The `{16, 1024}` arm there is unreachable now that `flashinfer.rocm.aiter_utils.AITER_MIN_VERSION` floors AITER at 0.1.21. It is only a hint: `plan()` probes with `_aiter_native_paging_available()` and degrades to flat-gather if the installed AITER cannot serve the config. **Non-native page sizes are NOT rejected** — they flat-gather. So the "{1, 16, 1024}" guidance from older docs is wrong. **Measured on amd-aiter 0.1.21.post2, all three advertised sizes fail the probe** (`no matching kernel found`) on gfx942 *and* gfx950, so every paged prefill flat-gathers and `mha_batch_prefill` is unreachable — benchmark the paged path as flat-gather until that is fixed. + - **"Native" page sizes** (no flat-gather): `_aiter_native_page_sizes()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) returns `{1, 16, 1024}` for any AITER >= 0.1.10; the `{16, 1024}` arm is the fallback for when `amd-aiter` metadata is unreadable, not dead code. What is *routed* native is narrower still: `_aiter_paged_route_page_sizes()` gives fp8 all three and fp16/bf16 only `{1024}`, because the gather measured equal or faster elsewhere. **Non-native page sizes are NOT rejected** -- they flat-gather. + - Measured on amd-aiter 0.1.21.post2, gfx950: native paged prefill at `page_size=1024` runs (62 TFLOPS at `kv=1024`). Beyond `kv=1024` it hits a GPU memory fault in `FmhaBatchPrefillWithPagedKVCacheKernel`, reproducible on 0.1.20 as well -- so bound the paged sweep at `kv=1024` rather than reading the fault as a regression. - Auto-selection (no explicit `backend=`) silently falls back to `fa2` for any of: `kv_layout != "NHD"`, custom mask, dtype not in `{fp16, bf16}`, `dtype_q != dtype_kv`, `head_dim_qk != head_dim_vo`, `pos_encoding_mode != "NONE"`, or `amd-aiter` not importable. See `_auto_select_prefill_backend()` in [`rocm/prefill.py`](../../../flashinfer/rocm/prefill.py) for the authoritative list; it returns `(backend, reason)`, and the reason names the constraint that forced the fallback. - **Always verify numerical parity before trusting perf numbers.** Compare default-HIP vs AITER outputs with `torch.testing.assert_close(rtol=1e-2, atol=1e-2)` for BF16/FP16 first. - **`gcnArchName` is the unambiguous arch marker.** Device strings show `cuda:0` on AMD too. Record `torch.cuda.get_device_properties(0).gcnArchName` and `torch.version.hip` alongside every number — a `gfx942` / ROCm 7.2 result is not comparable to a `gfx950` / ROCm 7.0.2 result. diff --git a/tests/rocm/test_aiter_source_probes.py b/tests/rocm/test_aiter_source_probes.py index bcb43758e3a..89a37332e70 100644 --- a/tests/rocm/test_aiter_source_probes.py +++ b/tests/rocm/test_aiter_source_probes.py @@ -164,6 +164,7 @@ def _blocked(name, *args, **kwargs): monkeypatch.setattr(builtins, "__import__", _blocked) def test_aiter_meta_is_preferred_when_present(self): + pytest.importorskip("aiter_meta") found = aiter_source._aiter_csrc_include_dir() assert found.name == "include" and found.parent.name == "csrc" @@ -179,7 +180,7 @@ def test_an_absent_aiter_meta_falls_back_to_aiter_csrc_dir( def test_an_incomplete_aiter_meta_falls_back_too(self, monkeypatch, tmp_path): """Present but pointing nowhere useful is the same failure as absent.""" - import aiter_meta + aiter_meta = pytest.importorskip("aiter_meta") inc = tmp_path / "csrc" / "include" inc.mkdir(parents=True) From 1d85803c1bb0f7c9f1827ff24c1abbcb293c0bee Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 17:47:39 -0400 Subject: [PATCH 24/41] fix(rocm): make only the jit directories world-writable, not the kernels a+rwX on the whole tree left every compiled kernel -rwxrwxrwx in the shipped image, so any uid inside the container could replace a .so the process later dlopens. AITER only needs to *create* files in that directory -- its get_user_jit_dir() check is os.access(dir, W_OK) -- so the write bit belongs on the directories alone. Measured in the image after the change: directory 777, artifact 755. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 172a8df6ecb..7a33f4713f9 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -242,11 +242,13 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # The chmod is load-bearing: AITER's get_user_jit_dir() falls back to copying the # whole jit tree into $HOME when its own directory is unwritable, so a container # run with --user of any other uid would miss every artifact built here -- or die -# outright when HOME is not writable either. +# outright when HOME is not writable either. Only the directories go a+w; leaving +# the .so world-writable would let any uid swap a compiled kernel. RUN set -eu; \ GPU_ARCHS="$AITER_GPU_ARCHS" \ python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ - chmod -R a+rwX "$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)" + JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ + chmod -R a+rX "$JIT"; find "$JIT" -type d -exec chmod a+w {} + WORKDIR /home/$USERNAME From 02bc2a3c8dc90e7b3c35d4a1ec6a5e5f72faa29f Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 18:11:58 -0400 Subject: [PATCH 25/41] fix(rocm): prebuild the two modules a source install otherwise leaves cold PREBUILD_KERNELS=0 ships none of AITER's own modules, and LOADER_MODULES covered only module_fmha_v3_fwd -- the one that throws at dlopen. The two that merely stall were missed, so the Dockerfile comment calling the lost prebuilts "no loss" was wrong on its own terms: it already noted aiter.mla loads one. Measured in the built image, running tests/rocm/test_mla_aiter.py against a 49-artifact jit dir: it came back with 51. The two AITER built in-process were module_mla_asm and module_aiter_core, the latter also built by the attention path, so both sit in front of a first request rather than off to one side. Both build GPU-free through the same driver path as the variants: ok module_aiter_core 28s ok module_mla_asm 15s 43 s of image build to take that cost out of every container start. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 6ac2990d4d8..4ed142fa979 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -48,11 +48,12 @@ GENERATE_PY = "example/ck_tile/01_fmha/generate.py" -# Whole AITER modules aiter_loader.cc dlopens by name, with no variant axes. The -# wheel shipped these prebuilt; a PREBUILD_KERNELS=0 source install ships none, -# so the image has to build them or the load throws. -# tests/rocm/test_prebuild_aiter_attention.py checks this against the loader. -LOADER_MODULES = ("module_fmha_v3_fwd",) +# Whole AITER modules with no variant axes. The wheel shipped these prebuilt; a +# PREBUILD_KERNELS=0 source install ships none, so the image builds them or the +# first caller pays for them -- fmha_v3_fwd throws at dlopen, the other two only +# stall. mla_asm is aiter.mla's, and aiter_core is on the common attention path. +# tests/rocm/test_prebuild_aiter_attention.py checks the first against the loader. +LOADER_MODULES = ("module_fmha_v3_fwd", "module_aiter_core", "module_mla_asm") class Variant(NamedTuple): From 6f481af763411d6bf9befa4ed22a6fb4014a2784 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 18:11:58 -0400 Subject: [PATCH 26/41] fix(rocm): close the sticky-bit hole, export GPU_ARCHS, correct three stale claims The previous commit narrowed the chmod to directories and claimed that stopped any uid swapping a compiled kernel. It did not: unlink is governed by the directory's write bit, not the file's mode, so a+w on the directory let any uid delete a .so and write its own in its place. Directories now go a+wt -- the sticky bit leaves removal to the owner -- which is what the comment claimed. GPU_ARCHS was a command prefix on the prebuild, so the --print-jit-dir call on the next line ran without it. It resolves the same path either way here, but a chmod aimed at a different directory than the one just populated is the exact silent failure this line exists to prevent; now exported for both. Three claims contradicted by this branch's own measurements: - The tutorial said the native paged probe "fails for every advertised size on both architectures, so this path always gathers". It serves page_size=1024, and 16 for fp8 -- which the new fp8 test depends on, since an fp8 query that lands on flat-gather raises rather than degrading. - refresh_aiter_jitspec's docstring still named the two-component cache tag that this branch replaced with a three-component one. - MLA's ImportError still said "Install a wheel", now the unsupported path. Also trims the soft-cap comment to its conclusion; the per-row numbers are in fa4383ec's message, which is where a measurement belongs. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 9 +++++---- examples/amd_flashinfer_rocm_tutorial.ipynb | 2 +- flashinfer/jit/rocm/aiter_source.py | 2 +- flashinfer/rocm/arch_caps.py | 8 ++------ flashinfer/rocm/mla.py | 6 +++--- 5 files changed, 12 insertions(+), 15 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7a33f4713f9..7e123436289 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -242,13 +242,14 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # The chmod is load-bearing: AITER's get_user_jit_dir() falls back to copying the # whole jit tree into $HOME when its own directory is unwritable, so a container # run with --user of any other uid would miss every artifact built here -- or die -# outright when HOME is not writable either. Only the directories go a+w; leaving -# the .so world-writable would let any uid swap a compiled kernel. +# outright when HOME is not writable either. Directories go a+wt rather than a+w: +# write alone lets any uid unlink a kernel and drop in its own, whatever the +# file's own mode says; the sticky bit leaves that to the owner. RUN set -eu; \ - GPU_ARCHS="$AITER_GPU_ARCHS" \ + export GPU_ARCHS="$AITER_GPU_ARCHS"; \ python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ - chmod -R a+rX "$JIT"; find "$JIT" -type d -exec chmod a+w {} + + chmod -R a+rX "$JIT"; find "$JIT" -type d -exec chmod a+wt {} + WORKDIR /home/$USERNAME diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index e344a060644..596cd03061e 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -164,7 +164,7 @@ "\n", "**Backend integration.** amd-flashinfer offers [AITER](https://github.com/ROCm/aiter)—AMD’s SOTA kernel library for large language model operators on ROCm—as an optional **`backend=\"aiter\"`**. When the `aiter` package is present and the requested operator is integrated with AITER in this port, work **delegates** to AITER; otherwise the default HIP path applies. The following cells exercise the AITER-backed prefill path.\n", "\n", - "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the probe fails for every advertised size on both architectures, so this path always gathers." + "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the native kernel does serve `page_size=1024` (and `16` for fp8); other sizes gather." ] }, { diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 81d02834b9f..18cbc8a6a82 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -264,7 +264,7 @@ 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/__aiter-/``, and reach the module only as an + ``aiter_libs/__aiter-__rocm-/``, 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 diff --git a/flashinfer/rocm/arch_caps.py b/flashinfer/rocm/arch_caps.py index c1e1a281854..b303bbc7bc7 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -271,12 +271,8 @@ def __post_init__(self) -> None: # Not a KnownBad row: those gate a whole (op, backend, arch) on toolchain # version and would also disable the clean logits_soft_cap=0 path. What varies # by architecture is whether the kernel is affected, not at which length. -# -# Re-measured on amd-aiter 0.1.21.post2, 2026-09-14, over non-square -# qo_len x kv_len x cap at head_dim=128 causal, vs an fp32 reference, with this -# guard disabled so the kernel runs. Unchanged: gfx950 wrong on all 12 capped -# rows (9 at 1.19-4.11 absolute, 3 returning NaN) where fa2 is 0.0001-0.0011 on -# the same shapes; gfx942 clean on all 12. cap=0 is clean on both. +# Re-measured on 0.1.21.post2 and unchanged: gfx950 wrong at every capped +# shape, gfx942 clean at every one, cap=0 clean on both. _AITER_SOFTCAP_DEFECT_ARCHS = {"gfx942": False, "gfx950": True} diff --git a/flashinfer/rocm/mla.py b/flashinfer/rocm/mla.py index d32adaeca09..4ed79385e50 100644 --- a/flashinfer/rocm/mla.py +++ b/flashinfer/rocm/mla.py @@ -134,9 +134,9 @@ def _require_aiter_mla(device: torch.device) -> None: from .aiter_utils import AITER_MIN_VERSION raise ImportError( - "The 'aiter' package is required for MLA on ROCm. Install a wheel >= " - f"{AITER_MIN_VERSION}; see docs/rocm/backends.md for the index and the " - "pinned version. A source build tracks master, whose C ABI does not " + "The 'aiter' package is required for MLA on ROCm. Build it from source " + f"at the pinned tag (>= {AITER_MIN_VERSION}); see docs/rocm/backends.md " + "for the recipe. Build by tag, never a branch: master's C ABI does not " "match the structs vendored here." ) from exc From c50dac0785a827d173d6299415aed2e475d7164b Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:29:34 -0400 Subject: [PATCH 27/41] fix(rocm): keep the installed AITER tree read-only, give the fallback a HOME Copilot is right that the previous two commits were solving this the wrong way round. Making aiter/jit world-writable puts a writable directory on Python's import path, and the sticky bit does not close it: +t stops one uid unlinking another's file, but nothing stops a uid *creating* a module_*.so that has not been built yet, which the next process then dlopens. The reason for making it writable at all does not survive measurement. AITER's fallback copies the tree to ~/.aiter and returns that, so the prebuilt artifacts are found either way -- the copy is 226 MB and takes 97 ms, against the 815 s the prebuild saves. It was never worth a writable code path. What it does need is somewhere to land: docker gives an unknown --user uid HOME=/, which is not writable, so the fallback died rather than copying. A sticky world-writable /home/devuser fixes that, and ~/.aiter is not on sys.path the way the package tree is. Verified on a derived image, foreign uid 12345, installed tree left at 0755: installed jit mode: 0755 writable: False resolved jit dir -> /home/devuser/.aiter/jit (1.51s) artifacts visible: 41 Co-Authored-By: Claude --- docker/Dockerfile.rocm | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7e123436289..10eec2847b4 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -239,17 +239,16 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # in-process, under a process-global lock, minutes at a time. The driver ends in # its own --check, so a missing or empty artifact fails the build here. # -# The chmod is load-bearing: AITER's get_user_jit_dir() falls back to copying the -# whole jit tree into $HOME when its own directory is unwritable, so a container -# run with --user of any other uid would miss every artifact built here -- or die -# outright when HOME is not writable either. Directories go a+wt rather than a+w: -# write alone lets any uid unlink a kernel and drop in its own, whatever the -# file's own mode says; the sticky bit leaves that to the owner. RUN set -eu; \ export GPU_ARCHS="$AITER_GPU_ARCHS"; \ - python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ - JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ - chmod -R a+rX "$JIT"; find "$JIT" -type d -exec chmod a+wt {} + + python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" + +# AITER copies its jit tree to ~/.aiter when the installed one is not writable. +# Docker hands an unknown --user uid HOME=/, which is not writable, so that +# fallback dies rather than copying. A sticky world-writable home lets it land; +# the installed tree stays read-only and off any writable import path. +ENV HOME=/home/$USERNAME +RUN chmod 1777 /home/$USERNAME WORKDIR /home/$USERNAME From a977c156e66bde1c6cbc9c5e3f11d88fc7f82781 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:29:34 -0400 Subject: [PATCH 28/41] test(rocm): make the token-order check actually check order The test claimed to catch "a renamed or reordered token" and only caught the first: _md_tokens returned a set, so the decomposition proved membership and said nothing about position. If AITER moves _mask ahead of _nbias, every prebuilt filename for that arm becomes undiscoverable while the test stays green -- the whole failure this file exists to prevent. Tokens now carry the index of their first appearance in AITER's own source (findall preserves order), and the decomposition asserts those indices never go backwards. A/B, swapping _nbias and _mask in the driver's batch-prefill composer: AssertionError: 'mha_batch_prefill_bf16_nlogits_nmask_nbias_nlse_...': '_nbias' is out of AITER's order -- it appears at index 5 after a token at 7 The set version passed that same swap. Co-Authored-By: Claude --- tests/rocm/test_prebuild_aiter_attention.py | 34 +++++++++++++-------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index b40a78ab0dc..944fa239231 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -217,35 +217,45 @@ def _mha_recipes_source() -> str: return src.read_text() -def _md_tokens(family: str) -> set: - """Every name token AITER can emit for this family, dtype excluded. +def _md_token_order(family: str) -> dict: + """Name token -> position of its first appearance in AITER's own source. - varlen delegates its suffix to `compose_mha_fwd_variant_suffix_and_filter`, - so its literals live in mha_recipes.py rather than the composer body. + findall preserves source order, so the index doubles as the token's slot in + the composed name. varlen delegates its suffix to + `compose_mha_fwd_variant_suffix_and_filter`, so its literals live in + mha_recipes.py rather than in the composer body. """ source = ( _mha_recipes_source() if family == "mha_varlen_fwd" else _composer_body(family) ) - return set(re.findall(r'"(_[a-z0-9]+)"', source)) + order = {} + for i, tok in enumerate(re.findall(r'"(_[a-z0-9]+)"', source)): + order.setdefault(tok, i) + return order @pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) def test_every_md_name_decomposes_into_aiter_token_literals(family): - """Catches a renamed or reordered token: each name we compose must be the - family base plus a concatenation of literals AITER itself appends.""" - tokens = _md_tokens(family) - assert tokens, f"no md_name += literals found in cmdGenFunc_{family}" + """Each composed name must be the family base plus AITER's own literals, in + AITER's own order. Membership alone is not enough: a reordered token leaves + every filename for that arm undiscoverable while still decomposing.""" + order = _md_token_order(family) + assert order, f"no md_name += literals found in cmdGenFunc_{family}" for v in driver.reachable_variants(): if v.family != family: continue assert v.md_name.startswith(f"{family}_{v.dtype}") rest = v.md_name[len(f"{family}_{v.dtype}") :] + seen = -1 while rest: - match = max( - (t for t in tokens if rest.startswith(t)), key=len, default=None - ) + match = max((t for t in order if rest.startswith(t)), key=len, default=None) assert match, f"{v.md_name!r}: no aiter token matches at {rest!r}" + assert order[match] >= seen, ( + f"{v.md_name!r}: {match!r} is out of AITER's order -- it appears " + f"at index {order[match]} after a token at {seen}" + ) + seen = order[match] rest = rest[len(match) :] From e84279f9afdad03d8d5320d10950109df05f9618 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:29:34 -0400 Subject: [PATCH 29/41] docs(rocm): install AITER's build requirements before the no-isolation build The recipe says it is what docker/Dockerfile.rocm performs, then drops the layer that makes it work: --no-build-isolation installs no build requirements, and the Dockerfile has already put setuptools, setuptools-scm and packaging in the image by the time it clones. Following the doc in a clean environment fails on setuptools_scm before it compiles anything. Co-Authored-By: Claude --- docs/rocm/backends.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index f5377168371..c398e4913bc 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -151,6 +151,9 @@ That says nothing about their ROCm support either way. `docker/Dockerfile.rocm` performs: ```bash +# --no-build-isolation means pip installs no build requirements for you. +pip install "setuptools>=80" "setuptools-scm>=9.2" "packaging>=24" + git clone --recursive --depth 1 --shallow-submodules \ --branch v0.1.21.post2 https://github.com/ROCm/aiter.git cd aiter && PREBUILD_KERNELS=0 GPU_ARCHS="gfx942;gfx950" \ From a4daaa550d2ff560f8575d282d2478e15ac7671e Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:52:04 -0400 Subject: [PATCH 30/41] fix(rocm): settle HOME before the prebuild resolves AITER's jit directory ENV HOME landed after the prebuild RUN, so the layer that resolves AITER's jit directory ran under whatever HOME the base image left. It happens not to bite -- the build user owns the installed tree, so get_user_jit_dir() takes the writable-in-place branch and the built image has all 49 artifacts there -- but the ordering made the artifacts' destination depend on a variable set two layers later. Set before the prebuild, where it is read. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 10eec2847b4..77d337ee8aa 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -232,6 +232,14 @@ if torch.version.hip is None: print(f"amd-aiter {got}; torch {torch.__version__} (hip={torch.version.hip})") PY +# AITER copies its jit tree to ~/.aiter when the installed one is not writable. +# Docker hands an unknown --user uid HOME=/, which is not writable, so that +# fallback dies rather than copying. A sticky world-writable home lets it land; +# the installed tree stays read-only and off any writable import path. Set before +# the prebuild, which resolves the same directory. +ENV HOME=/home/$USERNAME +RUN chmod 1777 /home/$USERNAME + # Placed after the install so editing the variant list re-runs only the prebuild. COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuild_aiter_attention.py @@ -243,13 +251,6 @@ RUN set -eu; \ export GPU_ARCHS="$AITER_GPU_ARCHS"; \ python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" -# AITER copies its jit tree to ~/.aiter when the installed one is not writable. -# Docker hands an unknown --user uid HOME=/, which is not writable, so that -# fallback dies rather than copying. A sticky world-writable home lets it land; -# the installed tree stays read-only and off any writable import path. -ENV HOME=/home/$USERNAME -RUN chmod 1777 /home/$USERNAME - WORKDIR /home/$USERNAME CMD ["/bin/bash"] From fb6bc3f14213bc3335627bdb1d669f83af215a5a Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:52:04 -0400 Subject: [PATCH 31/41] docs(rocm): the 0.1.21 floor is the norm/rope ABI move, not a struct change Three places said the vendored structs stop matching below 0.1.21 while the header itself said they are field-for-field identical to 0.1.20's. The header was right -- the AITER sources the shims compile against are byte-identical across the two revisions: mha_fwd.h 48cf8eac128502778b85dd3b79735c26 (both) mha_batch_prefill.h 1b82ad6ece78841892bbce5864cbc6b5 (both) fmha_fwd.hpp d5471e0ba156dbb7a70a04cc18e97bdf (both) The floor is still required, for the reason the migration commit gives: 0.1.21 moved rmsnorm and rope to the POD tensor, so those shims cannot resolve their symbols against 0.1.20. Stating it as a struct-layout change would send the next person re-vendoring headers that did not move. Co-Authored-By: Claude --- CLAUDE.md | 5 +++-- tests/rocm/test_aiter_version_gate.py | 8 ++++---- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 3cc37b94d1e..37a666f27a3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -114,8 +114,9 @@ gfx950 at `-n 1`, so 92% of a first run is compilation. Clearing **AITER is a source build at a pinned tag.** `docker/Dockerfile.rocm` clones `v0.1.21.post2` and installs it non-editable with `PREBUILD_KERNELS=0`, so the image needs no separate install. `aiter_utils.AITER_MIN_VERSION` (0.1.21) is -the hard floor: below it the vendored struct layouts stop matching, and 0.1.21 -moved `rmsnorm` and `rope` from `at::Tensor` to the POD `aiter_tensor_t`. +the hard floor: 0.1.21 moved `rmsnorm` and `rope` from `at::Tensor` to the POD +`aiter_tensor_t`, so the shims cannot resolve their symbols against anything +older. The vendored prefill structs are unchanged from 0.1.20. Build it by **tag**, never a branch — a clone that loses its tag makes `setuptools_scm` invent a `.dev` version that falls under the floor, and every diff --git a/tests/rocm/test_aiter_version_gate.py b/tests/rocm/test_aiter_version_gate.py index 4ec4ae2462d..3bbad55ac0f 100644 --- a/tests/rocm/test_aiter_version_gate.py +++ b/tests/rocm/test_aiter_version_gate.py @@ -3,10 +3,10 @@ # SPDX-License-Identifier: Apache-2.0 """The amd-aiter ABI floor. -The vendored structs under include/flashinfer/rocm/attention/aiter/ follow the 0.1.21 -layout and travel by value through dlsym'd pointers, so an older AITER shifts -field offsets instead of failing to load. Nothing downstream can detect that, -which is why the floor is enforced before routing rather than at the call. +The floor is enforced before routing rather than at the call: the vendored structs +travel by value through dlsym'd pointers, so a layout change shifts field offsets +instead of failing to load. The 0.1.21 bump is the norm/rope POD ABI move -- the +prefill structs are byte-identical to 0.1.20's. """ import pathlib From 4514b8a6ef7796119e5bbfc9980daf9e882059e6 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 22:52:04 -0400 Subject: [PATCH 32/41] test(rocm): pin the direction literal, the Dockerfile floor, and the fp8 range Three gaps, each A/B'd by breaking the thing it now catches. The receipt was pinned against AITER's source but the `-d` literal beside it was compared to nothing, so a typo there emits an empty instance set that still links and still passes --check: AssertionError: ... generate.py -d batch_prefil --receipt 200 ... varlen is excluded because it takes its command from AITER's own helper. The Dockerfile's FLOOR duplicated AITER_MIN_VERSION with nothing holding them in step; drifting it to 0.1.20 now fails. That is the pairing where an image builds green against an AITER the library then refuses at run time. The fp8 test hard-coded 240.0, which is e4m3fnuz's maximum, and applied it on gfx950 where the encoding is e4m3fn and the maximum is 448 -- wasting half the range and making the same tolerance a tighter gate on one architecture than the other. Now torch.finfo(dtype).max. It also re-implemented the encoding probe instead of importing prefill.py's, so the test could disagree with routing about which encoding the board uses. Co-Authored-By: Claude --- tests/rocm/test_fp8_paged_prefill_aiter.py | 30 ++++++++++++--------- tests/rocm/test_prebuild_aiter_attention.py | 30 +++++++++++++++++++++ 2 files changed, 48 insertions(+), 12 deletions(-) diff --git a/tests/rocm/test_fp8_paged_prefill_aiter.py b/tests/rocm/test_fp8_paged_prefill_aiter.py index 710cf86049c..33cbe5f55fa 100644 --- a/tests/rocm/test_fp8_paged_prefill_aiter.py +++ b/tests/rocm/test_fp8_paged_prefill_aiter.py @@ -21,6 +21,7 @@ from flashinfer.rocm.prefill import ( FP8_PREFILL_OUT_DTYPE, _aiter_paged_route_page_sizes, + _native_fp8_dtype, ) from tests.test_helpers.test_helpers import requires_aiter @@ -28,17 +29,19 @@ NUM_QO_HEADS, NUM_KV_HEADS, HEAD_DIM = 8, 2, 128 PAGE, KV_LEN, Q_LEN, BATCH = 16, 512, 64, 2 -FP8_MAX = 240.0 -def _native_fp8_dtype(): - """This GPU's fp8 encoding. The two are indistinguishable by shape or by the - .so name, and the wrong one is read under the wrong bias and returns NaN.""" - aiter = pytest.importorskip("aiter") - try: - return aiter.dtypes.fp8 - except AttributeError: +def _fp8_dtype(): + """This GPU's fp8 encoding, from the same probe routing uses. + + The two encodings are indistinguishable by shape or by the .so name, and the + wrong one is read under the wrong exponent bias and returns NaN. + """ + pytest.importorskip("aiter") + dtype = _native_fp8_dtype() + if dtype is None: pytest.skip("aiter.dtypes.fp8 is unreadable, so the encoding is unknown") + return dtype @pytest.fixture(scope="module") @@ -71,7 +74,7 @@ def _plan_and_run(workspace, q, kv, backend, **kw): def test_fp8_paged_prefill_matches_a_bf16_reference(workspace): - fp8 = _native_fp8_dtype() + fp8 = _fp8_dtype() if PAGE not in _aiter_paged_route_page_sizes(fp8): pytest.skip(f"page_size={PAGE} does not route to the native paged kernel") @@ -92,8 +95,11 @@ def test_fp8_paged_prefill_matches_a_bf16_reference(workspace): ref = _plan_and_run(workspace, q, kv, "fa2") - sq = q.abs().amax().float() / FP8_MAX - sk = kv.abs().amax().float() / FP8_MAX + # From the dtype, not a constant: e4m3fnuz maxes at 240 and e4m3fn at 448, + # so a fixed divisor wastes half the range on one of the two architectures. + fp8_max = torch.finfo(fp8).max + sq = q.abs().amax().float() / fp8_max + sk = kv.abs().amax().float() / fp8_max got = _plan_and_run( workspace, (q.float() / sq).to(fp8), @@ -114,7 +120,7 @@ def test_fp8_paged_prefill_matches_a_bf16_reference(workspace): def test_fp8_prefill_is_refused_on_fa2(workspace): """fa2 has no fp8 kernel, so the refusal must be an error, not a fallback.""" - fp8 = _native_fp8_dtype() + fp8 = _fp8_dtype() dev = "cuda:0" q = torch.zeros( BATCH * Q_LEN, NUM_QO_HEADS, HEAD_DIM, device=dev, dtype=torch.bfloat16 diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 944fa239231..207daca59a1 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -154,6 +154,28 @@ def test_every_name_token_the_loader_can_emit_is_one_the_driver_builds(): assert not missing, f"aiter_loader.cc can emit {sorted(missing)}, never built" +class _FakeCore: + """build_recipe reads only CK_DIR off core for the two composed families.""" + + CK_DIR = "/nonexistent" + + +def test_the_dockerfile_floor_matches_aiter_min_version(): + """Two literals, one meaning. If they drift the image builds green against an + AITER the library then refuses, and every path degrades to fa2 at run time.""" + from flashinfer.rocm import aiter_utils + + dockerfile = ( + Path(__file__).resolve().parents[2] / "docker" / "Dockerfile.rocm" + ).read_text() + found = re.findall(r'^FLOOR = "([0-9][^"]*)"', dockerfile, re.M) + assert found, "docker/Dockerfile.rocm no longer spells FLOOR; is this test stale?" + assert found == [aiter_utils.AITER_MIN_VERSION], ( + f"Dockerfile FLOOR {found} != AITER_MIN_VERSION " + f"{aiter_utils.AITER_MIN_VERSION!r}" + ) + + # --------------------------------------------------------------------------- # 2. receipts and token spellings match the installed aiter # --------------------------------------------------------------------------- @@ -180,6 +202,14 @@ def test_the_receipt_and_direction_match_aiter(family, direction, receipt): }[family] assert from_driver == receipt + # The receipt was pinned but the -d literal beside it was not: a typo there + # emits an empty instance set that still links and still passes --check. + # varlen is excluded because it takes its command from AITER's own helper. + if family != "mha_varlen_fwd": + v = next(x for x in driver.reachable_variants() if x.family == family) + cmd = driver.build_recipe(v, _FakeCore())["blob_gen_cmd"][0] + assert f"-d {direction} " in cmd, cmd + def test_mha_fwd_masks_on_underscore_m_not_underscore_mask(): """The one filter token that differs between families. Spelling `_mask*` here From d56e274296ce06f459d53608326507154a1283cc Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 23:09:19 -0400 Subject: [PATCH 33/41] fix(rocm): stop widening a shared directory, and correct two stated numbers Third round on the same theme, so this stops solving it in the image. 1777 on /home/devuser is the same defect as the world-writable jit tree it replaced, and arguably broader: HOME is the WORKDIR and contributes ~/.local user-site packages, so any uid could plant a module a later process imports. Nothing in the image is group- or world-writable now. A --user run whose uid is not USER_UID needs a writable HOME of its own, which the Dockerfile documents; that is also already the guidance for AITER work, where --user is omitted precisely because it makes AITER's jit directory unwritable. The fallback this was protecting costs 226 MB and 97 ms when it does run. CONTRIBUTING said ~23 min for the prebuild while this PR measured 28 min for 40 variants, before fp8 and the whole-module builds; now ~35 min. The tutorial said fp8 serves page_size 16, where _aiter_paged_route_page_sizes gives it all three of {1, 16, 1024}. Co-Authored-By: Claude --- CONTRIBUTING.md | 2 +- docker/Dockerfile.rocm | 13 ++++++------- examples/amd_flashinfer_rocm_tutorial.ipynb | 2 +- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index cdf41dc9fc3..af5389602b4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -23,7 +23,7 @@ docker build -t flashinfer-dev:rocm10.0 -f docker/Dockerfile.rocm . 10.0, 24.04, 3.12, and 2.12.0. They select the `rocm/pytorch` base image tag, so they are not independent knobs — any override has to name a tag that exists on Docker Hub. `AITER_REF` selects the AITER tag to build from source, and -`AITER_PREBUILD_JOBS` how many variants compile at once (~23 min at 2 on 32 +`AITER_PREBUILD_JOBS` how many variants compile at once (~35 min at 2 on 32 cores). `AITER_SOURCE=0` falls back to the `AITER_VERSION`/`AITER_INDEX` wheel install, which is an unblock only: no `amd-aiter` wheel targets ROCm 10, so it installs a retarget of another ROCm. Do not raise `TORCH_VERSION` to 2.13 — it drops a `c10` symbol AITER's diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 77d337ee8aa..afcbf26afb5 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -232,13 +232,12 @@ if torch.version.hip is None: print(f"amd-aiter {got}; torch {torch.__version__} (hip={torch.version.hip})") PY -# AITER copies its jit tree to ~/.aiter when the installed one is not writable. -# Docker hands an unknown --user uid HOME=/, which is not writable, so that -# fallback dies rather than copying. A sticky world-writable home lets it land; -# the installed tree stays read-only and off any writable import path. Set before -# the prebuild, which resolves the same directory. -ENV HOME=/home/$USERNAME -RUN chmod 1777 /home/$USERNAME +# Nothing here is made group- or world-writable. AITER falls back to copying its +# jit tree to ~/.aiter when the installed one is not writable, so a --user run +# with a uid that is not $USER_UID needs a writable HOME of its own: +# docker run --user "$(id -u)" -e HOME=/work -v "$PWD":/work ... +# Widening a shared directory instead puts a writable path on Python's import +# path, which is worse than the copy it avoids (226 MB, 97 ms). # Placed after the install so editing the variant list re-runs only the prebuild. COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuild_aiter_attention.py diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 596cd03061e..06c6da7e5a1 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -164,7 +164,7 @@ "\n", "**Backend integration.** amd-flashinfer offers [AITER](https://github.com/ROCm/aiter)—AMD’s SOTA kernel library for large language model operators on ROCm—as an optional **`backend=\"aiter\"`**. When the `aiter` package is present and the requested operator is integrated with AITER in this port, work **delegates** to AITER; otherwise the default HIP path applies. The following cells exercise the AITER-backed prefill path.\n", "\n", - "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the native kernel does serve `page_size=1024` (and `16` for fp8); other sizes gather." + "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the native kernel does serve `page_size=1024`, and `1`, `16` or `1024` for fp8; other sizes gather." ] }, { From a27cf1167f91e1d3f88b3fec5db63d1104778232 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Mon, 14 Sep 2026 23:23:51 -0400 Subject: [PATCH 34/41] fix(rocm): pin where the prebuild writes, and scan the loader's call sites too Removing ENV HOME in d56e27429 put back the ordering hazard a4daaa550 had fixed: the prebuild resolved AITER's jit directory under whatever HOME the base image left. Rather than manage HOME, assert the outcome. The driver already prints the directory it resolved, so the build now fails loudly if the artifacts land anywhere but the installed tree -- which is the only place a runtime container will look for them. Nothing is made writable to achieve it. The version-assertion heredoc is the one BuildKit-only instruction in the file, so it now declares the frontend instead of relying on the daemon default. _loader_name_tokens() scanned only between dtype_token and build_so_name, which excludes the prefix/infix/suffix literals the three call sites pass in -- the very place _nsink lives. A/B, dropping _nsink from the driver: AssertionError: aiter_loader.cc can emit ['nsink'], never built The old range could not see that at all. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 10 +++++++++- tests/rocm/test_prebuild_aiter_attention.py | 16 ++++++++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index afcbf26afb5..b47ac673959 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,3 +1,4 @@ +# syntax=docker/dockerfile:1.7 # ROCm development container for amd-flashinfer. # # The base image already provides ROCm, Python and a HIP build of PyTorch inside @@ -248,7 +249,14 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # RUN set -eu; \ export GPU_ARCHS="$AITER_GPU_ARCHS"; \ - python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" + python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ + JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ + case "$JIT" in \ + */site-packages/aiter/jit) echo "prebuilt into $JIT" ;; \ + *) echo "ERROR: artifacts went to $JIT, not the installed tree; a runtime" \ + "container resolves a different directory and would rebuild all of them" >&2; \ + exit 1 ;; \ + esac WORKDIR /home/$USERNAME diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 207daca59a1..41ab7704920 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -129,12 +129,24 @@ def test_check_does_not_demand_what_only_never_builds(tmp_path, monkeypatch): def _loader_name_tokens() -> set: - """String literals `build_so_name`/`dtype_token` can put in a filename.""" + """Every string literal the loader can put in a filename. + + Two sources, and missing the second is how `_nsink` could have slipped in: + the axis tokens inside `build_so_name`/`dtype_token`, and the + prefix/infix/suffix literals its three call sites pass in. + """ text = _LOADER_CC.read_text() start = text.find("const char* dtype_token(") end = text.find("\n}", text.find("std::string build_so_name(")) assert start != -1 and end != -1, "aiter_loader.cc: name composition not found" - return set(re.findall(r'"([a-z0-9_]+)"', text[start:end])) + tokens = set(re.findall(r'"([a-z0-9_]+)"', text[start:end])) + + call_sites = re.findall(r"build_so_name\(key,([^;]*?)\)", text, re.S) + assert call_sites, "aiter_loader.cc: no build_so_name call sites found" + for args in call_sites: + for lit in re.findall(r'"([a-z0-9_.]*)"', args): + tokens.update(s for s in lit.replace(".so", "").split("_") if s) + return tokens def test_every_name_token_the_loader_can_emit_is_one_the_driver_builds(): From e441b656e89ea7f872c26bbdd6794dfd26998c60 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 00:21:39 -0400 Subject: [PATCH 35/41] fix(rocm): drop the fp8 LSE arm, and detect an artifact with no kernels in it Both suppressed comments from the third review, and the first proves the second. fp8 was enumerated over both LSE arms. It has none: prefill.py:3334 raises "fp8 prefill cannot produce LSE: AITER builds no LSE instance of the fp8 kernel at any page size", and prefill.py:1019 already excluded it for the store route. The driver did not, so four of the eight fp8 builds asked CK for an instance set that does not exist. They compiled, linked, and passed --check: 518,440 ..._logits_nbias_mask_lse_...pertensor_nsink.so <- all four LSE 518,440 ..._logits_nbias_nmask_lse_... 5,130,416 ..._nlogits_nbias_nmask_nlse_... <- all four non-LSE 5,437,120 ..._logits_nbias_mask_nlse_... Ten times smaller, because they are a dispatcher with nothing behind it. 44 variants now, not 48. That is the failure mode three reviews called speculative, so _check now looks for it: an artifact carrying fewer than 200 `ck_tile` markers has no instances. Measured on the two arms of the same variant -- 20 markers against 8,325 -- which is what sets the threshold. Applied to the CK variants only; the LOADER_MODULES are built -DENABLE_CK=0 and legitimately have none. Also bounds each child's ninja parallelism. AITER computes MAX_JOBS from ~80% of *all* host CPUs and its own free-memory snapshot, per process, so --jobs N oversubscribed the box N-fold; the children now divide one budget. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 44 ++++++++++++++++++++- tests/rocm/test_prebuild_aiter_attention.py | 36 +++++++++++++++-- 2 files changed, 75 insertions(+), 5 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 4ed142fa979..86db77ab7d1 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -178,14 +178,20 @@ def _filter_batch_prefill(v: Variant) -> str: def reachable_variants() -> Tuple[Variant, ...]: - """Every ``.so`` the loader can ask for: 8 + 16 + 24 = 48 per architecture.""" + """Every ``.so`` the loader can ask for: 8 + 16 + 20 = 44 per architecture. + + fp8 contributes 4, not 8: AITER ships no LSE instance of the fp8 kernel at + any page size, and prefill.py raises on the combination rather than + dispatching it. + """ out: List[Variant] = [] for family, has_logits_axis in _HAS_LOGITS_AXIS.items(): logits_values = (False, True) if has_logits_axis else (False,) for dtype in _FAMILY_DTYPES[family]: + lse_values = (False,) if dtype == "fp8bf16" else (False, True) for has_logits_cap in logits_values: for needs_mask in (False, True): - for has_lse in (False, True): + for has_lse in lse_values: out.append( Variant(family, dtype, has_logits_cap, needs_mask, has_lse) ) @@ -394,6 +400,19 @@ def _run_jobs(variants: Sequence[Variant], jobs: int, extra: Sequence[str] = ()) return 1 if failed else 0 +def _child_env(jobs: int) -> Dict[str, str]: + """Divide ninja's parallelism across the children rather than per child. + + Left alone, each child sets MAX_JOBS from ~80% of *all* host CPUs and a + free-memory snapshot it took independently, so --jobs N oversubscribes the + box N-fold. AITER honours MAX_JOBS when it is already set. + """ + env = dict(os.environ) + if "MAX_JOBS" not in env: + env["MAX_JOBS"] = str(max(1, int((os.cpu_count() or 1) * 0.8) // max(1, jobs))) + return env + + def _drain(pending, running, failed, jobs, total) -> None: done = 0 while pending or running: @@ -411,6 +430,7 @@ def _drain(pending, running, failed, jobs, total) -> None: ], stdout=fd, stderr=subprocess.STDOUT, + env=_child_env(jobs), ) finally: os.close(fd) # Popen dup'd it; holding ours would exhaust the table @@ -456,6 +476,17 @@ def _decode(s: str) -> Variant: return Variant(family, dtype, bool(int(lc)), bool(int(mask)), bool(int(lse))) +_MIN_CK_MARKERS = 200 + + +def _has_ck_instances(path: Path) -> bool: + """Does this artifact carry CK-tile kernel instances, or only a dispatcher? + + Only meaningful for the CK variants: LOADER_MODULES are built -DENABLE_CK=0. + """ + return path.read_bytes().count(b"ck_tile") >= _MIN_CK_MARKERS + + def _check(variants: Sequence[Variant], extra: Sequence[str]) -> int: """Assert every selected artifact is present and non-empty. @@ -470,6 +501,15 @@ def _check(variants: Sequence[Variant], extra: Sequence[str]) -> int: if missing: print("missing or empty:\n " + "\n ".join(missing)) return 1 + # A --filter that selected no instances still compiles and links, so presence + # is not enough: the result is a dispatcher with no kernel behind it. + hollow = [v.so_name for v in variants if not _has_ck_instances(d / v.so_name)] + if hollow: + print( + f"{len(hollow)} artifact(s) carry no CK instances -- the filter " + "selected nothing:\n " + "\n ".join(hollow) + ) + return 1 return 0 diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 41ab7704920..5378051f053 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -84,15 +84,28 @@ def _fi_so_names() -> set: def test_the_driver_covers_flashinfers_table_and_the_fp8_arm_it_omits(): """flashinfer's table is the store's domain, which has no fp8 arm; the driver writes where the loader looks, so it must carry fp8 as well.""" - assert len(driver.reachable_variants()) == 48 + assert len(driver.reachable_variants()) == 44 assert _fi_so_names() < {v.so_name for v in driver.reachable_variants()} +def test_fp8_carries_no_lse_arm(): + """AITER ships no LSE instance of the fp8 kernel, and prefill.py raises on + the combination -- building it yields a dispatcher with no kernel behind it.""" + fp8 = [v for v in driver.reachable_variants() if v.dtype == "fp8bf16"] + assert len(fp8) == 4 + assert not [v for v in fp8 if v.has_lse] + # The guard this mirrors, so a change there shows up here. + prefill = ( + Path(__file__).resolve().parents[2] / "flashinfer" / "rocm" / "prefill.py" + ).read_text() + assert "fp8 prefill cannot produce LSE" in prefill + + def test_nothing_is_trimmed_from_the_built_set(): """A trim buys a few minutes of image build for a multi-minute stall on the arm that turns out to be missing.""" assert set(driver.selected_variants()) == set(driver.reachable_variants()) - assert len(driver.selected_variants()) == 48 + assert len(driver.selected_variants()) == 44 @pytest.mark.parametrize("family", sorted(_FAMILY_BY_NAME)) @@ -112,13 +125,30 @@ def test_check_does_not_demand_what_only_never_builds(tmp_path, monkeypatch): monkeypatch.setattr(driver, "jit_dir", lambda _core: tmp_path) picked = driver.selected_variants(only=["mha_fwd"]) for v in picked: - (tmp_path / v.so_name).write_bytes(b"x") + # Enough markers to clear the instance check; that is a separate test. + (tmp_path / v.so_name).write_bytes(b"ck_tile" * driver._MIN_CK_MARKERS) assert driver.main(["--check", "--only", "mha_fwd"]) == 0 # Unrestricted, the same directory is incomplete: the module is absent. assert driver.main(["--check"]) == 1 +def test_check_rejects_an_artifact_with_no_ck_instances(tmp_path, monkeypatch): + """A --filter that selected nothing still compiles and links. Measured on a + real one: 20 `ck_tile` markers with no instances, 1701 with them.""" + monkeypatch.setattr(driver, "_aiter_jit_core", lambda: object()) + monkeypatch.setattr(driver, "jit_dir", lambda _core: tmp_path) + picked = driver.selected_variants(only=["mha_fwd"]) + for v in picked: + (tmp_path / v.so_name).write_bytes(b"ck_tile" * driver._MIN_CK_MARKERS) + + assert driver.main(["--check", "--only", "mha_fwd"]) == 0 + + hollow = tmp_path / picked[0].so_name + hollow.write_bytes(b"ck_tile" * (driver._MIN_CK_MARKERS - 1)) + assert driver.main(["--check", "--only", "mha_fwd"]) == 1 + + # --------------------------------------------------------------------------- # 1b. every token the loader can compose is one the driver emits # --------------------------------------------------------------------------- From f3b52e1aef6ea08d3e8b8636f426dcd8b47e8db2 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 00:34:29 -0400 Subject: [PATCH 36/41] fix(rocm): put the gfx950 paged fault where routing can see it, and four others The memory fault this branch found while benchmarking was recorded only in the benchmark skill, but _aiter_paged_route_page_sizes sends bf16/fp16 to exactly the page size that faults, so the note now sits beside that return. Not gated: it reproduces on 0.1.20 as well, and the failing kv boundary has not been bounded -- gating on a guess would disable a path that works below it. _child_env deferred to a pre-set MAX_JOBS, which CLAUDE.md documents as the way to limit a build, so the documented knob defeated the split it was added for. It now divides that budget: MAX_JOBS=4 with --jobs 2 gives each child 2. _aiter_importable no longer probes for AITER's header tree. It answers "can AITER run?", and a shim that is already built or AOT-cached needs no headers; the build path raises with the real cause when they are what is missing. Raised in four separate reviews, each time correctly. The Dockerfile now resolves and checks the jit destination before the 35-minute prebuild rather than after it, _has_ck_instances reads in 1 MB chunks instead of pulling a whole multi-MB artifact into memory, and compose_cache_tag rejects the same character class its test asserts on. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 10 +++++----- docker/prebuild_aiter_attention.py | 24 ++++++++++++++++-------- flashinfer/jit/rocm/aiter_source.py | 4 +++- flashinfer/rocm/aiter_utils.py | 9 +++------ flashinfer/rocm/prefill.py | 4 ++++ 5 files changed, 31 insertions(+), 20 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b47ac673959..7ae35056029 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -249,14 +249,14 @@ COPY --chown=$USER_UID:$USER_GID docker/prebuild_aiter_attention.py /opt/prebuil # RUN set -eu; \ export GPU_ARCHS="$AITER_GPU_ARCHS"; \ - python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS"; \ JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ case "$JIT" in \ - */site-packages/aiter/jit) echo "prebuilt into $JIT" ;; \ - *) echo "ERROR: artifacts went to $JIT, not the installed tree; a runtime" \ - "container resolves a different directory and would rebuild all of them" >&2; \ + */site-packages/aiter/jit) echo "prebuilding into $JIT" ;; \ + *) echo "ERROR: AITER resolves its jit dir to $JIT, not the installed tree;" \ + "a runtime container would look elsewhere and rebuild everything" >&2; \ exit 1 ;; \ - esac + esac; \ + python3 /opt/prebuild_aiter_attention.py --jobs "$AITER_PREBUILD_JOBS" WORKDIR /home/$USERNAME diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 86db77ab7d1..08894064f8b 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -48,11 +48,9 @@ GENERATE_PY = "example/ck_tile/01_fmha/generate.py" -# Whole AITER modules with no variant axes. The wheel shipped these prebuilt; a -# PREBUILD_KERNELS=0 source install ships none, so the image builds them or the -# first caller pays for them -- fmha_v3_fwd throws at dlopen, the other two only -# stall. mla_asm is aiter.mla's, and aiter_core is on the common attention path. -# tests/rocm/test_prebuild_aiter_attention.py checks the first against the loader. +# Whole AITER modules, no variant axes. A PREBUILD_KERNELS=0 source install ships +# none, so the image builds them or the first caller pays: fmha_v3_fwd throws at +# dlopen, mla_asm and aiter_core only stall. LOADER_MODULES = ("module_fmha_v3_fwd", "module_aiter_core", "module_mla_asm") @@ -408,8 +406,11 @@ def _child_env(jobs: int) -> Dict[str, str]: box N-fold. AITER honours MAX_JOBS when it is already set. """ env = dict(os.environ) - if "MAX_JOBS" not in env: - env["MAX_JOBS"] = str(max(1, int((os.cpu_count() or 1) * 0.8) // max(1, jobs))) + try: + budget = int(env["MAX_JOBS"]) + except (KeyError, ValueError): + budget = int((os.cpu_count() or 1) * 0.8) + env["MAX_JOBS"] = str(max(1, budget // max(1, jobs))) return env @@ -484,7 +485,14 @@ def _has_ck_instances(path: Path) -> bool: Only meaningful for the CK variants: LOADER_MODULES are built -DENABLE_CK=0. """ - return path.read_bytes().count(b"ck_tile") >= _MIN_CK_MARKERS + marker, seen, tail = b"ck_tile", 0, b"" + with path.open("rb") as fh: + while chunk := fh.read(1 << 20): + seen += (tail + chunk).count(marker) + if seen >= _MIN_CK_MARKERS: + return True + tail = chunk[-(len(marker) - 1) :] + return False def _check(variants: Sequence[Variant], extra: Sequence[str]) -> int: diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 18cbc8a6a82..897d1a552b8 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -239,7 +239,9 @@ def compose_cache_tag(arch: Optional[str] = None) -> str: tag = f"{arch}__aiter-{version}__rocm-{_rocm_version()}" # The version strings come from package metadata and torch, neither of which # this module controls, so the composed tag gets the same check the arch did. - if tag != Path(tag).name or tag.startswith("."): + # Same character class the cache-tag test asserts on: the AITER and ROCm + # versions come from package metadata and torch, neither of which we control. + if tag != Path(tag).name or tag.startswith(".") or set(tag) & set("/:;, "): raise ValueError( f"refusing to build a cache directory name from {tag!r}: " f"not a single safe path component" diff --git a/flashinfer/rocm/aiter_utils.py b/flashinfer/rocm/aiter_utils.py index 76572af0c35..9b900f6d380 100644 --- a/flashinfer/rocm/aiter_utils.py +++ b/flashinfer/rocm/aiter_utils.py @@ -99,12 +99,9 @@ def _aiter_importable() -> bool: import aiter # noqa: F401 from aiter.jit import core as _core # noqa: F401 - # Not `import aiter_meta`: an editable AITER ships none, and - # jit.rocm.aiter_source._aiter_csrc_include_dir now falls back to - # AITER_CSRC_DIR. Requiring it here made that fallback unreachable. - from ..jit.rocm.aiter_source import _aiter_csrc_include_dir - - _aiter_csrc_include_dir() + # Deliberately no header-tree check: this answers "can AITER run?", and + # the shim may already be built or AOT-cached. The build path raises with + # the real cause when the headers are the thing that is missing. except Exception: return False return _aiter_version_supported() diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index 5b9b8803527..02d397e14d1 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -90,6 +90,10 @@ def _aiter_paged_route_page_sizes(dtype: torch.dtype) -> frozenset: faster than native at every batch size (docs/rocm/backends.md). Widening this for fp16/bf16 is a benchmark, not a one-line edit. """ + # Known, unfixed, and not ours: on gfx950 the native paged kernel faults + # (GPU memory fault in FmhaBatchPrefillWithPagedKVCacheKernel) past kv=1024 + # at page_size=1024, on amd-aiter 0.1.20 and 0.1.21.post2 alike. Not gated + # here because the failing kv boundary has not been bounded. native = _aiter_native_page_sizes() return native if dtype in FP8_PREFILL_DTYPES else native & {1024} From 714d451d947f75c50e57d95ed23af7c42e3697a5 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 00:53:08 -0400 Subject: [PATCH 37/41] fix(rocm): tell source installs how to get the asm module they do not have Making the source build the supported install left a gap the wheel did not have. PREBUILD_KERNELS=0 ships no AITER modules, and aiter_loader.cc dlopens module_fmha_v3_fwd directly -- there is no lazy JIT behind it, unlike the mha_fwd variants. The image is fine because the prebuild driver builds it; the standalone recipe in docs/rocm/backends.md stopped short, so a large enough single prefill fails at dlopen. The loader's hint still read "ships prebuilt in the amd-aiter wheel ... its absence means the wheel is incomplete", which sends a source user to reinstall the thing that cannot fix it. It now names the driver and the FLASHINFER_AITER_ASM_PREFILL=0 escape, which is read at include/flashinfer/rocm/attention/aiter/single_prefill.cuh:100. The fp8 test's docstring also still said eight variants and described --check as presence-only; both moved in e441b656e. Co-Authored-By: Claude --- csrc/rocm/aiter_loader.cc | 9 +++++---- docs/rocm/backends.md | 9 +++++++++ tests/rocm/test_fp8_paged_prefill_aiter.py | 7 +++---- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/csrc/rocm/aiter_loader.cc b/csrc/rocm/aiter_loader.cc index 50a78933f15..c6e10a084a1 100644 --- a/csrc/rocm/aiter_loader.cc +++ b/csrc/rocm/aiter_loader.cc @@ -298,10 +298,11 @@ void* get_aiter_mha_fwd_asm_handle() { "on CK Tile."); } return load_variant_sym(s_asm_mu, s_asm_cache, 0, kAsmModuleName, kMhaFwdSymbol, []() { - return " Hint: " + std::string(kAsmModuleName) + - " ships prebuilt in the amd-aiter wheel, so unlike the mha_fwd variants there is no " - "JIT build to trigger. Its absence means the wheel is incomplete or the JIT dir " - "points elsewhere." + + return " Hint: unlike the mha_fwd variants there is no JIT build to trigger for " + + std::string(kAsmModuleName) + + ". The supported source install (PREBUILD_KERNELS=0) ships no AITER modules, so it " + "must be built by docker/prebuild_aiter_attention.py; set " + "FLASHINFER_AITER_ASM_PREFILL=0 to stay on CK Tile instead." + kAbiPinNote; }); } diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index c398e4913bc..b74ffd9f8e8 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -170,6 +170,15 @@ loses its tag makes `setuptools_scm` invent a `.dev` version that falls under the floor below, and every AITER path silently degrades to `fa2`. Non-editable, not `setup.py develop`: the develop branch installs no `aiter_meta`, and `aiter_meta/csrc/include` is what every C++ shim compiles against. +`PREBUILD_KERNELS=0` leaves the install with **no AITER modules at all**, which +the image compensates for by running `docker/prebuild_aiter_attention.py`. A +standalone install must do the same, or single prefill fails at `dlopen` once a +call is large enough to take the asm arm: + +```bash +python3 docker/prebuild_aiter_attention.py --jobs 2 # or: export FLASHINFER_AITER_ASM_PREFILL=0 +``` + `PREBUILD_KERNELS=0` because `=1` needs a live GPU — and its prebuilt modules are no loss, since FlashInfer builds its own `lib.so` from the same sources. diff --git a/tests/rocm/test_fp8_paged_prefill_aiter.py b/tests/rocm/test_fp8_paged_prefill_aiter.py index 33cbe5f55fa..11588cde43b 100644 --- a/tests/rocm/test_fp8_paged_prefill_aiter.py +++ b/tests/rocm/test_fp8_paged_prefill_aiter.py @@ -8,10 +8,9 @@ route runs mha_varlen_fwd, which has no fp8 arm -- so AITER's native paged kernel is the whole of fp8 prefill support, and nothing else here exercised it. -It is also the only end-to-end check on the eight fp8 variants -docker/prebuild_aiter_attention.py composes: a CK ``--filter`` that matched no -instances still links, still passes the driver's own --check, and fails only -here, at dispatch. +It is also the only end-to-end check on the four fp8 variants +docker/prebuild_aiter_attention.py composes. The driver's ``--check`` now +rejects an artifact carrying no CK instances, but only this exercises one. """ import pytest From 724b9e7e042ce503fd652728b2354fc1f1757474 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 01:59:37 -0400 Subject: [PATCH 38/41] fix(rocm): delete a duplicate fp8 test, and make the install recipe actually work tests/rocm/test_fp8_paged_prefill.py already exists with 20 cases, including test_fp8_paged_prefill_matches_fp32_reference (a tighter reference than the bf16 one added here), test_every_routed_fp8_page_size_is_numerically_right across pages 1/16/1024, test_fp8_rejects_return_lse, and test_single_prefill_fp8_raises_instead_of_a_ninja_log. Both cases added in 77334430c have a direct equivalent there, so the new file was pure duplication on a compile-bound suite, and its docstring's claim that nothing else exercised fp8 paged prefill was simply wrong. Removed. The standalone prebuild recipe added in 714d451d9 omitted GPU_ARCHS, which the driver never sets -- the flat `from jit import core` import skips the aiter/__init__.py that would populate it, so every build asserts. Measured first-hand earlier in this branch: eight consecutive `FAIL` lines with the variable unset. Both the doc and the driver's usage line now set it. --jobs took int() at parser-construction time, so a non-numeric AITER_PREBUILD_JOBS crashed before --help could run; it is now validated in main() and reports the variable by name. _check runs even when a build failed, so one bad variant no longer hides the hollow artifacts among the other 43. Two stale comments: aiter_source's docstring still taught the at::Tensor forward declaration, which is the trap this branch removes, and the tutorial promised page_size=1024 without the gfx950 fault that prefill.py and the benchmark skill both record. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 15 ++- docs/rocm/backends.md | 3 +- examples/amd_flashinfer_rocm_tutorial.ipynb | 2 +- flashinfer/jit/rocm/aiter_source.py | 16 +-- tests/rocm/test_fp8_paged_prefill_aiter.py | 138 -------------------- 5 files changed, 22 insertions(+), 152 deletions(-) delete mode 100644 tests/rocm/test_fp8_paged_prefill_aiter.py diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index 08894064f8b..f19e6693098 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -13,7 +13,7 @@ deliberate copy of ``flashinfer.jit.rocm.aiter_variants``, which pulls torch; ``tests/rocm/test_prebuild_aiter_attention.py`` holds the two in step. - python3 docker/prebuild_aiter_attention.py --list|--check|--jobs 2 + GPU_ARCHS="gfx942;gfx950" python3 docker/prebuild_aiter_attention.py --jobs 2 """ from __future__ import annotations @@ -537,8 +537,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ) p.add_argument( "--jobs", - type=int, - default=int(os.environ.get("AITER_PREBUILD_JOBS", "1")), + # str, validated in main(): an int() here runs while the parser is being + # built, so a bad AITER_PREBUILD_JOBS would crash before --help works. + default=os.environ.get("AITER_PREBUILD_JOBS", "1"), help="concurrent build processes", ) p.add_argument( @@ -571,6 +572,10 @@ def main(argv: Optional[Sequence[str]] = None) -> int: if a.check: return _check(variants, extra) + try: + a.jobs = int(a.jobs) + except ValueError: + p.error(f"--jobs/AITER_PREBUILD_JOBS must be an integer, got {a.jobs!r}") if a.jobs < 1: p.error("--jobs must be >= 1") t0 = time.time() @@ -579,7 +584,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: f"\n{len(variants)} variant(s) + {len(extra)} module(s) in " f"{(time.time() - t0) / 60:.0f} min at --jobs {a.jobs}" ) - return rc or _check(variants, extra) + # Both, not `rc or`: a run that lost one variant should still report every + # other artifact that linked hollow, rather than one finding per rebuild. + return _check(variants, extra) or rc if __name__ == "__main__": diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index b74ffd9f8e8..44f623d4947 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -176,7 +176,8 @@ standalone install must do the same, or single prefill fails at `dlopen` once a call is large enough to take the asm arm: ```bash -python3 docker/prebuild_aiter_attention.py --jobs 2 # or: export FLASHINFER_AITER_ASM_PREFILL=0 +GPU_ARCHS="gfx942;gfx950" python3 docker/prebuild_aiter_attention.py --jobs 2 +# or, to skip it: export FLASHINFER_AITER_ASM_PREFILL=0 ``` `PREBUILD_KERNELS=0` because `=1` needs a live GPU — and its prebuilt modules diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 06c6da7e5a1..48487707178 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -164,7 +164,7 @@ "\n", "**Backend integration.** amd-flashinfer offers [AITER](https://github.com/ROCm/aiter)—AMD’s SOTA kernel library for large language model operators on ROCm—as an optional **`backend=\"aiter\"`**. When the `aiter` package is present and the requested operator is integrated with AITER in this port, work **delegates** to AITER; otherwise the default HIP path applies. The following cells exercise the AITER-backed prefill path.\n", "\n", - "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the native kernel does serve `page_size=1024`, and `1`, `16` or `1024` for fp8; other sizes gather." + "**Requirements and layout.** `amd-aiter` is built from source at tag `v0.1.21.post2`, which is also the hard floor; see [Backends on ROCm](https://github.com/AMD-Ecosystem/flashinfer/blob/amd-integration/docs/rocm/backends.md). These calls expect **NHD** key/value layout (`[sequence, heads, dim]`). For batched paged prefill, `plan()` probes AITER for a native page size and otherwise gathers K/V per run. On amd-aiter 0.1.21.post2 the native kernel does serve `page_size=1024`, and `1`, `16` or `1024` for fp8; other sizes gather. On gfx950 that native kernel faults past `kv=1024` at `page_size=1024` (reproducible on 0.1.20 too), so prefer a gathered page size there for long contexts." ] }, { diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 897d1a552b8..b2383e7d3b8 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -6,10 +6,11 @@ FlashInfer wraps AITER kernels by compiling a small ``csrc/rocm/*_aiter.cu`` shim that calls AITER's C++ entry point directly and links the symbol-visible AITER ``.so``. Prefer ``#include``-ing AITER's real header, so a signature change is a -compile error rather than a load-time ``undefined symbol``. Fall back to a -forward declaration only for headers that still pull in pybind11, which clashes -with FlashInfer's ``-DPy_LIMITED_API`` build (``rope.h``, ``rmsnorm.h`` as of -0.1.16); there ``torch::Tensor`` is ``at::Tensor``, so the linker still resolves. +compile error rather than a load-time ``undefined symbol``. A forward +declaration is the fallback for headers that still pull in pybind11, which +clashes with ``-DPy_LIMITED_API`` -- but declare the signature AITER actually +has: since 0.1.21 rmsnorm and rope take ``aiter_tensor_t&``, and the old +``at::Tensor`` spelling compiles and then fails at ``dlopen`` on the mangled name. AITER's installed wheel builds its modules with ``-fvisibility=hidden``, so the kernel symbols (e.g. ``rope_cached_positions_2c_fwd_impl``) are not linkable. This @@ -237,10 +238,9 @@ def compose_cache_tag(arch: Optional[str] = None) -> str: # ("0.1.21.post2") where a wheel's carried "+rocm..." and keyed ROCm by # accident; without it, two toolchains would share one tag. tag = f"{arch}__aiter-{version}__rocm-{_rocm_version()}" - # The version strings come from package metadata and torch, neither of which - # this module controls, so the composed tag gets the same check the arch did. - # Same character class the cache-tag test asserts on: the AITER and ROCm - # versions come from package metadata and torch, neither of which we control. + # These versions come from package metadata and torch, neither of which this + # module controls, so the tag gets the same check the arch did -- on the + # character class the cache-tag test asserts on. if tag != Path(tag).name or tag.startswith(".") or set(tag) & set("/:;, "): raise ValueError( f"refusing to build a cache directory name from {tag!r}: " diff --git a/tests/rocm/test_fp8_paged_prefill_aiter.py b/tests/rocm/test_fp8_paged_prefill_aiter.py deleted file mode 100644 index 11588cde43b..00000000000 --- a/tests/rocm/test_fp8_paged_prefill_aiter.py +++ /dev/null @@ -1,138 +0,0 @@ -# SPDX-FileCopyrightText: 2026 Advanced Micro Devices, Inc. -# -# SPDX-License-Identifier: Apache-2.0 - -"""fp8 paged prefill, the only route that serves fp8 at all. - -The in-tree fa2 kernel rejects 8-bit types at compile time and the flat-gather -route runs mha_varlen_fwd, which has no fp8 arm -- so AITER's native paged -kernel is the whole of fp8 prefill support, and nothing else here exercised it. - -It is also the only end-to-end check on the four fp8 variants -docker/prebuild_aiter_attention.py composes. The driver's ``--check`` now -rejects an artifact carrying no CK instances, but only this exercises one. -""" - -import pytest -import torch - -import flashinfer -from flashinfer.rocm.prefill import ( - FP8_PREFILL_OUT_DTYPE, - _aiter_paged_route_page_sizes, - _native_fp8_dtype, -) -from tests.test_helpers.test_helpers import requires_aiter - -pytestmark = requires_aiter - -NUM_QO_HEADS, NUM_KV_HEADS, HEAD_DIM = 8, 2, 128 -PAGE, KV_LEN, Q_LEN, BATCH = 16, 512, 64, 2 - - -def _fp8_dtype(): - """This GPU's fp8 encoding, from the same probe routing uses. - - The two encodings are indistinguishable by shape or by the .so name, and the - wrong one is read under the wrong exponent bias and returns NaN. - """ - pytest.importorskip("aiter") - dtype = _native_fp8_dtype() - if dtype is None: - pytest.skip("aiter.dtypes.fp8 is unreadable, so the encoding is unknown") - return dtype - - -@pytest.fixture(scope="module") -def workspace(): - return torch.empty(128 * 1024 * 1024, dtype=torch.uint8, device="cuda:0") - - -def _plan_and_run(workspace, q, kv, backend, **kw): - pages_per_seq = KV_LEN // PAGE - dev = q.device - w = flashinfer.BatchPrefillWithPagedKVCacheWrapper( - workspace, "NHD", backend=backend - ) - w.plan( - torch.arange(0, (BATCH + 1) * Q_LEN, Q_LEN, device=dev, dtype=torch.int32), - torch.arange( - 0, (BATCH + 1) * pages_per_seq, pages_per_seq, device=dev, dtype=torch.int32 - ), - torch.arange(BATCH * pages_per_seq, device=dev, dtype=torch.int32), - torch.full((BATCH,), PAGE, device=dev, dtype=torch.int32), - NUM_QO_HEADS, - NUM_KV_HEADS, - HEAD_DIM, - PAGE, - causal=True, - q_data_type=q.dtype, - kv_data_type=kv.dtype, - ) - return w.run(q, kv, **kw) - - -def test_fp8_paged_prefill_matches_a_bf16_reference(workspace): - fp8 = _fp8_dtype() - if PAGE not in _aiter_paged_route_page_sizes(fp8): - pytest.skip(f"page_size={PAGE} does not route to the native paged kernel") - - torch.manual_seed(0) - dev = "cuda:0" - q = torch.randn( - BATCH * Q_LEN, NUM_QO_HEADS, HEAD_DIM, device=dev, dtype=torch.bfloat16 - ) - kv = torch.randn( - BATCH * (KV_LEN // PAGE), - 2, - PAGE, - NUM_KV_HEADS, - HEAD_DIM, - device=dev, - dtype=torch.bfloat16, - ) - - ref = _plan_and_run(workspace, q, kv, "fa2") - - # From the dtype, not a constant: e4m3fnuz maxes at 240 and e4m3fn at 448, - # so a fixed divisor wastes half the range on one of the two architectures. - fp8_max = torch.finfo(fp8).max - sq = q.abs().amax().float() / fp8_max - sk = kv.abs().amax().float() / fp8_max - got = _plan_and_run( - workspace, - (q.float() / sq).to(fp8), - (kv.float() / sk).to(fp8), - "aiter", - scale_q=sq.reshape(1), - scale_k=sk.reshape(1), - scale_v=sk.reshape(1), - ) - - assert got.dtype == FP8_PREFILL_OUT_DTYPE - assert not torch.isnan(got).any(), "fp8 prefill returned NaN" - rel = (got.float() - ref.float()).abs().mean() / ref.float().abs().mean() - # Loose on purpose: this asserts the kernel ran and is not garbage, not that - # fp8 quantisation is accurate. - assert rel < 0.15, f"fp8 output differs from the bf16 reference by rel={rel:.4f}" - - -def test_fp8_prefill_is_refused_on_fa2(workspace): - """fa2 has no fp8 kernel, so the refusal must be an error, not a fallback.""" - fp8 = _fp8_dtype() - dev = "cuda:0" - q = torch.zeros( - BATCH * Q_LEN, NUM_QO_HEADS, HEAD_DIM, device=dev, dtype=torch.bfloat16 - ) - kv = torch.zeros( - BATCH * (KV_LEN // PAGE), - 2, - PAGE, - NUM_KV_HEADS, - HEAD_DIM, - device=dev, - dtype=torch.bfloat16, - ) - - with pytest.raises(NotImplementedError, match="fp8"): - _plan_and_run(workspace, q.to(fp8), kv.to(fp8), "fa2") From 3a423d1924fea9b7b1960599309fbff3f984b0cf Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 02:12:22 -0400 Subject: [PATCH 39/41] fix(rocm): gate HIP_CLANG_PATH on existence, and finish the wheel-to-source sweep _hip_clang_path claimed to mirror AITER's compile_ops and did not: core.py:1776 guards the export with `os.path.exists(hip_clang_path)`, so an inherited but stale MHA_HIP_CLANG_PATH is ignored by the lazy build and was being turned into an invalid compiler path by the prebuild. Same guard now. The standalone recipe ran `cd aiter && pip install .` unparenthesised, so the prebuild command in the next block resolved inside the AITER clone, where the FlashInfer script does not exist. Wrapped in a subshell and said why. Two more messages that still sent users to a wheel after this branch made the source build the supported install: prefill.py's _require_aiter_runtime, on the explicit-AITER path, and the tutorial's missing-package cell. mla.py's copy was already corrected; these are the rest of the set. Co-Authored-By: Claude --- docker/prebuild_aiter_attention.py | 4 +++- docs/rocm/backends.md | 7 +++++-- examples/amd_flashinfer_rocm_tutorial.ipynb | 4 ++-- flashinfer/rocm/prefill.py | 7 +++---- 4 files changed, 13 insertions(+), 9 deletions(-) diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index f19e6693098..ed45dd400a3 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -295,8 +295,10 @@ def _hip_clang_path(args: Dict) -> Iterator[None]: Only ``compile_ops`` sets it; calling ``build_module`` directly would leave the mha families on a different clang from their lazy builds. """ + # os.path.exists, as AITER's compile_ops does: a stale MHA_HIP_CLANG_PATH + # must be ignored here too, not turned into an invalid compiler path. path = args.get("hip_clang_path") - if not path: + if not path or not os.path.exists(path): yield return prev = os.environ.get("HIP_CLANG_PATH") diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index 44f623d4947..c264e0d5b81 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -156,10 +156,13 @@ pip install "setuptools>=80" "setuptools-scm>=9.2" "packaging>=24" git clone --recursive --depth 1 --shallow-submodules \ --branch v0.1.21.post2 https://github.com/ROCm/aiter.git -cd aiter && PREBUILD_KERNELS=0 GPU_ARCHS="gfx942;gfx950" \ - pip install --no-build-isolation . +( cd aiter && PREBUILD_KERNELS=0 GPU_ARCHS="gfx942;gfx950" \ + pip install --no-build-isolation . ) ``` +The subshell matters: the prebuild below is a FlashInfer script and has to run +from the FlashInfer checkout, not from the AITER clone. + Source rather than a wheel because no `amd-aiter` wheel is built against ROCm 10.0; the published ones are retargets of the same revision. A source build compiles against the ROCm actually present, and it is the only route to diff --git a/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 48487707178..e91b3bfdf13 100644 --- a/examples/amd_flashinfer_rocm_tutorial.ipynb +++ b/examples/amd_flashinfer_rocm_tutorial.ipynb @@ -208,8 +208,8 @@ "if importlib.util.find_spec(\"aiter.ops\") is None:\n", " raise RuntimeError(\n", " \"The AITER Python package is required for the prefill examples in this tutorial.\\n\"\n", - " \"Install it from the nightlies index; see docs/rocm/backends.md for the\\n\"\n", - " \"exact pinned version and why a loose specifier will not resolve it.\"\n", + " \"Build it from source at the pinned tag; docs/rocm/backends.md has the\\n\"\n", + " \"exact clone and pip command, and why a wheel is not the supported install.\"\n", " )\n", "print(\"AITER package available (aiter.ops).\")" ] diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index 02d397e14d1..af9af4b51a6 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -431,10 +431,9 @@ def _require_aiter_runtime(device: torch.device, op: str = "batch_prefill") -> N ) raise ImportError( "The 'aiter' package is required for the AITER backend and is not " - f"installed. Install a wheel >= {AITER_MIN_VERSION}; see " - "docs/rocm/backends.md for the index and the pinned version. A source " - "build tracks master, whose C ABI does not match the structs vendored " - "here." + f"installed. Build it from source at the pinned tag (>= {AITER_MIN_VERSION}); " + "docs/rocm/backends.md has the recipe. Build by tag, never a branch: " + "master's C ABI does not match the structs vendored here." ) From e2fdd10de70cef98520ffacb1efbf39fff156855 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 02:30:30 -0400 Subject: [PATCH 40/41] fix(rocm): cgroup-aware job budget, arch-safe header probe, and five smaller ones _aiter_csrc_include_dir's AITER_CSRC_DIR fallback was a third `import aiter` site that skipped _ensure_aiter_gpu_archs(). The other two call it first because AITER freezes its architecture at import, so reaching this one with GPU_ARCHS unset could pin the wrong arch for the rest of the process. _child_env sized its budget from os.cpu_count(), which reports the host's cores inside a cgroup-limited `docker build`; sched_getaffinity reports what the builder may actually use. AITER_SOURCE was compared literally to "1", so `--build-arg AITER_SOURCE=true` took the wheel branch and installed a ROCm 7.14 retarget into a ROCm 10 image without complaint. Now rejects anything that is not 0 or 1. The cache tag is checked against a positive character class rather than five denied spellings; a Popen that raises now unlinks its own log, which nothing else would have, since only started children reach `running`; the driver removes AITER's shared build/ root, not just the per-module subdirectory, so it stops shipping in the layer; and _loader_name_tokens asserts start < end, without which reordering the two functions in aiter_loader.cc would empty the slice and silently retire the parity check. Also drops test_fused_add_rmsnorm_aiter_is_correct_on_a_side_stream. Its own docstring recorded that it passes with the guard removed, so it banked coverage it did not have; the two cases that do detect it are A/B'd and stay. Co-Authored-By: Claude --- docker/Dockerfile.rocm | 3 +++ docker/prebuild_aiter_attention.py | 14 +++++++++-- flashinfer/jit/rocm/aiter_source.py | 9 ++++++- tests/rocm/test_aiter_stream_affinity.py | 26 +++------------------ tests/rocm/test_prebuild_aiter_attention.py | 1 + 5 files changed, 27 insertions(+), 26 deletions(-) diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7ae35056029..9acf0bfc90a 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -182,6 +182,9 @@ RUN pip install --no-cache-dir \ # prebuilt modules are no loss -- FlashInfer links its own lib.so built # from these same sources, and measured, only aiter.mla loads one of AITER's. RUN set -eu; \ + case "$AITER_SOURCE" in 0|1) ;; \ + *) echo "ERROR: AITER_SOURCE must be 0 or 1, got '$AITER_SOURCE'" >&2; exit 1 ;; \ + esac; \ if [ "$AITER_SOURCE" = "1" ]; then \ git clone --recursive --depth 1 --shallow-submodules \ --branch "$AITER_REF" "$AITER_REPO" /tmp/aiter && \ diff --git a/docker/prebuild_aiter_attention.py b/docker/prebuild_aiter_attention.py index ed45dd400a3..3a7c18edfbd 100644 --- a/docker/prebuild_aiter_attention.py +++ b/docker/prebuild_aiter_attention.py @@ -411,7 +411,9 @@ def _child_env(jobs: int) -> Dict[str, str]: try: budget = int(env["MAX_JOBS"]) except (KeyError, ValueError): - budget = int((os.cpu_count() or 1) * 0.8) + # sched_getaffinity, not cpu_count: the latter reports the host's cores + # inside a cgroup-limited `docker build`. + budget = int(len(os.sched_getaffinity(0)) * 0.8) env["MAX_JOBS"] = str(max(1, budget // max(1, jobs))) return env @@ -435,6 +437,9 @@ def _drain(pending, running, failed, jobs, total) -> None: stderr=subprocess.STDOUT, env=_child_env(jobs), ) + except BaseException: + os.unlink(log_path) # never reached `running`, so nothing else frees it + raise finally: os.close(fd) # Popen dup'd it; holding ours would exhaust the table running.append((v, proc, time.time(), log_path)) @@ -588,7 +593,12 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ) # Both, not `rc or`: a run that lost one variant should still report every # other artifact that linked hollow, rather than one finding per rebuild. - return _check(variants, extra) or rc + rc = _check(variants, extra) or rc + if rc == 0: + # AITER makes build/ at import and each build adds to it; only the + # per-module subdirectory was removed, so the root shipped in the layer. + shutil.rmtree(jit_dir(_aiter_jit_core()) / "build", ignore_errors=True) + return rc if __name__ == "__main__": diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index b2383e7d3b8..f0a715e1a95 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -53,6 +53,10 @@ # trailing newline, which would reach the ninja link line intact. _LIB_NAME_RE = re.compile(r"\A[A-Za-z0-9_][A-Za-z0-9_.+-]*\Z") +# The assembled cache tag becomes a directory name and reaches a -L/-rpath flag, +# so it is checked against a positive class, not a list of denied characters. +_TAG_RE = re.compile(r"\A[A-Za-z0-9._+-]+\Z") + # Guards the env-mutating build in _build_aiter_lib; see ensure_aiter_lib. _BUILD_LOCK = threading.Lock() @@ -241,7 +245,7 @@ def compose_cache_tag(arch: Optional[str] = None) -> str: # These versions come from package metadata and torch, neither of which this # module controls, so the tag gets the same check the arch did -- on the # character class the cache-tag test asserts on. - if tag != Path(tag).name or tag.startswith(".") or set(tag) & set("/:;, "): + if tag != Path(tag).name or tag.startswith(".") or not _TAG_RE.match(tag): raise ValueError( f"refusing to build a cache directory name from {tag!r}: " f"not a single safe path component" @@ -332,6 +336,9 @@ def _aiter_csrc_include_dir() -> Path: # only case where importing aiter (and needing a live GPU for it) is no # worse than the RuntimeError below. try: + from ...rocm.aiter_utils import _ensure_aiter_gpu_archs + + _ensure_aiter_gpu_archs() from aiter.jit.core import AITER_CSRC_DIR inc = Path(AITER_CSRC_DIR) / "include" diff --git a/tests/rocm/test_aiter_stream_affinity.py b/tests/rocm/test_aiter_stream_affinity.py index 2a196bb6e3d..8ba8a416865 100644 --- a/tests/rocm/test_aiter_stream_affinity.py +++ b/tests/rocm/test_aiter_stream_affinity.py @@ -13,6 +13,9 @@ vLLM and SGLang both drive a rank inside `with torch.cuda.stream(s)`, so this is the production configuration, not an exotic one. + +Only cases that actually fail with the guard removed live here: A/B'd on gfx950, +rmsnorm at hidden=4096 and the rope case do, the in-place fused_add path did not. """ import pytest @@ -58,29 +61,6 @@ def test_rmsnorm_aiter_is_correct_on_a_side_stream(hidden): torch.testing.assert_close(got, expected, rtol=1e-2, atol=1e-2) -def test_fused_add_rmsnorm_aiter_is_correct_on_a_side_stream(): - """Covers the in-place path, which stages into a fresh buffer and copies back. - - Measured: this case still passes with the guard removed. rmsnorm at 4096 and - the rope case are the two that actually detect a missing guard. - """ - hidden = 1024 - x = torch.randn(256, hidden, device="cuda", dtype=DT) - res = torch.randn(256, hidden, device="cuda", dtype=DT) - w = torch.randn(hidden, device="cuda", dtype=DT) - - x_ref, res_ref = x.clone(), res.clone() - flashinfer.fused_add_rmsnorm(x_ref, res_ref, w, backend="aiter") - - x_got, res_got = x.clone(), res.clone() - _on_side_stream( - lambda: flashinfer.fused_add_rmsnorm(x_got, res_got, w, backend="aiter") - ) - - torch.testing.assert_close(x_got, x_ref, rtol=1e-2, atol=1e-2) - torch.testing.assert_close(res_got, res_ref, rtol=1e-2, atol=1e-2) - - def test_rope_aiter_is_correct_on_a_side_stream(): nnz, heads, head_dim = 512, 8, 128 q = torch.randn(nnz, heads * head_dim, device="cuda", dtype=DT) diff --git a/tests/rocm/test_prebuild_aiter_attention.py b/tests/rocm/test_prebuild_aiter_attention.py index 5378051f053..9420a5a46ee 100644 --- a/tests/rocm/test_prebuild_aiter_attention.py +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -169,6 +169,7 @@ def _loader_name_tokens() -> set: start = text.find("const char* dtype_token(") end = text.find("\n}", text.find("std::string build_so_name(")) assert start != -1 and end != -1, "aiter_loader.cc: name composition not found" + assert start < end, "aiter_loader.cc: dtype_token now follows build_so_name" tokens = set(re.findall(r'"([a-z0-9_]+)"', text[start:end])) call_sites = re.findall(r"build_so_name\(key,([^;]*?)\)", text, re.S) From 644d725f1632b958ce9a80e7ecc7b3bc84398b16 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Tue, 15 Sep 2026 02:59:56 -0400 Subject: [PATCH 41/41] fix(rocm): include AITER's real norm and rope headers instead of declaring them The forward declarations existed because rmsnorm_quant.h and rope.h pulled in and so pybind11, which clashes with -DPy_LIMITED_API. At 0.1.21.post2 that is no longer true -- both are POD-only: rmsnorm_quant.h 1 include (aiter_tensor.h), 0 torch/pybind refs, 18 aiter_tensor_t rope.h 3 includes (POD + cstdint), 0 torch/pybind refs, 82 aiter_tensor_t So the reason for hand-writing the signatures is gone, and keeping them defeats the compile-time ABI check this module's own docstring prescribes -- the next signature change would again compile clean and fail at dlopen on the mangled name, which is the failure this branch exists to remove. Verified on a cold JIT cache so both shims really recompile: 73 passed across test_rmsnorm_aiter, test_rope_aiter and test_aiter_stream_affinity. Co-Authored-By: Claude --- csrc/rocm/norm_aiter.cu | 18 +++++------------- csrc/rocm/rope_aiter.cu | 16 +++++----------- flashinfer/jit/rocm/aiter_source.py | 9 ++++----- 3 files changed, 14 insertions(+), 29 deletions(-) diff --git a/csrc/rocm/norm_aiter.cu b/csrc/rocm/norm_aiter.cu index 2ee8bd49c16..bf72da3ad8a 100644 --- a/csrc/rocm/norm_aiter.cu +++ b/csrc/rocm/norm_aiter.cu @@ -17,20 +17,12 @@ #include #include -#include "aiter_tensor_compat.h" +// The real header, not a forward declaration: at 0.1.21 rmsnorm_quant.h is +// POD-only and no longer pulls , so a signature change is a +// compile error here rather than a dlopen failure on the mangled name. +#include -// AITER's public header (rmsnorm_quant.h) pulls in → full -// pybind11, which clashes with FlashInfer's -DPy_LIMITED_API, so forward-declare -// the entry points; the linker resolves them against the symbol-visible AITER -// .so. 0.1.21 moved both onto the POD aiter_tensor_t API -- declaring them with -// at::Tensor compiles and then fails at dlopen on the mangled name. -namespace aiter { -void add_rmsnorm(aiter_tensor_t& out, aiter_tensor_t& input, aiter_tensor_t& residual_in, - aiter_tensor_t& residual_out, aiter_tensor_t& weight, double epsilon, - bool gemma_norm); -void rmsnorm(aiter_tensor_t& out, aiter_tensor_t& input, aiter_tensor_t& weight, double epsilon, - bool gemma_norm); -} // namespace aiter +#include "aiter_tensor_compat.h" void fused_add_rmsnorm_aiter(at::Tensor input, at::Tensor residual, at::Tensor weight, double eps) { const c10::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(input.device()); diff --git a/csrc/rocm/rope_aiter.cu b/csrc/rocm/rope_aiter.cu index dd50bad342f..49bb1a62b4e 100644 --- a/csrc/rocm/rope_aiter.cu +++ b/csrc/rocm/rope_aiter.cu @@ -18,18 +18,12 @@ #include #include -#include "aiter_tensor_compat.h" +// The real header, not a forward declaration: at 0.1.21 rope.h is POD-only and +// no longer pulls , so a signature change is a compile error +// here rather than a dlopen failure on the mangled name. +#include -// AITER's public header (rope.h) pulls in → full pybind11, -// which clashes with FlashInfer's -DPy_LIMITED_API, so forward-declare the entry -// point instead; the linker resolves it against the symbol-visible AITER .so. -// 0.1.21 moved it onto the POD aiter_tensor_t API -- declaring it with at::Tensor -// compiles and then fails at dlopen on the mangled name. -void rope_cached_positions_2c_fwd_impl(aiter_tensor_t& output_x, aiter_tensor_t& output_y, - const aiter_tensor_t& input_x, const aiter_tensor_t& input_y, - const aiter_tensor_t& cos, const aiter_tensor_t& sin, - const aiter_tensor_t& positions, const int32_t rotate_style, - const bool reuse_freqs_front_part, const bool nope_first); +#include "aiter_tensor_compat.h" void apply_rope_pos_ids_cos_sin_cache_aiter(at::Tensor query, at::Tensor key, at::Tensor query_out, at::Tensor key_out, at::Tensor cos_sin_cache, diff --git a/flashinfer/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index f0a715e1a95..a24ff811592 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -6,11 +6,10 @@ FlashInfer wraps AITER kernels by compiling a small ``csrc/rocm/*_aiter.cu`` shim that calls AITER's C++ entry point directly and links the symbol-visible AITER ``.so``. Prefer ``#include``-ing AITER's real header, so a signature change is a -compile error rather than a load-time ``undefined symbol``. A forward -declaration is the fallback for headers that still pull in pybind11, which -clashes with ``-DPy_LIMITED_API`` -- but declare the signature AITER actually -has: since 0.1.21 rmsnorm and rope take ``aiter_tensor_t&``, and the old -``at::Tensor`` spelling compiles and then fails at ``dlopen`` on the mangled name. +compile error rather than a load-time ``undefined symbol``. Since 0.1.21 the +norm and rope headers are POD-only, so every shim now includes the real one; a +forward declaration remains the fallback only for a header that still pulls in +pybind11, which clashes with ``-DPy_LIMITED_API``. AITER's installed wheel builds its modules with ``-fvisibility=hidden``, so the kernel symbols (e.g. ``rope_cached_positions_2c_fwd_impl``) are not linkable. This