Build AITER from source at a pinned tag and prebuild every attention variant

.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:

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 <site-packages>/aiter/jit/module_<x>.so | grep '<fn>('

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 —

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
 

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<Key, void*, Hash>&
 // 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 <variant.so> | grep mha_fwd
+// Pinned to amd-aiter 0.1.21.post2. Regenerate with: nm -D <variant.so> | 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;
   });
 }

csrc/rocm/norm_aiter.cu

@@ -17,17 +17,12 @@
 #include <c10/core/GradMode.h>
 #include <c10/hip/HIPGuard.h>
 
-// AITER's public header (rmsnorm_quant.h) pulls in <torch/extension.h> → 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 <torch/extension.h>, so a signature change is a
+// compile error here rather than a dlopen failure on the mangled name.
+#include <rmsnorm_quant.h>
+
+#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);
 }

csrc/rocm/rope_aiter.cu

@@ -18,15 +18,12 @@
 #include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>
 #include <c10/hip/HIPGuard.h>
 
-// AITER's public header (rope.h) pulls in <torch/extension.h> → 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 <torch/extension.h>, so a signature change is a compile error
+// here rather than a dlopen failure on the mangled name.
+#include <rope.h>
+
+#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);