Reach AITER's asm prefill arm, and route to it on CDNA4 - #365
Merged
Conversation
Resolving an AITER dispatcher symbol does not mean you get its arms. AITER splits one entry point across modules built with different flags: aiter::mha_fwd tries asm then CK-Tile in source (csrc/cpp_itfs/mha_fwd.cu), but module_mha_fwd compiles -DFAV2_ON=1 only and the asm arm lives in a separate module_fmha_v3_fwd (-DFAV3_ON=1 -DENABLE_CK=0). Measured on amd-aiter 0.1.20+rocm10.1.0a20260819.3135022: `strings` on the shipped mha_varlen_fwd variant .so gives 0 hits for fmha_fwd_v3 and 0 for AiterAsmKernel, and objdump shows aiter::mha_fwd reading byte 1 of its args (v3_api_check) and never byte 0 (use_asm_v3). The subsequent commits fix the dead store that fact leaves in single_prefill.cuh. Also records the rule the miss came from: bind AITER in csrc/rocm/, never through aiter.ops.* from the Python API, which is for capability probes, JIT bootstrap and backend routing only. mla.py is the one remaining exception. The python one-liner in CLAUDE.md was run against the pinned wheel; it prints ["'-DFAV2_ON=1'"]. Co-Authored-By: Claude <noreply@anthropic.com>
…sproved The variant .so the loader dlopens is built -DFAV2_ON=1 with no -DFAV3_ON, so the asm arm is not in that binary at all -- 0 hits for fmha_fwd_v3 under `strings` on amd-aiter 0.1.20. Every statement that a soft cap, or anything else, "disables AITER's asm paths" has the causality backwards: CK Tile was always what ran, so there was nothing to disable. test_single_prefill_aiter_bf16 claimed to exercise "the ASM v3 (bf16+hd128) fast path". It never did, and its qo_len values top out at 577, below the threshold the later commits add, so it still will not. Renamed to what it actually covers and pointed at the test that does cover the asm arm. The other three are comments describing routing as "the asm path" where the distinction being tested is whether the soft-cap gate fires. Reworded to say that instead, which is both true and what the assertions check. PR #361 corrected the same premise in _aiter_softcap_defect's docstring independently, so nothing here touches prefill.py. Co-Authored-By: Claude <noreply@anthropic.com>
Neither existed. aiter_loader.cc could only reach the CK-Tile variant .so files, and nothing under csrc/rocm or include/flashinfer/rocm/attention read gcnArchName at all -- the only device queries in the tree are the CU count and the LDS sizes in gpu_runtime_compat.hpp. get_aiter_mha_fwd_asm_handle() opens module_fmha_v3_fwd.so, which exports the same mangled aiter::mha_fwd symbol the CK variants do (verified with nm -D on the pinned wheel), so the call site needs no second function pointer type. It takes no VariantKey: the module ships prebuilt and is trait-generic, since AITER resolves dtype/hdim/mask/mode from its own config table per call. That also means there is no JIT bootstrap to trigger, so the hint text deliberately does not send the reader off to build one. The AITER_ASM_DIR guard is not defensive padding. AITER reads its asm kernels as .co files from that directory and reaches std::abort() via AITER_CHECK when it cannot, which would take the host process down with no traceback. The variable is a side effect of importing the aiter Python package, which no pure-C++ consumer guarantees, so the loader refuses up front and lets the caller stay on CK Tile. getGcnArchName strips at the first ':' so gfx950:sramecc+:xnack- matches what arch_caps.normalize_arch produces on the Python side, which is where the routing policy is keyed. AITER has the same helper in aiter_hip_common.h but its headers pull in pybind11, which cannot be included from a -DPy_LIMITED_API translation unit -- the same constraint norm_aiter.cu already works around by forward-declaring. Syntax-checked both TUs against hipcc/ROCm 10.0 in the dev image. Co-Authored-By: Claude <noreply@anthropic.com>
Sibling of _AITER_SOFTCAP_DEFECT_MIN_KV_LEN and shaped the same way: a private
dict keyed on the bare gfx name plus a public accessor taking an arch string,
returning Optional[int] through normalize_arch. It is not a KnownBad row for
the same reason that one is not -- KnownBad gates a whole (op, backend, arch)
on toolchain version, and this gates a shape range.
The asymmetry is the point. asm needs enough q-side parallelism to fill the
device, and the two architectures reach that differently. Over a 60-cell sweep
(batch 1 and 4 x 16/32/64 q-heads x seqlen 256..6144, bf16 head_dim 128,
A/A-controlled at +/-2%):
gfx950 per-seqlen geomean is monotone above 1024; the 24 cells at or above
2048 give geomean 1.210 with the worst cell at 1.00. Re-verified on
an idle node at geomean 1.222, min 1.01.
gfx942 non-monotonic. 1.34x at seqlen 1024 drops to 0.90x at 1536 and
recovers, and both reproduce against the noise floor. The same
threshold there scores geomean 1.032 with 8 of 24 cells regressing
and a worst case of 0.80, so gfx942 gets None.
None therefore means "never route to asm", and covers an unrecognised arch as
well as a measured loss -- the safe answer is the same for both.
qo_len is the gate variable, not kv_len: the gfx950 prefill-with-history shape
(qo 512, kv 4096) measures 0.73x, and gating on qo_len excludes it while still
admitting the 4096 non-causal case at 1.24x.
README regenerated from the note; the arch-support-matrix hook enforces it.
Co-Authored-By: Claude <noreply@anthropic.com>
Replaces the dead store. args.use_asm_v3 was computed from AiterAsmV3Eligible and then handed to a .so built -DFAV2_ON=1 with no -DFAV3_ON, which does not contain the code that reads it. The flag now selects between two modules instead of being ignored by one. Dispatch order is pre-filter, threshold, probe, launch: AiterAsmV3Eligible stays as the cheap pre-filter rather than the authoritative test. AITER's own v3_api_check is authoritative -- it resolves the real config table and cannot drift from it -- but it needs the asm .so dlopened and fmha_fwd_v3 logs a warning on every rejection, so screening fp16 and hd64 callers first keeps both costs off the common path. The probe runs once per trait set, not per call. Only needs_mask varies the lookup here: head dim is a template parameter, dtype is pinned by the pre-filter, and is_group_mode is false whenever there is no soft cap. The error contract is the part worth reading twice. The asm module is built -DENABLE_CK=0, so it has two distinct failure modes and neither is a hipError_t: a config-table miss returns a negative value having launched nothing, and a failed launch throws out of ck_tile_shim's launch_and_check after that code has already consumed hipGetLastError(). So the asm arm is wrapped in try/catch, a post-asm hipGetLastError() is never trusted, and every failure falls through to CK Tile rather than surfacing an AITER-internal message in place of a working fallback. getGcnArchName is inside a try for the same reason: FI_HIP_CALL throws, and this function returns hipError_t. Checking the CK return closes a hole that predates this change: a negative return there also means nothing launched, which left the caller's output buffer untouched while hipGetLastError() stayed clean. It now reports hipErrorNoBinaryForGpu, and the TORCH_CHECK distinguishes "nothing ran" from "the launch failed". FLASHINFER_AITER_ASM_PREFILL=0 pins CK Tile. AITER reaches std::abort() via AITER_CHECK on several asm failure modes, so an operator-side off switch is part of the contract rather than a convenience. Built cold on gfx942 and ran a 2048x2048 bf16 causal prefill: the CK variant .so loads and the asm arm is correctly not taken, since the gfx942 threshold is 0. Co-Authored-By: Claude <noreply@anthropic.com>
First coverage, not a regression test: the arm was unreachable before the previous commit, and test_single_prefill_aiter_bf16 -- which claimed to exercise it -- tops out at qo_len 577, below the threshold. test_asm_gate_numerics straddles the threshold deliberately. 512 is CK Tile on every architecture, 2048 and 4096 take asm on gfx950, and (512, 4096) pins the gate to qo_len rather than kv_len: that shape measured 0.73x, so routing it to asm would be a regression, and a gate keyed on kv_len would do exactly that. return_lse is the case that needed a test rather than a benchmark. Every measurement behind this change ran return_lse=False, so the asm arm's LSE was unverified; it comes from a different kernel than CK Tile's while single_prefill_aiter.cu divides the result by log(2) unconditionally, so a differing base would corrupt every LSE silently. naive_attention already returns log2 LSE in FlashInfer's layout, so the reference is reused rather than rewritten. test_cuh_mirrors_arch_caps guards the one duplicated constant. It is a substring check on a marker comment, matching the shape of test_aiter_version_gate.py's header check -- parsing the C++ literal would fail on reformatting rather than on a real divergence. The kill-switch test spawns subprocesses because the switch is read once into a function-local static, so it cannot be toggled in-process. On gfx942 it doubles as the assertion that the gate never fires there. Co-Authored-By: Claude <noreply@anthropic.com>
Modelled on bench_norm.py rather than bench_aiter_prefill.py: this is a routing decision, so it needs an A/A noise floor and a two-arm ratio, not a single-arm roofline. It measures AITER's two kernels against each other directly instead of through single_prefill_with_kv_cache. The shim's overhead is identical on both arms and the threshold is a claim about the kernels, so that is the unit the decision is made in. It also sidesteps FLASHINFER_AITER_ASM_PREFILL being read once into a C++ static, which makes a single-process A/B through the shim impossible. The sweep is two-dimensional because asm needs q-side parallelism to fill the device: a seqlen-only sweep at one batch/head count yields a threshold that does not generalise, since batch 1 with 16 heads and batch 4 with 64 heads sit at opposite ends of the same effect. That is also why the summary prints the count of cells below 1.00 rather than just a geomean -- the geomean alone would have hidden the eight regressions that rule gfx942 out. --accuracy is separate because the arms do not agree bit-for-bit and on gfx942 the asm kernel is the more accurate of the two, which a timing table hides. how_v3_bf16_cvt is pinned to 0 to match what single_prefill.cuh sends. It is not a free knob: it selects a different .co on gfx942, and get_kernel_name_key ignores it entirely on gfx950. Co-Authored-By: Claude <noreply@anthropic.com>
The backends.md note keeps the house shape for a measurement-justified routing decision: counts per architecture, a mechanistic reason rather than just a ratio, the arch divergence stated outright, and a literal re-run command with the read-the-A/A-first instruction. The mechanism is the part worth writing down. asm tiles 256 rows of Q at a time, so at batch 1 with few heads there are not enough workgroups to fill the device and it loses. That is why gfx942's surface is non-monotonic rather than simply worse, and why the gate reads qo_len rather than kv_len -- a prefill-with-history shape (qo 512, kv 4096) measures 0.73x on gfx950 and has to stay on CK Tile even though its kv_len is well past the threshold. README gets a hand-written callout beside the soft-cap one, in the same shape: bold lede, mechanism in two clauses, absolute link into backends.md. The generated matrix row above it came from the arch_caps note in an earlier commit. Co-Authored-By: Claude <noreply@anthropic.com>
Five of these are mine from the preceding commits; the review is the reason they are not in the PR. The try/catch could not do what its comment claimed. AITER reaches std::abort() through AITER_CHECK for a missing .co, a failed hipModuleLoad and most other internal failures, because the thread_local that would make it throw defaults to false and is per-.so under RTLD_LOCAL. Only this shim's own exceptions are catchable. The comment now says so, and the mitigations are the ones that actually work: the loader's AITER_ASM_DIR pre-check, the new capture guard, and the kill switch. Graph capture is now excluded outright. The first asm call loads a .co lazily, HIP rejects a module load during capture, and AITER turns that into an abort -- so a vLLM/SGLang worker capturing a long prefill would have died. This is the open question the plan deferred; hipStreamIsCapturing is the answer. test_cuh_mirrors_arch_caps was theatre. It grepped a marker comment, so editing `return 2048u` to 4096 kept it green -- the single failure it existed to catch. The threshold is now a macro that is both what the test reads and what the function returns. A/B'd: drifting the macro to 4096 fails the test, restoring it passes. Nothing could observe which arm ran, so every test in the new file passed vacuously if asm was never reached -- the numerics are identical either way. FLASHINFER_AITER_ASM_VERBOSE=1 now names the arm, and two tests assert on it instead of on a scalar sum that cancelled ~8M near-zero terms. Also: memoize a failed handle load, since swallowing the exception turned a one-off into a per-call tax on the hot path; demote the probe when the real call disagrees with it; normalize a negative window sentinel to -1, as -2 would miss AITER's config lookup and poison the probe cache for later valid calls; return hipGetLastError() on the asm path as the CK path does; drop the test's head count to 8 so the fp32 reference at 4096 is ~1.6 GB rather than ~6.3 GB on a card several xdist workers share; and raise the benchmark's _MAX_ELEMS, which pruned b4/hq64/s6144 and so swept 59 cells while the docs cited 60. Records one thing measurement has not settled: the sweep is square apart from a single point, so chunked prefill (qo_len >= 2048 over a long cached context) takes the asm arm on square-shape evidence.
There was a problem hiding this comment.
🟡 Changes recommended
The routing gate exceeds the measured performance surface, and key numerical, arm-selection, and graph-capture tests need strengthening.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Routes eligible gfx950 single-prefill workloads to AITER’s asm kernel while retaining CK Tile fallback behavior.
Changes:
- Adds architecture-aware asm dispatch, loading, probing, and fallback handling.
- Adds routing policy, benchmarks, and numerical/routing tests.
- Documents performance evidence, controls, and AITER integration guidance.
File summaries
| File | Description |
|---|---|
csrc/rocm/aiter_loader.cc |
Loads the standalone AITER asm module. |
csrc/rocm/single_prefill_aiter.cu |
Improves unmatched-kernel error handling. |
include/flashinfer/rocm/attention/aiter/aiter_loader.h |
Declares and documents the asm loader. |
include/flashinfer/rocm/attention/aiter/single_prefill.cuh |
Implements asm routing, probing, and fallback. |
include/flashinfer/rocm/gpu_runtime_compat.hpp |
Adds cached architecture detection. |
flashinfer/rocm/arch_caps.py |
Defines the gfx950 routing threshold. |
tests/rocm/test_aiter_asm_routing.py |
Adds policy, routing, and numerical coverage. |
tests/rocm/test_single_prefill_kernels.py |
Corrects outdated asm-path descriptions. |
benchmarks/rocm/bench_asm_vs_cktile.py |
Adds comparative performance and accuracy sweeps. |
README.md |
Documents routing and environment controls. |
docs/rocm/backends.md |
Records dispatch rationale and measurements. |
CLAUDE.md |
Documents AITER dispatcher-module behavior. |
.claude/skills/add-rocm-kernel/SKILL.md |
Adds guidance for AITER loading styles. |
Review details
Suppressed comments (1)
tests/rocm/test_aiter_asm_routing.py:133
- These numerical cases can all pass through the CK Tile fallback, so they do not establish the stated asm coverage for non-causal calls or
return_lse=True. The separate reachability test only observes a causal, no-LSE invocation. Please make arm selection observable for each above-threshold trait combination (for example, extend the subprocess probe to parameterize causality and LSE) and assert that"launched"was emitted before accepting its numerics.
res = flashinfer.single_prefill_with_kv_cache(
q, k, v, causal=causal, backend="aiter", return_lse=return_lse
)
- Files reviewed: 13/13 changed files
- Comments generated: 3
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
naive_attention does not upcast -- the existing prefill tests all pass q.float(), k.float(), v.float(), and this file passed bf16. The reference was therefore computed in the dtype under test, which hides exactly the kernel error the comparison exists to catch. Reachability was only asserted for the causal, no-LSE call, so the other three trait combinations in test_asm_gate_numerics could each have been served by CK Tile while still passing -- the two arms agree numerically, which is the whole reason an explicit signal exists. The probe is now parametrized over causal and return_lse and asserts "launched" for each. Adds the graph-capture regression test the guard was missing. It runs cold and in a subprocess on purpose: the guard exists because AITER loads its .co on the first asm call and turns HIP's rejection of a module load during capture into std::abort(), so a broken guard takes the process down rather than raising, and only the subprocess exit code catches that. Co-Authored-By: Claude <noreply@anthropic.com>
Collaborator
Author
|
Suppressed comment (tests/rocm/test_aiter_asm_routing.py:133) — fixed in 0f0e2e0: the reachability probe is now parametrized over causal and return_lse and asserts |
There was a problem hiding this comment.
🟡 Changes recommended
The new loader references undefined helpers and contains additional fallback and dispatch issues.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 4
- Review effort level: Balanced
The merge of amd-integration replaced get_jit_dir()/load_and_cache_sym() with jit_dir_candidates()/load_variant_sym(), but kept this PR's call sites, so aiter_loader.cc no longer compiled: error: use of undeclared identifier 'get_jit_dir' error: use of undeclared identifier 'load_and_cache_sym' load_variant_sym() is the right replacement rather than load_path_sym(): the asm module lives alongside the variants and should be searched across every configured store, not resolved against one baked directory. Also treat an empty AITER_ASM_DIR as unset, matching the env_dir helper a few lines up. `export AITER_ASM_DIR=` otherwise passed the guard, and AITER then roots its .co path at "/<arch>/..." and aborts the process on the failed load -- exactly what the guard exists to prevent. hipcc -fsyntax-only: clean after, two errors before. Co-Authored-By: Claude <noreply@anthropic.com>
hipGetDevice, getGcnArchName and hipStreamIsCapturing all ran unconditionally, so every AITER single-prefill paid them -- including fp16, hd64, windowed and below-threshold callers that can never take the asm arm. The && in asm_wanted short-circuits, but the queries had already happened above it. That also contradicted the dispatch order the PR documents. Now: traits, then arch and threshold, then capture status. Co-Authored-By: Claude <noreply@anthropic.com>
The mirror test asserted `return <macro>;` appeared somewhere in the file, which swapping the gfx942 and gfx950 branches still satisfies -- both macros are still returned, just by the wrong arch. That reverses the routing policy silently. A/B against the swap: fails with the new regex, passed with the old one. This is the second time this test was too weak in the same way; the first version grepped a marker comment rather than the returned value. Co-Authored-By: Claude <noreply@anthropic.com>
The success path resets args.use_asm_v3 before falling through to CK Tile; the catch path did not, so a throw from asm_fn left it set on the call that follows. Inert today because the CK variant is built -DFAV2_ON=1 and ignores the field, but it re-enters the arm that just threw the moment a variant ships both. <stdexcept> was reached only via gpu_runtime_compat.hpp on the next line, so reordering the includes would have broken the TU. Co-Authored-By: Claude <noreply@anthropic.com>
- `assert "CAPTURE_OK" not in err` compared a stdout marker against stderr, so it could never fail and nothing asserted the graph was captured at all. Now a positive check against stdout. - The probe filename was built from the override KEYS, so the kill switch's on and off runs shared a path; a second process exiting 0 without writing left the first run's tensor to be compared against itself. - The capture probe made its first-ever flashinfer call inside the capture, so a cold JIT cache compiled and dlopened mid-capture. Warm on a below-threshold shape first, which leaves AITER's .co cold -- the actual condition under test. - The subprocess resolved flashinfer independently of the parent; it now prints its import path and the parent asserts they match. Not hypothetical: a stale PYTHONPATH=/wt in a verification container sent a benchmark at base code today. Adds (2048, 8192): above the threshold and non-square, so causal exercises bottom-right masking on the asm arm. No other case reached it, while the docs claimed correctness was covered either way. Co-Authored-By: Claude <noreply@anthropic.com>
The capability note promised asm for "an unwindowed bf16 head_dim 128 call at qo_len >= 2048", but the code also excludes logits_soft_cap (which forces the varlen handle and fails AiterAsmV3Eligible) and HIP graph capture. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟡 Changes recommended
The memoized fallback still incurs repeated exception overhead, and benchmark provenance omits the ROCm version.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
Previously missed (1) — in code that hasn't changed since the last review.
benchmarks/rocm/bench_asm_vs_cktile.py:112
- The provenance header omits the ROCm runtime version, so results from different ROCm stacks can look comparable even though kernel behavior and timing are version-sensitive. Record
torch.version.hipalongside the PyTorch and AITER versions to make the threshold evidence reproducible.
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Balanced
Collaborator
Author
|
Suppressed comment (benchmarks/rocm/bench_asm_vs_cktile.py:112) — fixed in 91356c6: the provenance header now records |
There was a problem hiding this comment.
🟡 Changes recommended
The asymmetric causal test can silently use CK Tile and therefore does not verify the claimed asm coverage.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟡 Changes recommended
The tests do not verify ASM-produced LSE values or prove that asymmetric chunked-prefill inputs actually reach the ASM arm.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (1)
tests/rocm/test_aiter_asm_routing.py:169
- This probe always makes K/V length equal to
qo_len, sotest_asm_arm_is_actually_reachedonly proves routing for square inputs. The(2048, 8192)numerical case elsewhere can silently fall back to CK Tile, even though the test and documentation say chunked prefill reaches asm. Parameterize the probe's KV length and add an above-threshold asymmetric case that asserts thelaunchedsignal.
q = torch.randn(qo, 8, 128, dtype=torch.bfloat16, device=d)
k = torch.randn(qo, 2, 128, dtype=torch.bfloat16, device=d)
v = torch.randn(qo, 2, 128, dtype=torch.bfloat16, device=d)
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Balanced
The (2048, 8192) case added last commit sat in test_asm_gate_numerics, which
runs in-process and cannot see which arm served it -- so the bottom-right-mask
coverage it was supposed to give could have been CK Tile's. _PROBE also built
K/V at kv_len == qo_len always, and discarded the LSE, so neither the asymmetric
shape nor the LSE path was ever observed on asm.
The probe now takes a KV length and saves {out, lse}; the reachability test
parametrizes kv_mult 1 and 4 and compares both against an fp32 reference in the
same execution that emitted the `launched` marker.
LSE is the case the module docstring calls out: the asm kernel writes it from a
different path than CK Tile and single_prefill_aiter.cu divides by log(2)
unconditionally, so a differing base corrupts it silently.
Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🔵 Needs a closer look
Cross-library asm dispatch has unrecoverable abort paths, while final dual-architecture regression and benchmark validation remain pending.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Balanced
Collaborator
Author
|
Suppressed comment (tests/rocm/test_aiter_asm_routing.py:169) — stale, that review ran against 91356c6. Fixed in 679c890: the probe takes a KV length and saves the LSE, and |
There was a problem hiding this comment.
🟡 Changes recommended
The default-enabled ASM path can still abort the process when its module and code-object directory are mismatched or incomplete.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 1
- Review effort level: Balanced
There was a problem hiding this comment.
🟢 Approval recommended
The updated implementation addresses prior findings and provides targeted coverage for routing, fallback, graph capture, asymmetric contexts, and LSE correctness.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Balanced
There was a problem hiding this comment.
🟢 Approval recommended
The implementation includes defensive fallbacks and focused coverage for routing, capture safety, asymmetric contexts, LSE correctness, and operator disablement.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Balanced
AITER opens its .co with AITER_CHECK, which calls std::abort() rather than returning, so a stale or incomplete AITER_ASM_DIR took the process down before the CK Tile fallback could run. The existing guard only rejected unset/empty. Our eligibility pins arch, dtype, head dim and mask, so the .co name is fully determined -- stat it and stay on CK Tile when it is absent, memoized in the same per-trait slot as the probe. A/B on gfx950 with AITER_ASM_DIR pointed at an empty directory: guard disabled: EXIT=134, Aborted (core dumped) guard enabled: "code object missing under AITER_ASM_DIR; using CK Tile", exit 0 Note the variable has to be set *after* `import aiter`, which overwrites it with the installed tree unconditionally -- the first attempt at this A/B tested nothing because the value never reached the gate. Co-Authored-By: Claude <noreply@anthropic.com>
There was a problem hiding this comment.
🟢 Approval recommended
The revised implementation addresses prior findings and provides focused coverage for dispatch, fallback, numerics, capture safety, and configuration controls.
Review details
- Files reviewed: 13/13 changed files
- Comments generated: 0 new
- Review effort level: Balanced
4 tasks
demandal25
added a commit
that referenced
this pull request
Sep 13, 2026
#365 landed while the previous merge was being verified. Only README.md conflicted; it is generated, so it was regenerated from arch_caps.py rather than hand-resolved, and `scripts/gen_arch_support_matrix.py --check` passes. aiter_loader.cc auto-merged cleanly -- #365's asm handle uses a fixed module name and never calls build_so_name, so this branch's fourth argument does not reach it. Also folded in the findings a `/code-review xhigh` raised against the *first* merge: **#362's probe-site test had been silently retired.** It picked its page size from `_aiter_native_page_sizes()`, which the previous merge stopped being the predicate that disarms the gate: native set : [1, 16, 1024] route set : [1024] test picks page_size = 1 does the selector disarm the gate at that page? False So the selector returned fa2 and the probe site never ran, while all four assertions still passed at the selector site. It now picks from the route set; A/B confirms it: disabling the probe-site re-check fails it, where before the fix that deletion was invisible. **fp8 at a non-routed page size with a short query raised the wrong message.** The short-query gate reaches fa2 first, so the caller was told fp8 "has no in-tree fa2 kernel ... use BatchPrefillWithPagedKVCacheWrapper" -- advice they were already following. The page-size refusal now runs first: NotImplementedError: fp8 prefill ... is not supported on the paged flat-gather route The remaining edits are small: the ragged fp8 guard moves above the jit-module split so a supplied module is covered too; the two now-unreachable `_reject_fp8_on_fa2` re-guards are dropped; `_aiter_paged_route_page_sizes` gains `functools.cache` (three calls per plan, each allocating a frozenset, on a per-step path); two tests stop mutating `global PAGE` when the helper already takes `page=`; and three comments that still justified a page size by the capability set now name the route set, which is how the retired test happened. Declined: making `_native_fp8_dtype` device-aware (flashinfer is one GPU per process), the partial `_backend` mutation before a guard raises, and the dead `scale_q/k/v` triple in `aiter_paged_run` -- all pre-existing and out of scope for a conflict resolution. Co-Authored-By: Claude <noreply@anthropic.com>
demandal25
added a commit
that referenced
this pull request
Sep 13, 2026
…370) ## Summary The asm prefill speedup published by #365 was measured with a benchmark that released the allocator cache between timing the two arms, so the second arm paid `hipMalloc` inside its own measurement. #365 fixed the harness but shipped the numbers taken before the fix. These are the numbers the corrected harness gives, re-measured at `e802be34e`. ## What changed - **`docs/rocm/backends.md`** — gfx950 gate-region figures, the per-`seqlen` curves for both architectures, the A/A floor they are judged against, and the non-square correctness claim. - **`README.md`** — the headline 1.21× becomes 1.17×. - **`flashinfer/rocm/arch_caps.py`** — the policy comment behind the threshold, which carried every superseded claim and is the single source of truth for it. - **`benchmarks/rocm/bench_asm_vs_cktile.py`** — docstring only; it told readers to disregard the gfx950 `s=256` column as noise, which the tightened floor no longer supports. No code changes; the routing gate and its threshold are unaffected. ## Architecture / design notes The threshold does not move. 23 of 24 gate-region cells still win or hold and the geomean is comfortably above 1, so `qo_len >= 2048` on gfx950 remains the right cut. What changes is that the docs can no longer claim a clean sweep: one admitted cell is a small loss, and saying so is the difference between a measurement and a sales figure. Two claims in the same paragraph turned out to need re-deriving, not just renumbering — they were written against the old harness and nothing in the first pass re-checked them: | claim | corrected | | --- | --- | | gfx950 "per-`seqlen` geomean rises monotonically above 1024" | 1.01, 1.09, 1.15, **1.14**, 1.19, 1.19 — a step down at 3072, inside that length's A/A floor of 0.995-1.003 | | gfx942 "1.34× at `seqlen` 1024 falls to 0.90× at 1536" | those points are **0.955** and **0.971**; the non-monotonic conclusion survives on different data (1.08 at 2048, 0.97 at 3072, 1.06 at 4096, 0.99 at 6144) | The non-square hedge also goes. #365 opened by saying correctness on chunked-prefill shapes was "covered either way" and closed by adding `kv_len == 4 * qo_len` coverage on the asm arm — the doc still carried the opening position. ## Benchmark results AITER asm vs its CK Tile arm, bf16 `head_dim` 128, 60 cells (batch 1 and 4 × 16/32/64 q-heads × seqlen 256-6144), causal. Gate region is `qo_len >= 2048`, 24 cells. `single_prefill_with_kv_cache` serves one request and the dispatcher pins `args.batch = 1`, so only the batch-1 half of the sweep is reachable through it. The batch-1 column is the one the routing decision rests on; the mixed column is what #365 published. **gfx950 / MI350X** | | geomean | worst cell | cells below 1.00 | | --- | ---: | ---: | ---: | | published by #365 (mixed, 24 cells) | 1.210 | 1.00 | 0 | | corrected harness, mixed 24 cells | 1.166 | 0.975 | 1 | | **corrected, reachable batch-1, 12 cells** | **1.138** | **0.975** | **1** | The regressing cell is `b1 hq16 s2048` at 0.9752, outside the batch-1 A/A floor of 0.995-1.008. **gfx942 / MI300X** — conclusion unchanged, and the reachable subset makes it stronger: the gate region is a net *loss* there. | | geomean | worst cell | cells below 1.00 | | --- | ---: | ---: | ---: | | published by #365 (mixed, 24 cells) | 1.032 | 0.80 | 8 | | corrected harness, mixed 24 cells | 1.024 | 0.807 | 9 | | **corrected, reachable batch-1, 12 cells** | **0.976** | **0.807** | **8** | **A/A control, same box and session.** This is the part that makes the 0.975 cell meaningful rather than noise: | | A/A range (whole map) | gate-region geomean | | --- | --- | ---: | | old harness, gfx942 | 0.903 - 1.099 | 0.995 | | corrected, gfx942 | **0.987 - 1.013** | 0.998 | | corrected, gfx950 | **0.989 - 1.012** | 1.000 | Removing `empty_cache()` from the per-arm loop cut the A/A spread from roughly ±10% to ±1.3%, which is the direct evidence that the old placement was injecting the noise rather than merely being untidy. Reproduce: `python benchmarks/rocm/bench_asm_vs_cktile.py --aa` then without `--aa`. ## Test plan - [x] gfx942 — A/A, sweep and the seven-file prefill suite at the pinned commit, suite `rc=0` - [x] gfx950 — A/A and sweep at `e802be34e`; suite running, will update if it is not clean - [x] Numbers in the diff cross-checked against the raw sweep logs rather than retyped - [x] `pre-commit run` on both changed files --------- Co-authored-by: Claude <noreply@anthropic.com>
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
args.use_asm_v3insingle_prefill.cuhwas a dead store: the variant.sothe loader dlopens is built-DFAV2_ON=1with no-DFAV3_ON, so AITER's asm arm is not in that binary and every AITER prefill number this port has recorded came off CK Tile. This reaches the asm module for real and routes to it where measurement says it wins —qo_len >= 2048on gfx950 only, ~1.21× — while gfx942 stays on CK Tile because its win/loss surface is non-monotonic and no threshold holds.What changed
Reaching the arm
csrc/rocm/aiter_loader.cc,include/flashinfer/rocm/attention/aiter/aiter_loader.h—get_aiter_mha_fwd_asm_handle()opensmodule_fmha_v3_fwd.so, which exports the same mangledaiter::mha_fwdsymbol as the CK variants. NoVariantKey: the module ships prebuilt in the wheel and is trait-generic. Guards onAITER_ASM_DIR, which AITER needs to find its.cofiles and aborts without.include/flashinfer/rocm/gpu_runtime_compat.hpp—getGcnArchName(), cached per device. Nothing undercsrc/rocmorinclude/flashinfer/rocm/attentionread the arch before this.include/flashinfer/rocm/attention/aiter/single_prefill.cuh— the gate, the probe, and the fallbacks.Policy and evidence
flashinfer/rocm/arch_caps.py—_AITER_ASM_PREFILL_MIN_QO_LENplus its accessor, shaped like the soft-cap pair next to it. Single source of truth for the threshold.benchmarks/rocm/bench_asm_vs_cktile.py— the re-runnable evidence, modelled onbench_norm.py(A/A noise floor, provenance header, warm-both-arms-then-time).tests/rocm/test_aiter_asm_routing.py— first coverage of the arm, includingreturn_lse.docs/rocm/backends.md,README.md,CLAUDE.md,.claude/skills/add-rocm-kernel/SKILL.md— the measurement, the two new env vars, and the rule this miss came from.Four in-tree statements asserting the asm arm was reachable are corrected; PR #361 fixed a fifth independently while this was in flight.
Architecture / design notes
Why the threshold lives in two places. Policy is in
arch_caps.py; C++ carries a macro mirror because there is no way to hand an integer from Python to this call site without widening arun()signature that is generated from a jinja shared with the fa2 path.test_cuh_mirrors_arch_capscompares the Python table against the macros the function actually returns — an earlier version compared against a marker comment, which would have passed while the returned value said something else.Dispatch order. Cheap pre-filter → arch/
qo_lenthreshold → AITER's ownv3_api_checkprobe → launch. The probe is authoritative because it resolves AITER's real config table and cannot drift from it; the pre-filter stays becausefmha_fwd_v3logs a warning on every rejection and probing needs the.sodlopened, neither of which should happen for an fp16 or hd64 caller.The error contract is the load-bearing part. The asm module is built
-DENABLE_CK=0and has failure modes that are nothipError_t:.so/ unsetAITER_ASM_DIRstd::runtime_errorfrom this shim.co, failedhipModuleLoadstd::abort()AITER's
AITER_CHECKaborts rather than throws, because the thread_local that would make it throw defaults to false and is per-.sounderRTLD_LOCAL. That is why graph capture is excluded outright (hipStreamIsCapturing) rather than caught: the first asm call loads a.co, HIP rejects a module load mid-capture, and AITER would take the process down.Gate on
qo_len, notkv_len. The gfx950 prefill-with-history shape (qo512,kv4096) measures 0.73×; gating onkv_lenwould route it to asm.Benchmark results
AITER asm vs its CK Tile arm, bf16
head_dim128, GQA 4, causal. 60 cells (batch 1 and 4 × 16/32/64 q-heads × seqlen 256-6144), A/A-controlled at a ±2% floor.Ratio at and above the shipping threshold (
qo_len >= 2048):gfx950 per-seqlen geomean — monotone above 1024, which is what makes a threshold possible:
† the 256 column did not survive its A/A control (one re-measured cell went 1.82 → 0.98); it is below the threshold and unused.
gfx942 is non-monotonic, which is why it ships nothing. At fixed batch and heads the ratio swings 1.34× at seqlen 1024 → 0.90× at 1536 → recovers, reproduced against the noise floor. Whole-map geomean 0.919, so a blanket swap would be a regression.
Re-verified on an idle node (the sweep ran under load ~200): gate-region cells geomean 1.222, min 1.01, A/A 0.99-1.02.
Reproduce:
python benchmarks/rocm/bench_asm_vs_cktile.py --aathen without--aa.Test plan
test_aiter_asm_routing.py— new; policy tests CPU-only, numerics straddle the threshold,return_lseat and above itLoadKernelonfwd_hd128_bf16_causal.coandfwd_hd128_bf16.co; no asm load below the threshold, none forqo=512 kv=4096, none under the kill switchlse_err = 0.0000against fp32 on the asm armhipcc -fsyntax-onlyon the revised header/code-review xhigh— run per round; 15 findings on9b159d9, 13 more on the review-response commits, of which two were defects in fixes written for the previous roundtest_gated_architecture_really_is_defectivefailures that predate this branch (fixed separately in Let the soft-cap gate test reach the kernel it measures #371). Re-confirmed after merge at5b4924ac5: gfx942 2999/0 failed, gfx950 2944/0 failedtest_aiter_asm_routing.py37 passed / 0 skipped, including the non-squarekv_len == 4 * qo_lencase on the asm armpre-commit run -aKnown gap, stated rather than hidden: the speed evidence is square (
qo_len == kv_len) apart from one point, so chunked prefill takes the asm arm on square-shape timings. Correctness there is no longer a gap — review rounds addedkv_len == 4 * qo_lencoverage that asserts output and LSE against fp32 on the execution that reached asm.FLASHINFER_AITER_ASM_VERBOSE=1shows which arm ran.