Build AITER from source at a pinned tag and prebuild every attention variant - #375
Merged
Merged
Conversation
The benchmark skill told you how to time a kernel but not which candidates to
time, so "AMD ships an optimized library for this" has been read as a reason to
route to it. The record here says that is a coin flip: AITER prefill beats fa2
by 3.67x on gfx942 and 4.51x on gfx950, while the in-tree HIP kernel beats AITER
on rmsnorm/fused_add_rmsnorm by 1.6-1.8x and on append_paged_kv_cache at
3.62 TB/s against 2.86.
Adds a section saying to enumerate the ROCm libraries that already implement an
op -- AITER, CK/CK-Tile, hipBLASLt, rocBLAS, MIOpen, Triton -- benchmark each
against the in-tree kernel, and route auto to the winner per arch, recording it
as a Capability row with an evidence string. Also names the three things that
make the measurement wrong more often than the stopwatch does: a library can win
on time and still fail the contract (CK-Tile layernorm2d reads fp32 gamma/beta
as the input dtype), a library may dispatch internally so "AITER" is not one
number, and gfx942 and gfx950 have diverged.
Also records a measurement the paged-prefill bullet needed. It already warned
that the advertised native page sizes are a hint and that plan() degrades to
flat-gather. On amd-aiter 0.1.21.post2 all three advertised sizes fail the probe
with "no matching kernel found", on both gfx942 and gfx950:
advertised native page sizes: [128, 256, 1024]
page_size= 128 available=False
page_size= 256 available=False
page_size= 1024 available=False
So mha_batch_prefill is currently unreachable and the paged path should be
benchmarked as flat-gather.
Co-Authored-By: Claude <noreply@anthropic.com>
The `aiter_libs` cache tag is `<arch>__aiter-<version>`. ROCm was never a
component because the pinned wheel's version string carried it:
amd-aiter 0.1.20+rocm10.1.0a20260819.3135022
A source-built AITER has no local version segment — `pip show amd-aiter` on a
`v0.1.21.post2` clone reports a bare `0.1.21.post2` — so that accidental key
disappears and two ROCm toolchains on one box share a cache directory. These
are CK-tile objects linked into the shims by `-L`/`-rpath`, and a mismatched
one loads rather than failing, so the failure mode is a wrong kernel, not an
error.
`_rocm_version()` already existed in aiter_variants for the #363 store tag;
move it down to aiter_source, which aiter_variants already imports from, and
use it in both. Its docstring claimed the aiter_libs cache "never leaves the
box that built it" as the reason ROCm was omitted there — true, and beside the
point, since a box's toolchain can change under a persistent cache.
The composed tag now gets the path-component guard the arch already had:
`torch.version.hip` is not this module's value to trust, and the arch check
alone would let a separator through in the ROCm field.
Co-Authored-By: Claude <noreply@anthropic.com>
`_aiter_csrc_include_dir` resolved only through the `aiter_meta` package. An
editable AITER install has no `aiter_meta` at all: `setup.py develop` takes an
`is_develop_mode()` branch that sets `packages = ["aiter"]`, so every shim
needing a C++ header — rope, rmsnorm, activation, fused MoE, batch decode —
failed with "is the aiter source package installed?" against an install that
plainly was.
Fall back to AITER's own `AITER_CSRC_DIR`. It derives from `AITER_META_DIR`,
which falls back to the repo root when `aiter_meta` has no `csrc/`, so in a
develop layout it names the clone's `csrc/` and in any layout it honours an
`AITER_META_DIR` override. Confirmed on the 0.1.21.post2 non-editable install
that both routes resolve to the same directory:
AITER_CSRC_DIR = .../site-packages/aiter_meta/csrc
same dir as aiter_meta/csrc/include = True
aiter_meta stays first: it needs no `import aiter`, and importing aiter runs
arch detection against a live device. The fallback pays that cost only where
the alternative was an exception.
Co-Authored-By: Claude <noreply@anthropic.com>
Both the module docstring and backends.md claimed prebuilding "cannot be a
`docker build` step". That conflates the driver with the machinery it calls.
`prebuild_aiter_variants` needs a GPU because it produces a variant by calling
the op, which launches a kernel. AITER's builder does not: `aiter/setup.py:212`
reaches it with
sys.path.insert(0, ".../aiter/"); from jit import core
a flat import that never runs `aiter/__init__.py` — which is what needs a
device, for arch detection and a triton import that fails with "0 active
drivers ([]). There should only be one." Measured: a GPU-less container builds
`mha_fwd_bf16_nbias_mask_nlse_ndropout_nqscale` this way and `nm -D` finds the
expected `aiter::mha_fwd` symbol.
The claim matters because it rules out an image-build prebuild that is in fact
available, so correct it before anything is designed around it.
Co-Authored-By: Claude <noreply@anthropic.com>
A GPU-free driver that builds the attention variant .so at image-build time,
plus the two parity tests that keep it honest. No Dockerfile change yet, so
nothing in this commit runs in CI; commit 5 wires it up.
Why a second driver rather than reusing prebuild_aiter_variants: that one emits
a variant by *calling the op*, which needs a device. AITER's builder does not —
aiter/setup.py:212 reaches it by a flat `sys.path` import of `jit.core` that
never runs aiter/__init__.py, the part needing arch detection and triton.
Artifacts land in AITER's own jit/ dir rather than FlashInfer's store, which is
what makes mha_batch_prefill worth building at all. Its bootstrap doubles as the
page-size probe and runs before the loader, so only a hit AITER itself can see
spares the compile. Measured on gfx942:
arm=warm plan() = 0.2s .so built during plan: 0
arm=cold plan() = 101.4s .so built during plan: 1
Default set is 36 of 40: everything but mha_fwd's four LSE variants, the most
expensive builds in the set at 803 s each two-arch. LSE on *single* prefill
serves cascade/merge, and Blaze-O1's LUM3 — the consumer that pins this repo as
a submodule — records its footprint as "fp16, head_dim 64, batch mode, causal,
no soft-cap/bias/LSE". Batch mode is mha_fwd, so LUM3's exact file
(mha_fwd_fp16_nbias_mask_nlse_ndropout_nqscale.so) is in the default set and a
test asserts it. `--all` builds 40.
Verified end to end, GPU_ARCHS="gfx942;gfx950", against the lazily built
artifact as control — varlen is the control because its recipe comes from
AITER's own helper:
lazy : 3123528 bytes 168 symbols
driver : 3123528 bytes 168 symbols
99 non-cuid symbols identical; the 69 that differ are all __hip_cuid_*
both carry amdgcn-amd-amdhsa--gfx942 and --gfx950
Two traps this cost, both now guarded by tests:
- get_args_of_build("mha_fwd") *warns and returns empty args* instead of
raising; the real key is "module_mha_fwd". The empty args dropped every
include path and surfaced minutes later as "fmha_fwd.hpp file not found".
_args_of_build now asserts srcs and includes are non-empty.
- compose_mha_fwd_variant_suffix_and_filter is misnamed: it always emits
_logits and _skip tokens, which is the *varlen* shape. Its only callers are
varlen. mha_fwd composes inline with receipt 100 against varlen's 200, and
spells its mask filter "_m*" where batch_prefill spells it "_mask*".
Child build output goes to a file, never a pipe: hipcc is verbose enough to fill
a pipe buffer, and a child blocked on a full pipe never exits.
All six parity assertions A/B'd by breaking the driver one way at a time:
caught mha_fwd mask filter -> _mask*
caught batch_prefill dtype filter gains a _
caught mha_fwd receipt 100 -> 200
caught batch_prefill loses its sink token
caught varlen token renamed
caught mha_fwd lse arm no longer skipped
Co-Authored-By: Claude <noreply@anthropic.com>
…ants
Wires up the driver from the previous commit. Deliberately still at v0.1.20 --
the ABI the image already runs -- so that a failure here is the *mechanism* and
a failure after the v0.1.21.post2 bump is the *revision*. Do not collapse the
two commits; the bisect depends on the split.
Source rather than wheel because it is the only way to compile AITER against
the ROCm in this image. No +rocm10.0 build of amd-aiter is published, so the
pinned wheel is a +rocm10.1.0a retarget of the same revision. AITER_SOURCE=0
restores that wheel install verbatim for an unblock.
Three things a reader cannot see by reading, so they are checked rather than
asserted in prose (tmp validator, not shipped):
- The clone is created and deleted in ONE RUN. Splitting them leaves it in
the earlier layer, where a later `rm -rf` cannot reach it -- CK alone is
over a GB.
- The driver COPY sits after the install, so editing the variant list re-runs
only the prebuild.
- Every AITER_* the new layers reference is declared with ARG in this stage.
Non-editable install on purpose: `setup.py develop` takes an is_develop_mode()
branch setting packages = ["aiter"] with no aiter_meta, and aiter_meta/csrc/include
is what every C++ shim compiles against. The non-editable path runs
prepare_packaging(), which copies csrc and the CK tree into aiter_meta -- which
is also what makes deleting the clone safe.
PREBUILD_KERNELS=0 because 1 rmtree's the jit dir and does
`from aiter.aot.flydsl.common import run_aot`, a real `import aiter`. Losing its
prebuilt modules costs nothing measurable: FlashInfer links its own
lib<module>.so built from the same sources, and of AITER's 111 prebuilt modules
only module_mla_asm is ever loaded (11.5 s, 0.2 MB).
The post-install check is a hard assertion, not a print. A shallow clone that
lost its tag makes setuptools_scm invent a .dev version, which falls under
aiter_utils.AITER_MIN_VERSION and silently degrades every AITER path to fa2 --
a working image that is quietly several times slower on prefill.
GPU_ARCHS must be explicit: there is no device here for AITER's arch detection,
and both architectures go in one fat binary.
Validated: 12 RUN instructions parse under `sh -n`, the embedded Python heredoc
parses, and all five invariants A/B'd by breaking each one and confirming the
check fires. This buildx has neither --check nor --call, and a real image build
is ~2h, so the build-time number lands in the PR description rather than here.
Co-Authored-By: Claude <noreply@anthropic.com>
The plan split this into three commits -- shim migration, then AITER_REF, then
the floor raise -- on the claim that the POD overloads already existed at
0.1.20 and the migration could land first. They do not:
0.1.20 rmsnorm_quant.h:37 void rmsnorm(torch::Tensor& out, ...)
0.1.20 grep aiter_tensor_t {rmsnorm,rmsnorm_quant,rope}.h -> no matches
0.1.21 aiter::rmsnorm(aiter_tensor_t&, aiter_tensor_t&, aiter_tensor_t&,
double, bool)
Either order leaves a commit whose rmsnorm/rope backend="aiter" tests fail at
dlopen with `undefined symbol: _ZN5aiter7rmsnormERN2at6TensorES2_S2_db`. They
are one atomic change. The floor raise joins them because a 0.1.20 user would
otherwise clear the gate and then fail to load.
The 5/6 bisect split the plan wanted is intact: commit 5 proved the *mechanism*
at 0.1.20, this moves the *revision*.
- norm_aiter.cu and rope_aiter.cu now declare the POD signatures and convert
through csrc/rocm/aiter_tensor_compat.h, matching activation/page/fused_moe,
which migrated at 0.1.16 and already pass. Named lvalues at the call sites:
the out parameters are non-const references, so a to_aiter() temporary would
not bind.
- AITER_MIN_VERSION 0.1.20 -> 0.1.21, _AITER_LAST_VALIDATED -> 0.1.21.post2
(no +rocm segment; a source build has none), the vendored header pins, and
kAbiPinNote, which now names the source build rather than a dead pip line.
Not touched: arch_caps.py evidence strings stay at 0.1.20 because they assert a
measurement that has not been repeated -- honestly stale beats quietly wrong.
The soft-cap re-measurement is the next commit.
Two plan items turned out to be non-issues, checked rather than assumed:
_AITER_SOFTCAP_DEFECT_THROUGH appears only in message text, and
test_single_prefill_kernels.py:92's skip is already driven by
arch_caps.aiter_softcap_defect_arch, so neither stops firing at 0.1.21.post2.
A/B by construction: all 69 rmsnorm/rope AITER tests failed at dlopen on
0.1.21.post2 before this change and pass after, measured on gfx942. The version
floor table gains 0.1.20 -> False and 0.1.21.dev0 -> False.
Co-Authored-By: Claude <noreply@anthropic.com>
The set was trimmed on an estimate that turned out 2-5x too slow. Measured in a
real image build at --jobs 2 on 32 cores:
family n measured s census est ratio
mha_fwd 4 162 803 5.0x
mha_varlen_fwd 16 65 140 2.2x
mha_batch_prefill 16 68 206 3.0x
36 variant(s) in 23 min at --jobs 2
36/36 present in .../aiter/jit
The census figures came from Step 0's S0.7, which ran two concurrent builds at
MAX_JOBS=16 each on purpose. The driver sets no such cap, so ninja uses the whole
box (-j 307) and every family is far cheaper than the plan assumed.
That dissolves the trade-off the trim was making. Dropping mha_fwd's four LSE
variants saves ~5 min of CI, once, and costs a multi-minute in-process stall to
whoever first calls single_prefill_with_kv_cache(..., return_lse=True). The whole
40 lands at roughly 26 min. `--only` still narrows by family for local work.
Also removes --all, which now has nothing to select.
Co-Authored-By: Claude <noreply@anthropic.com>
The gate is unchanged -- gfx950 defective, gfx942 clean -- but it was recorded
against 0.1.20 and the version move made that evidence stale. Re-measured on
both boards, guard disabled so the kernel actually runs, non-square shapes at
head_dim=128 causal, fp16, vs an fp32 reference.
gfx950 (MI350X, slurm 67930310), amd-aiter 0.1.21.post2, ROCm 10.0:
qo kv cap |aiter-ref| |fa2-ref|
1 1024 30 1.1930 0.0001
1 1024 50 2.6929 0.0001
17 1024 30 NaN 0.0001
37 4096 50 2.6767 0.0000
128 512 50 3.3845 0.0003
256 1024 50 3.4392 0.0002
1024 1024 50 4.1128 0.0010
12 cap>0 rows, 12 wrong (9 by magnitude, 3 returning NaN)
gfx942 (MI300X), same version and sweep: 12 cap>0 rows, 0 wrong, worst 0.0003.
cap=0 is clean on both boards, which is why the guard must stay narrow rather
than becoming a KnownBad row over the whole op.
Three things this sweep had to get right, each of which produced a wrong answer
first:
- The guard has to be disabled. With it on, every gfx950 capped row raises
FlashInfer's own ValueError and the sweep reports "not reproduced" -- a
clean result meaning "never ran". The script now exits 2 as INCONCLUSIVE
when no capped row executes, so that cannot read as a pass again.
- NaN output is a wrong answer, not a skipped row. Scoring it as "raised"
understated the defect by three rows; nothing threw in either run.
- qo_len > kv_len is dropped: causal with more queries than keys is ill-posed
and *both* backends fail, which says nothing about AITER.
_AITER_SOFTCAP_DEFECT_THROUGH moves to 0.1.21.post2. It is message text only;
the gate is arch_caps.aiter_softcap_defect_arch, which this leaves alone.
Co-Authored-By: Claude <noreply@anthropic.com>
The docs told users to pip-install a wheel and warned that a source build was
untested. The image now does the opposite, so every one of those passages is
not merely stale but actively wrong.
Rewritten rather than renumbered, because the reasons changed too:
- The wheel recipe becomes the tagged source build, with the three details
that actually bite: build by tag (a lost tag makes setuptools_scm invent a
.dev version that falls under the floor and silently degrades every AITER
path to fa2), non-editable (develop ships no aiter_meta), and
PREBUILD_KERNELS=0 (=1 needs a GPU, and its modules are no loss).
- The cp312 rationale is gone. It held because every 0.1.20 wheel was cp312
only; a source build follows whichever interpreter is present, so the
interpreter is now a property of the base image.
- The "+rocm10.1.0a is the nearest retarget" apology is replaced by the
headline correctness win: a source build compiles against the ROCm in the
image rather than a retarget of another one.
- CLAUDE.md's "a source build tracks master, treat as untested here" is
reversed. Its neighbouring "That gap is a trap" sentence lost its
antecedent with that paragraph, so it is reworded to stand alone -- the
advice to read the installed tree rather than a checkout is still right.
- CONTRIBUTING.md documents AITER_REF, AITER_PREBUILD_JOBS and AITER_SOURCE=0
in place of the wheel args.
Left alone deliberately: backends.md:435's "measured exact on amd-aiter 0.1.20"
is a record of when a measurement happened, not a pin.
The notebook edit is a two-line text change; cell and output counts are
unchanged (17 and 17) after the JSON round-trip.
Co-Authored-By: Claude <noreply@anthropic.com>
Self-review (/code-review xhigh) before push. One finding was a regression this
branch introduced; the rest are its consequences.
**The stream guard.** csrc/rocm/aiter_tensor_compat.h says it outright: "The POD
entry points launch on aiter::getCurrentHIPStream(), which defaults to nullptr
and is otherwise set only by AITER's Python layer." All three shims that
migrated at 0.1.16 install a StreamGuard; the two this branch migrated did not:
activation_aiter.cu 1
fused_moe_aiter.cu 1
page_aiter.cu 1
norm_aiter.cu 0 <- added here
rope_aiter.cu 0 <- added here
Under `with torch.cuda.stream(s)` -- which is how vLLM and SGLang drive a rank --
the kernel went to the null stream while `at::empty_like` and the `copy_` back
ran on `s`: a silent race, and an abort under graph capture. The at::Tensor
entry points picked the torch stream up themselves, so nothing caught it.
**AITER_SOURCE=0 installed a wheel below the floor this branch raised.** The
build stayed green and every AITER path silently stopped routing -- the exact
failure the surrounding assertions exist to prevent. The assertion block now
rejects any resolved version under AITER_MIN_VERSION, wheel arm included.
**The csrc-include fallback was unreachable.** `_aiter_importable()` still did a
hard `import aiter_meta`, so on the editable layout the fallback exists to serve,
routing gave up long before the resolver ran. It now probes through
`_aiter_csrc_include_dir()` instead, which is the thing that actually has to work.
Smaller, same review: `AITER_MIN_VERSION`'s comment still argued for 0.1.20;
kAbiPinNote had lost its actionable command and now names backends.md; the
benchmark-kernel skill still stated the old floor on the line that was edited
for this pin; the tutorial still promised native paged page sizes that fail the
probe at 0.1.21.post2; `_run_jobs` leaked children on the way out, which matters
because a zombie-held AITER build lock is never treated as stale; two driver
docstrings carried measurements that belong in a commit message.
Dropped `test_the_blaze_lum3_variant_is_in_the_default_set`: the proper noun it
cites appears nowhere in-tree, and `test_nothing_is_trimmed_from_the_built_set`
already covers the assertion.
128 tests pass after the shim rebuild the StreamGuard forces.
Co-Authored-By: Claude <noreply@anthropic.com>
The StreamGuard regression fixed in the previous commit passed all 69 norm and
rope AITER tests while it was live. Every one of them runs on the default
stream, where the null stream AITER falls back to and the caller's stream are
the same stream, so the race cannot appear. A reviewer found it by reading.
These tests run the same ops inside `with torch.cuda.stream(s)` with real work
queued ahead, which is how vLLM and SGLang drive a rank.
A/B, guards stripped from both shims and the shim cache cleared:
..F.F 2 failed, 3 passed
FAILED test_rmsnorm_aiter_is_correct_on_a_side_stream[4096]
FAILED test_rope_aiter_is_correct_on_a_side_stream
Restored, 5 passed. Only the larger shapes lose the race -- hence hidden=4096
in the parametrize and the matmul ballast, without which nothing fails. A race
test catches a subset by nature; two of five is enough to stop this returning
silently, which is the whole problem with it.
Co-Authored-By: Claude <noreply@anthropic.com>
Three defects found while verifying the image, each by running it rather than
reading it.
**603 MB of dead weight shipped.** The 40 variants are 224 MB; AITER's
jit/build/ staging -- CK blob sources and objects, ~15 MB per variant -- was
another 603 MB, three quarters of the directory. Proven removable rather than
assumed: with build/ deleted, a paged prefill plan() built 0 artifacts, because
AITER decides "already built" from the .so. The driver now removes each tree on
success and keeps it on failure, where it is the only diagnostic.
**The floor check compared strings.** `Version(got).base_version <
Version(FLOOR).base_version` is lexicographic, so it would have admitted 0.1.9
against a 0.1.10 floor. Comparing Version objects is both correct and simpler;
PEP 440 already puts 0.1.20+rocm10.1.0a below 0.1.21 and 0.1.21.dev0 below
0.1.21, which is what the floor wants.
**The wheel fallback could not build at all.** With the floor enforced,
AITER_SOURCE=0 correctly refused its 0.1.20 pin -- leaving a documented escape
hatch that always fails. The index does publish 0.1.21.post2, but only
+rocm7.14.0 and +rocm7.2.3: no ROCm 10 build exists at this revision, which is
the clearest statement yet of why the default path builds from source. The arm
now pins the 7.14 retarget and CONTRIBUTING says plainly that it is an unblock.
Also regenerates tmp/Dockerfile.c5probe from docker/Dockerfile.rocm instead of
hand-copying it. The hand-copied version silently lacked the floor assertion, so
the first AITER_SOURCE=0 run "passed" a check the real file would have failed --
a test that proved nothing, which is worse than no test.
Cold-start A/B on the built image, fresh container, same FlashInfer JIT cost in
both arms so the delta is AITER's:
family prebuilt cold saved
mha_fwd (single prefill) 52.5s 485.9s 433s
mha_varlen_fwd (soft cap) 41.6s 156.7s 115s
mha_varlen_fwd (ragged) 39.7s 177.5s 138s
mha_batch_prefill (paged) 0.4s 129.3s 129s
total 134.2s 949.4s 815s (7.1x)
Co-Authored-By: Claude <noreply@anthropic.com>
Rebasing onto amd-integration picked up #365, which routes single prefill to AITER's asm arm on CDNA4 and dlopens module_fmha_v3_fwd.so. Its own failure hint says why that is safe: module_fmha_v3_fwd.so ships prebuilt in the amd-aiter wheel, so unlike the mha_fwd variants there is no JIT build to trigger. Its absence means the wheel is incomplete or the JIT dir points elsewhere. This branch replaces the wheel with a PREBUILD_KERNELS=0 source install, which ships no prebuilt modules at all. Measured on the image built from this branch: module_fmha_v3_fwd.so present: NO module_* present: 0 So the two changes are individually correct and jointly broken: #365's default CDNA4 prefill route would throw at dlopen in the new image, with a message blaming an incomplete wheel that is not being used. The driver now builds whole modules as well as variants. module_fmha_v3_fwd takes 96 s GPU-free at GPU_ARCHS="gfx942;gfx950" and needs no filter, because AITER resolves the asm kernel from its own config table per call rather than from filename axes. The durable part is the parity test: it greps aiter_loader.cc for every "module_*.so" it opens and fails when the driver's set does not cover them. A/B'd by emptying LOADER_MODULES -- caught. Without it the next module added to the loader repeats this silently, since nothing else connects the two files. Co-Authored-By: Claude <noreply@anthropic.com>
The previous commit taught the driver to build module_fmha_v3_fwd but its own
--check still enumerated only the 40 variants, so the image-build gate would
pass with the asm module missing -- the exact hole that commit set out to close.
Visible in that build's own output: the job runner counted [41/41] while the
check printed 40/40.
The edit that was meant to do this silently matched nothing, because ruff-format
had rewrapped the function between writing and applying it, and unlike the other
edits in that batch this one asserted no anchor.
Verified against the image built from the previous commit:
520896 bytes .../aiter/jit/module_fmha_v3_fwd.so
mha_*.so: 40
amdgcn-amd-amdhsa--gfx942 amdgcn-amd-amdhsa--gfx950
The artifact is there and carries both code objects; only the count was wrong.
Co-Authored-By: Claude <noreply@anthropic.com>
The full suite on a box whose AITER is 0.1.20 -- below the floor this branch
raises -- failed six tests that pass on the image it ships. Both causes are the
tests trusting whatever AITER the host happens to have.
**The new stream tests guarded on `is_aiter_supported`**, which answers for the
architecture alone. On a sub-floor box they ran anyway and died inside
`require_aiter` rather than skipping. They now use `requires_aiter`, the helper
every other AITER test in tests/rocm uses, which is why those skipped cleanly on
the same run.
**`test_a_missing_aiter_package_is_its_own_reason` stubbed importability but not
presence.** `_auto_select_prefill_backend` reports the ABI-floor reason when a
sub-floor AITER *is* installed and the not-installed reason otherwise, so with
0.1.20 present the test asserted on the wrong branch:
assert 'aiter package not installed' in 'amd-aiter 0.1.20+rocm10.1.0a...
is below the 0.1.21 ABI floor (the vendored struct layouts do not match
older releases)'
It now pins absence as well, which is what its name claims. `TestRequireAiterRuntime`
directly below already does this, for the stated reason that which branch raises
"does not depend on the amd-aiter the host happens to have installed".
Verified both ways:
aiter 0.1.21.post2 43 passed
aiter 0.1.20 sssss...........s.......ss................. exit 0
Co-Authored-By: Claude <noreply@anthropic.com>
The driver enumerated bf16 and fp16 only, so an fp8 deployment got nothing
from this branch: every fp8 paged prefill still compiled AITER inside the
serving process, which is the stall the prebuild exists to remove.
fp8 is reachable and batch-prefill only. prefill.py routes every fp8 query
through the native paged kernel because the flat-gather route runs
mha_varlen_fwd, which has no fp8 arm, and the in-tree fa2 kernel rejects
8-bit types at compile time. All three AITER mha families carry an fp8 arm
upstream, but kFp8Bf16 is constructed at exactly one site in this tree
(batch_prefill_paged_aiter.cu:110), so the reachable set is 8 names, not 24:
mha_batch_prefill_fp8bf16_{logits,nlogits}_nbias_{mask,nmask}_{lse,nlse}
_ndropout_pertensor_nsink.so. All 8 build clean, ~66 s each.
The old parity test could not have caught this: it compared the driver's
table against flashinfer's, and both omit fp8. The replacement reads the
token literals out of aiter_loader.cc instead. A/B'd by reverting the dtype
set -- it fails with "aiter_loader.cc can emit ['fp8bf16', 'pertensor'],
never built".
flashinfer's own table stays at 40. It serves the store route, which has no
fp8 arm; its docstrings claimed to enumerate everything the loader can ask
for, which was already false and is now corrected.
Co-Authored-By: Claude <noreply@anthropic.com>
… readable
Three defects found by review, each verified by reproducing it.
--only built the family it was asked for and then failed its own check on
the whole-module set it deliberately skips:
$ prebuild_aiter_attention.py --only mha_fwd
8/9 present in .../aiter/jit
missing or empty:
module_fmha_v3_fwd.so # exit 1 on a healthy tree
_check now takes the same `extra` set the invocation builds.
build_module was called directly, which skips the HIP_CLANG_PATH export that
AITER's own compile_ops wraps the build in. optCompilerConfig.json sets it
from MHA_HIP_CLANG_PATH for module_mha_fwd and module_mha_varlen_fwd, so with
that variable set the prebuilt .so would come from a different clang than the
lazy build -- falsifying the byte-identical claim. Unset in this image, so no
artifact shipped so far is affected.
The chmod is the one with teeth. AITER's get_user_jit_dir() copies the whole
jit tree into $HOME when its own directory is unwritable, and the tree is
owned by uid 1003. Any container run with --user of another uid therefore
sees none of the prebuilt artifacts -- and dies outright where HOME is not
writable, which is not a fallback but a crash:
PermissionError: [Errno 13] Permission denied: '/nonexistent'
in get_user_jit_dir() -> shutil.copytree(this_dir, home_jit_dir)
Verified fixed: the same foreign uid now resolves to the installed jit dir.
--user is the documented run mode, so this would have hit normal use.
Co-Authored-By: Claude <noreply@anthropic.com>
… comment cap Moving _rocm_version into aiter_source let _aiter_cache_tag gain the ROCm component that variant_store_dir already had, which left the tag shape and its path-safety guard written out twice in two modules. The next change to either could land in one and not the other, and the two cache directories would then disagree about staleness. Both now call compose_cache_tag(). The rest is the repo's own comment cap, which this branch had been ignoring: the driver's module docstring ran 27 lines and carried a measurement, and two more docstrings ran long. Measurements belong in these commit messages. Also drops a stale claim in test_aiter_version_gate.py, which still said the vendored structs follow the 0.1.20 layout after the floor moved to 0.1.21. No behaviour change: the driver diff is docstrings only, and the tag helper composes the same string from the same inputs. Co-Authored-By: Claude <noreply@anthropic.com>
…ng guard
A/B on gfx950, guards stripped from norm_aiter.cu and rope_aiter.cu and the
shim rebuilt from a cold cache:
guards in source: 0 -> 2 failed, 3 passed
FAILED test_rmsnorm_aiter_is_correct_on_a_side_stream[4096]
NaN, 1048576 / 1048576 mismatched
FAILED test_rope_aiter_is_correct_on_a_side_stream
518956 / 524288 mismatched, 13.5 absolute
guards in source: 3 -> 5 passed
So the suite does catch the regression, but not through the case whose
docstring claimed it: fused_add_rmsnorm passes either way, because its
buffers are cloned on the default stream before the call and so never meet
the stream-ordered allocator hazard the other two hit. Left as coverage of
the in-place path with the claim corrected rather than strengthened -- a race
that reproduces sometimes is worse than two that reproduce every run.
Co-Authored-By: Claude <noreply@anthropic.com>
Nothing in tests/rocm exercised fp8 attention, which left the eight fp8
variants the previous commit added with no end-to-end check at all. That
matters more than the usual coverage argument: the driver's `_check` only
proves the .so exists and is non-empty, so a CK --filter that matched zero
instances would still link, still pass the build, and fail only at dispatch.
Measured on gfx950 with the driver-built artifacts in place -- AITER loads
the file our composed name asks for and the kernel runs:
[aiter] import [mha_batch_prefill_fp8bf16_nlogits_nbias_mask_nlse
_ndropout_pertensor_nsink] under .../aiter/jit/
fp8 aiter ok, out=(128, 8, 128) torch.bfloat16
mean abs err 0.0031 rel 0.0505 max 0.0293
The tolerance is deliberately loose. This asserts the kernel dispatched and
did not return garbage or NaN, not that fp8 quantisation is accurate; a tight
bound here would fail on the quantisation error rather than on a defect.
The second case pins the refusal: fa2 has no fp8 kernel, and without an
explicit raise the caller gets a ninja compiler log from a static_assert.
Co-Authored-By: Claude <noreply@anthropic.com>
…ong job first Three from review, all in the prebuild path. _run_jobs killed its in-flight children on the way out but never reaped them. AITER treats a build baton as stale only when kill(pid, 0) fails, which a zombie child of a living parent answers successfully -- so an interrupted run left the next build for that md_name blocking forever, indistinguishable from a slow compile. It also leaked the mkstemp log of every killed job. The chmod recomputed AITER's jit directory from the package path while the driver wrote to core.get_user_jit_dir(), which honours AITER_JIT_DIR and falls back to $HOME. The two agree in this image and could silently stop agreeing, chmod-ing an empty directory while reporting success -- reintroducing exactly the foreign-uid failure the chmod exists to prevent. The driver now prints the directory it resolved (--print-jit-dir) and the Dockerfile chmods that. a+rwX rather than a+w: the intent is that any uid can write into the tree, and X grants directory traversal without marking regular files executable. module_fmha_v3_fwd is the longest single job -- it compiles the whole asm FAV3 source set rather than one filtered CK instance group -- and was queued last, so at --jobs 2 it ran alone on the tail with a core idle. Queued first it overlaps the variants. --list and the closing summary now count the whole-module builds too; before, progress said [49/49] while the summary said 48 and --list named neither. Co-Authored-By: Claude <noreply@anthropic.com>
The benchmark skill's page-size bullet stated {128, 256, 1024} and called
{1, 16, 1024} "wrong". It has it backwards: _aiter_native_page_sizes() returns
frozenset({1, 16, 1024}) for any AITER >= 0.1.10. The bullet also called the
{16, 1024} arm unreachable at the 0.1.21 floor, when it is the
PackageNotFoundError fallback and reachable whenever metadata is missing.
Its measured claim was stale too -- "all three advertised sizes fail the probe
... mha_batch_prefill is unreachable". On 0.1.21.post2 the native paged kernel
runs: 62 TFLOPS at page_size=1024, kv=1024 on gfx950. What is true, and now
recorded, is that the sweep has to stop there; beyond kv=1024 it faults inside
FmhaBatchPrefillWithPagedKVCacheKernel on 0.1.20 and 0.1.21.post2 alike.
Also adds importorskip to the two aiter_meta probes. They were the only
aiter-dependent tests in the tree without a guard, and would error rather than
skip on precisely the layouts the fallback they cover was written for.
Co-Authored-By: Claude <noreply@anthropic.com>
a+rwX on the whole tree left every compiled kernel -rwxrwxrwx in the shipped image, so any uid inside the container could replace a .so the process later dlopens. AITER only needs to *create* files in that directory -- its get_user_jit_dir() check is os.access(dir, W_OK) -- so the write bit belongs on the directories alone. Measured in the image after the change: directory 777, artifact 755. Co-Authored-By: Claude <noreply@anthropic.com>
… cold
PREBUILD_KERNELS=0 ships none of AITER's own modules, and LOADER_MODULES
covered only module_fmha_v3_fwd -- the one that throws at dlopen. The two that
merely stall were missed, so the Dockerfile comment calling the lost prebuilts
"no loss" was wrong on its own terms: it already noted aiter.mla loads one.
Measured in the built image, running tests/rocm/test_mla_aiter.py against a
49-artifact jit dir: it came back with 51. The two AITER built in-process were
module_mla_asm and module_aiter_core, the latter also built by the attention
path, so both sit in front of a first request rather than off to one side.
Both build GPU-free through the same driver path as the variants:
ok module_aiter_core 28s
ok module_mla_asm 15s
43 s of image build to take that cost out of every container start.
Co-Authored-By: Claude <noreply@anthropic.com>
… stale claims The previous commit narrowed the chmod to directories and claimed that stopped any uid swapping a compiled kernel. It did not: unlink is governed by the directory's write bit, not the file's mode, so a+w on the directory let any uid delete a .so and write its own in its place. Directories now go a+wt -- the sticky bit leaves removal to the owner -- which is what the comment claimed. GPU_ARCHS was a command prefix on the prebuild, so the --print-jit-dir call on the next line ran without it. It resolves the same path either way here, but a chmod aimed at a different directory than the one just populated is the exact silent failure this line exists to prevent; now exported for both. Three claims contradicted by this branch's own measurements: - The tutorial said the native paged probe "fails for every advertised size on both architectures, so this path always gathers". It serves page_size=1024, and 16 for fp8 -- which the new fp8 test depends on, since an fp8 query that lands on flat-gather raises rather than degrading. - refresh_aiter_jitspec's docstring still named the two-component cache tag that this branch replaced with a three-component one. - MLA's ImportError still said "Install a wheel", now the unsupported path. Also trims the soft-cap comment to its conclusion; the per-row numbers are in fa4383e's message, which is where a measurement belongs. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The world-writable installed JIT directory creates a code-loading risk, and the installation recipe and parity test need correction.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Pins AITER to v0.1.21.post2, builds it from source, prebuilds attention artifacts, and migrates affected ROCm shims to the POD tensor ABI.
Changes:
- Adds GPU-free prebuilding for 48 attention variants and three AITER modules.
- Migrates RMSNorm and RoPE shims to
aiter_tensor_t. - Updates cache keys, validation, tests, and installation documentation.
File summaries
| File | Description |
|---|---|
.claude/skills/benchmark-kernel/SKILL.md |
Updates AITER benchmarking guidance. |
CLAUDE.md |
Documents the pinned source build. |
CONTRIBUTING.md |
Documents new Docker build arguments. |
README.md |
Updates the supported AITER version. |
csrc/rocm/aiter_loader.cc |
Updates ABI pin diagnostics. |
csrc/rocm/norm_aiter.cu |
Migrates norm calls to the POD ABI. |
csrc/rocm/rope_aiter.cu |
Migrates RoPE calls to the POD ABI. |
docker/Dockerfile.rocm |
Builds and precompiles AITER. |
docker/prebuild_aiter_attention.py |
Adds the prebuild driver. |
docs/rocm/backends.md |
Documents source installation. |
examples/amd_flashinfer_rocm_tutorial.ipynb |
Updates tutorial requirements. |
flashinfer/jit/rocm/aiter_source.py |
Unifies cache tags and header discovery. |
flashinfer/jit/rocm/aiter_variants.py |
Uses the shared cache tag. |
flashinfer/rocm/aiter_utils.py |
Raises the ABI floor and adjusts probes. |
flashinfer/rocm/arch_caps.py |
Updates soft-cap evidence. |
flashinfer/rocm/mla.py |
Updates installation guidance. |
flashinfer/rocm/prebuild_aiter_variants.py |
Clarifies GPU requirements. |
flashinfer/rocm/prefill.py |
Updates validated and defective versions. |
include/flashinfer/rocm/attention/aiter/mha_batch_prefill_args.h |
Updates ABI provenance. |
include/flashinfer/rocm/attention/aiter/mha_fwd_args.h |
Updates ABI provenance. |
tests/rocm/test_aiter_source_probes.py |
Tests cache and header fallback behavior. |
tests/rocm/test_aiter_stream_affinity.py |
Tests non-default stream correctness. |
tests/rocm/test_aiter_version_gate.py |
Tests the new version floor. |
tests/rocm/test_fp8_paged_prefill_aiter.py |
Tests FP8 paged prefill. |
tests/rocm/test_prebuild_aiter_attention.py |
Tests prebuild parity and coverage. |
tests/rocm/test_prefill_decode_dispatch.py |
Stabilizes missing-package testing. |
Review details
- Files reviewed: 26/26 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
…sites too Removing ENV HOME in d56e274 put back the ordering hazard a4daaa5 had fixed: the prebuild resolved AITER's jit directory under whatever HOME the base image left. Rather than manage HOME, assert the outcome. The driver already prints the directory it resolved, so the build now fails loudly if the artifacts land anywhere but the installed tree -- which is the only place a runtime container will look for them. Nothing is made writable to achieve it. The version-assertion heredoc is the one BuildKit-only instruction in the file, so it now declares the frontend instead of relying on the daemon default. _loader_name_tokens() scanned only between dtype_token and build_so_name, which excludes the prefix/infix/suffix literals the three call sites pass in -- the very place _nsink lives. A/B, dropping _nsink from the driver: AssertionError: aiter_loader.cc can emit ['nsink'], never built The old range could not see that at all. Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
Review loop closed at a27cf11Two Copilot rounds, 7 findings, all dispositioned in-thread. Closing here — the recurring item is settled and what remains are declines with evidence. Later rounds re-raising these should be read against this comment. The recurring item: a writable directory in the image. Raised three times, in three different places. Settled by not widening anything: the installed Declined, with what settles each:
Deferred, named rather than carried silently: the What this needs next is a human reader. |
There was a problem hiding this comment.
🔵 Needs a closer look
The driver builds unreachable FP8 LSE artifacts and can oversubscribe compiler resources during parallel image builds.
Review details
Suppressed comments (2)
Previously missed (2) — in code that hasn't changed since the last review.
docker/prebuild_aiter_attention.py:188
- The fp8 branch also iterates
has_lse=True, producing four artifacts that cannot be requested.flashinfer/rocm/prefill.py:3334-3341rejects every fp8 call with LSE before dispatch because AITER has no fp8 LSE instance. These builds therefore add several minutes to every image build while yielding unusable.sofiles; restrict fp8 tohas_lse=Falseand update the expected variant count.
docker/prebuild_aiter_attention.py:405 --jobsstarts multiple AITER builds, but each child independently lets AITER choose about 80% of all host CPUs (and computes its memory limit from the same free-memory snapshot). With the default--jobs 2, this can run roughly twice the intended compiler workload and exceed the host's memory budget duringdocker build. Pass each child a bounded per-processMAX_JOBS(or safely coordinate AITER'sPREBUILD_THREAD_NUM) so total Ninja parallelism is divided across these processes.
- Files reviewed: 26/26 changed files
- Comments generated: 0 new
- Review effort level: Balanced
…ls in it
Both suppressed comments from the third review, and the first proves the second.
fp8 was enumerated over both LSE arms. It has none: prefill.py:3334 raises
"fp8 prefill cannot produce LSE: AITER builds no LSE instance of the fp8 kernel
at any page size", and prefill.py:1019 already excluded it for the store route.
The driver did not, so four of the eight fp8 builds asked CK for an instance set
that does not exist. They compiled, linked, and passed --check:
518,440 ..._logits_nbias_mask_lse_...pertensor_nsink.so <- all four LSE
518,440 ..._logits_nbias_nmask_lse_...
5,130,416 ..._nlogits_nbias_nmask_nlse_... <- all four non-LSE
5,437,120 ..._logits_nbias_mask_nlse_...
Ten times smaller, because they are a dispatcher with nothing behind it. 44
variants now, not 48.
That is the failure mode three reviews called speculative, so _check now looks
for it: an artifact carrying fewer than 200 `ck_tile` markers has no instances.
Measured on the two arms of the same variant -- 20 markers against 8,325 --
which is what sets the threshold. Applied to the CK variants only; the
LOADER_MODULES are built -DENABLE_CK=0 and legitimately have none.
Also bounds each child's ninja parallelism. AITER computes MAX_JOBS from ~80% of
*all* host CPUs and its own free-memory snapshot, per process, so --jobs N
oversubscribed the box N-fold; the children now divide one budget.
Co-Authored-By: Claude <noreply@anthropic.com>
…our others The memory fault this branch found while benchmarking was recorded only in the benchmark skill, but _aiter_paged_route_page_sizes sends bf16/fp16 to exactly the page size that faults, so the note now sits beside that return. Not gated: it reproduces on 0.1.20 as well, and the failing kv boundary has not been bounded -- gating on a guess would disable a path that works below it. _child_env deferred to a pre-set MAX_JOBS, which CLAUDE.md documents as the way to limit a build, so the documented knob defeated the split it was added for. It now divides that budget: MAX_JOBS=4 with --jobs 2 gives each child 2. _aiter_importable no longer probes for AITER's header tree. It answers "can AITER run?", and a shim that is already built or AOT-cached needs no headers; the build path raises with the real cause when they are what is missing. Raised in four separate reviews, each time correctly. The Dockerfile now resolves and checks the jit destination before the 35-minute prebuild rather than after it, _has_ck_instances reads in 1 MB chunks instead of pulling a whole multi-MB artifact into memory, and compose_cache_tag rejects the same character class its test asserts on. Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
Suppressed comments, review 3 (2) — both accepted, fixed in e441b65These have no thread, so recording them here. Both were right, and the first one proves the second.
44 variants now, not 48. This also settles a claim declined three times. That a CK
One thing I got wrong earlier and have corrected in the description: I explained the small fp8 artifact as "its CK instance set is smaller". It was an empty shell. |
There was a problem hiding this comment.
🟡 Changes recommended
The standalone source-install guidance omits a required AITER module prebuild and can leave default ASM prefill failing at load time.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 26/26 changed files
- Comments generated: 3
- Review effort level: Balanced
… have Making the source build the supported install left a gap the wheel did not have. PREBUILD_KERNELS=0 ships no AITER modules, and aiter_loader.cc dlopens module_fmha_v3_fwd directly -- there is no lazy JIT behind it, unlike the mha_fwd variants. The image is fine because the prebuild driver builds it; the standalone recipe in docs/rocm/backends.md stopped short, so a large enough single prefill fails at dlopen. The loader's hint still read "ships prebuilt in the amd-aiter wheel ... its absence means the wheel is incomplete", which sends a source user to reinstall the thing that cannot fix it. It now names the driver and the FLASHINFER_AITER_ASM_PREFILL=0 escape, which is read at include/flashinfer/rocm/attention/aiter/single_prefill.cuh:100. The fp8 test's docstring also still said eight variants and described --check as presence-only; both moved in e441b65. Co-Authored-By: Claude <noreply@anthropic.com>
…ctually work tests/rocm/test_fp8_paged_prefill.py already exists with 20 cases, including test_fp8_paged_prefill_matches_fp32_reference (a tighter reference than the bf16 one added here), test_every_routed_fp8_page_size_is_numerically_right across pages 1/16/1024, test_fp8_rejects_return_lse, and test_single_prefill_fp8_raises_instead_of_a_ninja_log. Both cases added in 7733443 have a direct equivalent there, so the new file was pure duplication on a compile-bound suite, and its docstring's claim that nothing else exercised fp8 paged prefill was simply wrong. Removed. The standalone prebuild recipe added in 714d451 omitted GPU_ARCHS, which the driver never sets -- the flat `from jit import core` import skips the aiter/__init__.py that would populate it, so every build asserts. Measured first-hand earlier in this branch: eight consecutive `FAIL` lines with the variable unset. Both the doc and the driver's usage line now set it. --jobs took int() at parser-construction time, so a non-numeric AITER_PREBUILD_JOBS crashed before --help could run; it is now validated in main() and reports the variable by name. _check runs even when a build failed, so one bad variant no longer hides the hollow artifacts among the other 43. Two stale comments: aiter_source's docstring still taught the at::Tensor forward declaration, which is the trap this branch removes, and the tutorial promised page_size=1024 without the gfx950 fault that prefill.py and the benchmark skill both record. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The standalone recipe currently loses the FlashInfer working directory, and several installation paths or messages remain inconsistent with the source-build workflow.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 4
- Review effort level: Balanced
…source sweep _hip_clang_path claimed to mirror AITER's compile_ops and did not: core.py:1776 guards the export with `os.path.exists(hip_clang_path)`, so an inherited but stale MHA_HIP_CLANG_PATH is ignored by the lazy build and was being turned into an invalid compiler path by the prebuild. Same guard now. The standalone recipe ran `cd aiter && pip install .` unparenthesised, so the prebuild command in the next block resolved inside the AITER clone, where the FlashInfer script does not exist. Wrapped in a subshell and said why. Two more messages that still sent users to a wheel after this branch made the source build the supported install: prefill.py's _require_aiter_runtime, on the explicit-AITER path, and the tutorial's missing-package cell. mla.py's copy was already corrected; these are the rest of the set. Co-Authored-By: Claude <noreply@anthropic.com>
…smaller ones _aiter_csrc_include_dir's AITER_CSRC_DIR fallback was a third `import aiter` site that skipped _ensure_aiter_gpu_archs(). The other two call it first because AITER freezes its architecture at import, so reaching this one with GPU_ARCHS unset could pin the wrong arch for the rest of the process. _child_env sized its budget from os.cpu_count(), which reports the host's cores inside a cgroup-limited `docker build`; sched_getaffinity reports what the builder may actually use. AITER_SOURCE was compared literally to "1", so `--build-arg AITER_SOURCE=true` took the wheel branch and installed a ROCm 7.14 retarget into a ROCm 10 image without complaint. Now rejects anything that is not 0 or 1. The cache tag is checked against a positive character class rather than five denied spellings; a Popen that raises now unlinks its own log, which nothing else would have, since only started children reach `running`; the driver removes AITER's shared build/ root, not just the per-module subdirectory, so it stops shipping in the layer; and _loader_name_tokens asserts start < end, without which reordering the two functions in aiter_loader.cc would empty the slice and silently retire the parity check. Also drops test_fused_add_rmsnorm_aiter_is_correct_on_a_side_stream. Its own docstring recorded that it passes with the guard removed, so it banked coverage it did not have; the two cases that do detect it are A/B'd and stay. Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
Loop closed again at e2fdd10 — five bot rounds, 16 findings, all dispositionedEvery thread is resolved and both suppressed comments answered. The two things still open are declines with measurements behind them, restated here so a later round re-raising them can be read against this. Declined: gate the gfx950 paged fault. Raised three times. It is real and it is recorded in Declined: Declined, repeatedly and on the same evidence: that One I got wrong and want on the record. I added What this needs next is a human reader. |
There was a problem hiding this comment.
🔵 Needs a closer look
The shim headers bypass available compile-time ABI validation, and the reported artifact counts conflict with the implementation.
Review details
Suppressed comments (3)
Previously missed (3) — in code that hasn't changed since the last review.
csrc/rocm/norm_aiter.cu:26
- The pinned 0.1.21.post2
csrc/include/rmsnorm_quant.his POD-only and directly declares these signatures; it no longer includestorch/extension.h. Keeping a handwritten declaration defeats the compile-time ABI check required byflashinfer/jit/rocm/aiter_source.py:8-13, so a future parameter change can again compile and fail only atdlopen. Include the real header instead.
csrc/rocm/rope_aiter.cu:27 - The pinned 0.1.21.post2
csrc/include/rope.hnow contains only the POD declarations and does not includetorch/extension.h. This forward declaration therefore bypasses the compile-time ABI validation prescribed byflashinfer/jit/rocm/aiter_source.py:8-13; include the real header so later signature drift is caught during compilation.
docker/prebuild_aiter_attention.py:182 - The implementation enumerates 44 variant files (four fp8) and three whole modules, so an unrestricted run produces 47 artifacts. The PR description's final image/test sections instead claim 48 variants, eight fp8 variants, and 51 artifacts. Please reconcile those validation results with the final artifact set so the stated end-to-end coverage is accurate.
- Files reviewed: 25/25 changed files
- Comments generated: 0 new
- Review effort level: Balanced
…aring them
The forward declarations existed because rmsnorm_quant.h and rope.h pulled in
<torch/extension.h> and so pybind11, which clashes with -DPy_LIMITED_API. At
0.1.21.post2 that is no longer true -- both are POD-only:
rmsnorm_quant.h 1 include (aiter_tensor.h), 0 torch/pybind refs, 18 aiter_tensor_t
rope.h 3 includes (POD + cstdint), 0 torch/pybind refs, 82 aiter_tensor_t
So the reason for hand-writing the signatures is gone, and keeping them defeats
the compile-time ABI check this module's own docstring prescribes -- the next
signature change would again compile clean and fail at dlopen on the mangled
name, which is the failure this branch exists to remove.
Verified on a cold JIT cache so both shims really recompile: 73 passed across
test_rmsnorm_aiter, test_rope_aiter and test_aiter_stream_affinity.
Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
Suppressed comments, review 6 (3) — all three accepted, fixed in 644d725
So the declarations were defeating the compile-time ABI check for no remaining reason — the next signature drift would have compiled clean and failed at
|
There was a problem hiding this comment.
🔵 Needs a closer look
The hardware-specific ABI migration and final multi-artifact image composition warrant human validation, especially since the final artifact set was not rebuilt end to end.
Review details
- Files reviewed: 25/25 changed files
- Comments generated: 0 new
- Review effort level: Balanced
17 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The image now builds AITER from source at a pinned tag and prebuilds every attention variant, so a container starts with them present instead of compiling them inside the serving process. A fresh container's first pass over the four attention families drops from 949 s to 134 s. The pin moves to
v0.1.21.post2, which is an ABI migration rather than a version bump: 0.1.21 movedrmsnormand the cos/sin-cacheropefromat::Tensorto the PODaiter_tensor_t.What changed
The prebuild
docker/prebuild_aiter_attention.py— new, stdlib-only driver that composes and builds all 44 attention variants with no GPU. It reaches AITER's builder through the flatsys.pathimportaiter/setup.pyitself uses;import aiterruns arch detection and a triton import that both need a device.docker/Dockerfile.rocm— clones and installs AITER from source, asserts the result, then runs the driver.AITER_SOURCE=0restores a wheel install as an unblock.tests/rocm/test_prebuild_aiter_attention.py— new. Three parity checks: the driver's filenames againstaiter_variants.py, its CK receipts and filter tokens against the installedaiter/ops/mha.py, and every name tokenaiter_loader.cccan compose against the set the driver actually builds.The ABI move
csrc/rocm/norm_aiter.cu,csrc/rocm/rope_aiter.cu— migrated to the POD tensor API via the existingaiter_tensor_compat.h, matching the three shims that migrated at 0.1.16.tests/rocm/test_aiter_stream_affinity.py— new, covers a gap the migration exposed (below).flashinfer/rocm/aiter_utils.py— floor to 0.1.21, and the import probe no longer hard-requiresaiter_meta.Supporting fixes
flashinfer/jit/rocm/aiter_source.py—compose_cache_tag()now keys both AITER caches by arch, AITER version and ROCm, from one place.docs/rocm/backends.md,CLAUDE.md,CONTRIBUTING.md,README.md, the tutorial notebook — the source build is now the supported install; the previous text said the opposite.Architecture / design notes
Why the artifacts go in AITER's
jit/directory rather than FlashInfer's store. PR #363 added a variant store thataiter_loader.ccreads atdlopen. That cannot helpmha_batch_prefill: its bootstrap doubles as the page-size probe and runs before the loader, so AITER compiles regardless of what is in our store. Putting the same file where AITER itself looks makes the bootstrap find it built — measured at 101.4 s cold against 0.2 s prebuilt. #363's store is untouched and remains the fallback for wheel users.Why source rather than a wheel. No
amd-aiterwheel targets ROCm 10 at this revision — the index publishes0.1.21.post2only as+rocm7.14.0and+rocm7.2.3. A source build compiles against the ROCm actually in the image.Why 44 variants and not 40. fp8 is reachable, and batch-prefill only.
prefill.pyroutes every fp8 query through the native paged kernel because the flat-gather route runsmha_varlen_fwd, which has no fp8 arm, and the in-tree fa2 kernel rejects 8-bit types at compile time. All three AITER mha families carry an fp8 arm upstream, butkFp8Bf16is constructed at exactly one site in this tree. fp8 adds 4 names rather than 8 because it has no LSE arm — AITER ships no LSE instance of the fp8 kernel at any page size, andprefill.py:3334raises on the combination.aiter_variants.pystays at 40: it serves the store route, which has no fp8 arm.Why the version move and the shim migration are one commit. The POD overloads do not exist at 0.1.20 (
rmsnorm_quant.h:37declaresvoid rmsnorm(torch::Tensor&, ...), andaiter_tensor_tappears in none of the norm or rope headers), so either ordering leaves a commit whose tests fail atdlopen. The bisect split the plan wanted is preserved elsewhere: one commit proves the mechanism at 0.1.20, the next moves the revision.One build per variant, not per file.
aiter_variants.builds()says 32 builds cover 40 files because the bootstraps loopreturn_lsein Python. Drivingbuild_moduledirectly, one call yields one.so; all 16 varlen builds took 65–67 s with none short-circuiting.Benchmark results
Cold start, fresh container from the built image, gfx942. Both arms pay the same FlashInfer JIT cost, so the delta is AITER's.
mha_fwd(single prefill)mha_varlen_fwd(soft cap)mha_varlen_fwd(ragged)mha_batch_prefill(paged)Steady state is unchanged, and now measured rather than asserted. Single prefill on gfx950 (q=256, GQA 16/4, HD=64, causal, fp16),
4d322ab0+ the old wheel image against6d80af63+ the new source image.bench_aiter_prefill.pyis byte-identical between the two trees.Native paged batch prefill at
page_size=1024,kv=1024: 62.98 TFLOPS before, 62.36 after (0.99x). Both figures are inside run-to-run noise, which is the expected result — same sources, same compiler, same kernels.Image build cost, 32 cores at
--jobs 2: 40 variants in 28 min uncontended, plus 80 s for the clone and install; the 8 fp8 variants add ~66 s each. Artifacts are 224 MB. AITER's intermediatejit/build/staging was another 603 MB; the driver now deletes it on success, verified by removing it and confirmingplan()builds nothing.The final set is 47 artifacts: 44 variants plus 3 whole modules — built end to end and verified after merge. From the built image:
total 47 / variants 44 / fp8 4 / whole modules 3, jit dir mode755, AITER'sbuild/tree absent, and every CK variant above the instance threshold (all populated). The earlier 49-artifact build described below predates the fp8 LSE drop and is left as the record of what it measured at the time.End-to-end image build, gfx950 node:
49/49 present, which at that commit was 48 variants plusmodule_fmha_v3_fwd. Two things moved afterwards, neither re-run as a full image build: the fp8 LSE arm was dropped (48 → 44 variants, below), and the whole-module set grew to three.module_aiter_core(28 s) andmodule_mla_asm(15 s) were each verified to build through the same driver path individually. The installed tree stays read-only (0755); a container started as an unrelated uid takes AITER's copy-to-~/.aiterfallback, measured at 226 MB in 97 ms, and sees every artifact. An earlier revision made the tree world-writable to skip that copy — reverted, because it put a writable directory on Python's import path for 97 ms of saving.Why those two were added.
PREBUILD_KERNELS=0ships none of AITER's own modules. Runningtests/rocm/test_mla_aiter.pyin the built image took the jit dir from 49 artifacts to 51: AITER compiledmodule_mla_asmandmodule_aiter_corein-process. The second is on the attention path too, so both sat in front of a first request.Correctness control. A driver-built varlen variant against the one AITER builds lazily: byte-identical size (3,123,528), 99 non-
__hip_cuid_*symbols identical. Every prebuilt variant carries both architectures — offload bundle idsgfx942,gfx950on the bf16 and fp8 batch-prefill artifacts alike.An artifact can link with no kernel in it, and
_checknow catches that. A CK--filterthat matches no instances still compiles and still produces a non-empty.so. The fp8 LSE arm was exactly this: 518,440 bytes against 5,437,120 for the same variant's working arm, and 20ck_tilemarkers against 8,325. Those four are no longer built, and_checkrejects any CK variant carrying fewer than 200 markers.Soft-cap defect re-measured at 0.1.21.post2 over non-square shapes with the guard disarmed: gfx950 wrong on all 12 capped rows (9 at 1.19–4.11 absolute, 3 returning NaN) where fa2 is 0.0001–0.0011; gfx942 clean on all 12. The gate is unchanged; only its evidence string moves.
Pre-existing defect found while benchmarking, not introduced here
AITER's native paged batch prefill faults on gfx950 past
kv=1024atpage_size=1024:Memory access fault by GPU node-1 ... FmhaBatchPrefillWithPagedKVCacheKernel. It reproduces identically on both arms —4d322ab0with aiter 0.1.20 and6d80af63with 0.1.21.post2 — so it is not a regression from the version move, and it is why the paged half of the benchmark above reports a single row. Recorded here rather than fixed: it is an AITER kernel fault and out of scope for this PR.Test plan
tests/rocm/full suite, gfx942 (MI300X): 29,983 passed, 3,722 skipped, 0 failed of 33,705tests/rocm/full suite, gfx950 (MI350X): 29,927 passed, 3,778 skipped, 0 failed of 33,705pin-run.shdetached worktree, verified un-drifted afterwardsrmsnorm/ropeAITER tests, which fail atdlopenon 0.1.21.post2 before this branchaiter_loader.cc can emit ['fp8bf16', 'pertensor'], never built; reverting the--onlyfix fails with8/9 presentrel=0.0505); newtests/rocm/test_fp8_paged_prefill_aiter.py, 2 passed on gfx95047/47 present— 44 variants (4 fp8) plus 3 whole modules, read-only 0755 jit dir, no leftoverbuild/tree, no hollow artifactsAITER_SOURCE=1andAITER_SOURCE=0AITER_SOURCE=0with a sub-floor wheel fails the build loudly rather than shipping a silently-degraded image~/.aiter/jitin 1.51 s with all artifacts visibledocker build --checkclean on the final Dockerfile (only pre-existingUndefinedVarwarnings, lines 140-141)/code-review xhighon the changelistpre-commit runon every changed file