diff --git a/.claude/skills/benchmark-kernel/SKILL.md b/.claude/skills/benchmark-kernel/SKILL.md index 16de4710685..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.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): `_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. @@ -83,6 +84,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: diff --git a/CLAUDE.md b/CLAUDE.md index ad79cee4f25..37a666f27a3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -111,29 +111,23 @@ 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: 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 +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..af5389602b4 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -22,9 +22,11 @@ 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 (~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 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/csrc/rocm/aiter_loader.cc b/csrc/rocm/aiter_loader.cc index c54f9eedb2d..c6e10a084a1 100644 --- a/csrc/rocm/aiter_loader.cc +++ b/csrc/rocm/aiter_loader.cc @@ -21,13 +21,13 @@ 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 or re-pin the symbols in csrc/rocm/aiter_loader.cc."; + "\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 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. // @@ -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" @@ -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/csrc/rocm/norm_aiter.cu b/csrc/rocm/norm_aiter.cu index a4344ad9755..bf72da3ad8a 100644 --- a/csrc/rocm/norm_aiter.cu +++ b/csrc/rocm/norm_aiter.cu @@ -17,17 +17,12 @@ #include #include -// 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. -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, - bool gemma_norm); -} // namespace aiter +// 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 + +#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()); @@ -42,7 +37,16 @@ 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); + 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); residual.copy_(residual_out); } @@ -54,11 +58,16 @@ 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); + const flashinfer::aiter_compat::StreamGuard stream_guard(at::hip::getCurrentHIPStream()); 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..49bb1a62b4e 100644 --- a/csrc/rocm/rope_aiter.cu +++ b/csrc/rocm/rope_aiter.cu @@ -18,15 +18,12 @@ #include #include -// 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, - const bool reuse_freqs_front_part, const bool nope_first); +// 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 + +#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, @@ -77,7 +74,15 @@ 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); + 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), /*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 68010dda823..9acf0bfc90a 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 @@ -14,6 +15,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,11 +37,25 @@ 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. -ARG AITER_VERSION=0.1.20+rocm10.1.0a20260819.3135022 +# 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.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. 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 + +# 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. @@ -133,6 +153,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 +165,101 @@ 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; \ + 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 && \ + ( 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}") + +# 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) < 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")) +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 + +# 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 + +# 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; \ + export GPU_ARCHS="$AITER_GPU_ARCHS"; \ + JIT="$(python3 /opt/prebuild_aiter_attention.py --print-jit-dir)"; \ + case "$JIT" in \ + */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; \ + 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 new file mode 100644 index 00000000000..3a7c18edfbd --- /dev/null +++ b/docker/prebuild_aiter_attention.py @@ -0,0 +1,605 @@ +#!/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. + +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. + + GPU_ARCHS="gfx942;gfx950" python3 docker/prebuild_aiter_attention.py --jobs 2 +""" + +from __future__ import annotations + +import argparse +import contextlib +import os +import shutil +import subprocess +import sys +import tempfile +import time +from pathlib import Path +from typing import Dict, Iterator, List, NamedTuple, Optional, Sequence, Tuple + +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. +RECEIPT_MHA_FWD = 100 +RECEIPT_VARLEN = 200 +RECEIPT_BATCH_PREFILL = 200 + +GENERATE_PY = "example/ck_tile/01_fmha/generate.py" + +# 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") + + +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" + + @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.""" + 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" + + f"_{v.qscale_token}" + + "_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*" + + f"_{v.qscale_token}*" + + "_nsink*" + ) + + +def reachable_variants() -> Tuple[Variant, ...]: + """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 lse_values: + out.append( + Variant(family, dtype, has_logits_cap, needs_mask, has_lse) + ) + return tuple(out) + + +def selected_variants(only: Optional[Sequence[str]] = None) -> Tuple[Variant, ...]: + """Every reachable variant; ``only`` narrows to named families. + + Nothing is trimmed: the whole set costs minutes, a missing arm costs a stall. + """ + variants = reachable_variants() + 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 ``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"): + raise RuntimeError( + f"aiter's optCompilerConfig.json gave no srcs/includes for " + f"module_{family}; the config key has been renamed" + ) + 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. + """ + # 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 or not os.path.exists(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}") + 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}") + 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() + recipe = build_recipe(v, core) + args = _args_of_build(core, v.family) + 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}") + # 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 + + +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(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. + running: List[Tuple[Variant, subprocess.Popen, float, str]] = [] + failed: List[str] = [] + total = len(pending) + + try: + _drain(pending, running, failed, jobs, total) + finally: + 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 + + +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) + try: + budget = int(env["MAX_JOBS"]) + except (KeyError, ValueError): + # 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 + + +def _drain(pending, running, failed, jobs, total) -> None: + done = 0 + while pending or running: + while pending and len(running) < jobs: + v = pending.pop(0) + 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", + v if isinstance(v, str) else _encode(v), + ], + stdout=fd, + 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)) + 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 + label = f"{v}.so" if isinstance(v, str) else v.so_name + if proc.returncode == 0: + print( + f"[{done}/{total}] ok {label} {time.time() - started:.0f}s", + flush=True, + ) + os.unlink(log_path) + else: + failed.append(label) + tail = Path(log_path).read_text(errors="replace").splitlines()[-15:] + print( + f"[{done}/{total}] FAIL {label} {time.time() - started:.0f}s" + f" (full log: {log_path})\n" + "\n".join(tail), + flush=True, + ) + + +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))) + + +_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. + """ + 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: + """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 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: + 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 + + +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( + "--only", + action="append", + choices=sorted(_MD_NAME), + help="restrict to a family (repeatable)", + ) + p.add_argument( + "--jobs", + # 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( + "--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)) + else: + build_module_by_name(a.build_one) + 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 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) + + 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() + rc = _run_jobs(variants, a.jobs, extra=extra) + print( + f"\n{len(variants)} variant(s) + {len(extra)} module(s) in " + f"{(time.time() - t0) / 60:.0f} min at --jobs {a.jobs}" + ) + # 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. + 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__": + raise SystemExit(main()) diff --git a/docs/rocm/backends.md b/docs/rocm/backends.md index 22dc275047f..c264e0d5b81 100644 --- a/docs/rocm/backends.md +++ b/docs/rocm/backends.md @@ -147,49 +147,66 @@ 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/ -``` - -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"`. +# --no-build-isolation means pip installs no build requirements for you. +pip install "setuptools>=80" "setuptools-scm>=9.2" "packaging>=24" -**`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. +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 . ) +``` -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. +The subshell matters: the prebuild below is a FlashInfer script and has to run +from the FlashInfer checkout, not from the AITER clone. -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`). +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. -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: +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` 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 -git clone --recursive https://github.com/ROCm/aiter.git -cd aiter && python3 setup.py develop +GPU_ARCHS="gfx942;gfx950" python3 docker/prebuild_aiter_attention.py --jobs 2 +# or, to skip it: export FLASHINFER_AITER_ASM_PREFILL=0 ``` -Nothing stops you running one, but treat it as untested here. +`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. + +Check afterwards that the install did not pull a CPU-only torch over your ROCm +one: `python -c "import torch; assert torch.version.hip"`. + +**`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 @@ -272,9 +289,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/examples/amd_flashinfer_rocm_tutorial.ipynb b/examples/amd_flashinfer_rocm_tutorial.ipynb index 56a09a088fb..e91b3bfdf13 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, `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." ] }, { @@ -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/jit/rocm/aiter_source.py b/flashinfer/jit/rocm/aiter_source.py index 12ccbfa8883..a24ff811592 100644 --- a/flashinfer/jit/rocm/aiter_source.py +++ b/flashinfer/jit/rocm/aiter_source.py @@ -6,10 +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``. 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``. 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 @@ -52,6 +52,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() @@ -201,25 +205,27 @@ def resolve_aiter_build_arch() -> str: return _DEFAULT_BUILD_ARCH -def _aiter_cache_tag() -> str: - """A filesystem-safe tag keying the lib cache by target arch and AITER version. - - 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. - - 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 _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 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 " @@ -231,13 +237,29 @@ def _aiter_cache_tag() -> str: version = _md.version("amd-aiter") except Exception: version = "unknown" - return f"{arch}__aiter-{version}" + # 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()}" + # 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 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" + ) + 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 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 @@ -247,7 +269,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 @@ -289,17 +311,44 @@ 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 ...rocm.aiter_utils import _ensure_aiter_gpu_archs + + _ensure_aiter_gpu_archs() + 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/flashinfer/jit/rocm/aiter_variants.py b/flashinfer/jit/rocm/aiter_variants.py index ec77fcd6be2..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 resolve_aiter_build_arch +from .aiter_source import compose_cache_tag __all__ = [ "Family", @@ -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 @@ -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-/``. @@ -224,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/flashinfer/rocm/aiter_utils.py b/flashinfer/rocm/aiter_utils.py index 3b4f1c49637..9b900f6d380 100644 --- a/flashinfer/rocm/aiter_utils.py +++ b/flashinfer/rocm/aiter_utils.py @@ -50,11 +50,10 @@ 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. -AITER_MIN_VERSION = "0.1.20" +# 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" def _aiter_installed_version() -> Optional[str]: @@ -98,8 +97,11 @@ 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 + + # 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/arch_caps.py b/flashinfer/rocm/arch_caps.py index 23255ac3380..6ca0415d312 100644 --- a/flashinfer/rocm/arch_caps.py +++ b/flashinfer/rocm/arch_caps.py @@ -272,7 +272,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. -# Measured on amd-aiter 0.1.20 over qo_len x kv_len x cap, vs an fp32 reference. +# 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 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 diff --git a/flashinfer/rocm/prefill.py b/flashinfer/rocm/prefill.py index 80969830348..af9af4b51a6 100755 --- a/flashinfer/rocm/prefill.py +++ b/flashinfer/rocm/prefill.py @@ -64,10 +64,11 @@ # 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" +# 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 @@ -89,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} @@ -426,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." ) 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_source_probes.py b/tests/rocm/test_aiter_source_probes.py index ce803376829..89a37332e70 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) @@ -83,12 +129,71 @@ def _uncached(self): yield aiter_source._aiter_csrc_include_dir.cache_clear() - def test_a_missing_header_tree_names_the_package(self, monkeypatch): - import aiter_meta + @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): + pytest.importorskip("aiter_meta") + 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.""" + aiter_meta = pytest.importorskip("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() diff --git a/tests/rocm/test_aiter_stream_affinity.py b/tests/rocm/test_aiter_stream_affinity.py new file mode 100644 index 00000000000..8ba8a416865 --- /dev/null +++ b/tests/rocm/test_aiter_stream_affinity.py @@ -0,0 +1,80 @@ +# 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. + +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 +import torch + +import flashinfer +from tests.test_helpers.test_helpers import requires_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 + + +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_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) diff --git a/tests/rocm/test_aiter_version_gate.py b/tests/rocm/test_aiter_version_gate.py index 90b34a8223b..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.20 -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 @@ -42,9 +42,12 @@ 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. + # 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 new file mode 100644 index 00000000000..9420a5a46ee --- /dev/null +++ b/tests/rocm/test_prebuild_aiter_attention.py @@ -0,0 +1,360 @@ +# 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 + + +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() + + +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_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()) == 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()) == 44 + + +@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} <= { + v.so_name for v in driver.reachable_variants() + } + 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: + # 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 +# --------------------------------------------------------------------------- + +# 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: + """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" + 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) + 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(): + """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" + + +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 +# --------------------------------------------------------------------------- + + +@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 + + # 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 + 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_token_order(family: str) -> dict: + """Name token -> position of its first appearance in AITER's own source. + + 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) + ) + 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): + """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 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) :] + + +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 + + +# --------------------------------------------------------------------------- +# 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" + ) diff --git a/tests/rocm/test_prefill_decode_dispatch.py b/tests/rocm/test_prefill_decode_dispatch.py index 89eef4660bf..9eb046a22c1 100644 --- a/tests/rocm/test_prefill_decode_dispatch.py +++ b/tests/rocm/test_prefill_decode_dispatch.py @@ -318,8 +318,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)