Route non-causal windowed prefill to AITER's masked kernel by demandal25 · Pull Request #339 · AMD-Ecosystem/flashinfer

demandal25 · 2026-08-31T22:28:32Z

Summary

AITER prefill silently ignored window_left whenever causal=False: it returned the unwindowed result, 0.66-0.69 abs error against an fp32 windowed reference, on both the flat-gather and native-paged routes and on both architectures. Nothing raised and nothing warned — the numbers were plausible, and only a reference comparison caught it. Found while measuring #337 and deferred to its own PR because it predates that work; it reproduces unchanged on the older ROCm 7.14 / amd-aiter 0.1.16 image.

The defect

Our variant key drove the .so mask axis off causal alone (aiter_loader.cc, name += key.causal ? "_mask" : "_nmask"), so a non-causal windowed request loaded a _nmask binary. Those are compiled with SimplifiedGenericAttentionMask<false> and have no masking logic at all, so window_size_left arrives and has nothing to act on. AITER splits that .so on a broader condition — "is anything masked":

no_mask = (not causal) and (window_size_left == -1) and (window_size_right == -1)
use_mask = not no_mask
                                          # aiter/ops/mha.py:766-769, amd-aiter 0.1.20

We compounded it by setting mask_type = causal ? bottom_right : none at the three args sites, so even the right binary would have been told not to mask.

Design: reuse the key field rather than add one

The obvious shape — add bool has_window to VariantKey — is wrong. build_so_name would only ever OR it with causal, so it could never select a different .so: a dead hash axis whose only effect is breaking every existing caller of the probes. Instead VariantKey::causal becomes needs_mask, fed causal || window_left >= 0. mask_type wants the same OR; only window_size_right wants the real causal. On the Python side that same value goes into the existing parameter, so no signature changes, no new lru_cache axis, and no broken callers.

Merging the two cache entries is correct rather than a collision: a causal=True request and a causal=False, windowed request need the same .so, and the bootstrap exists only to make AITER emit it.

Why `window_size_right` stays `-1`

FlashInfer's window is a left-bound-only, bottom-right-anchored band applied independently of causal (include/flashinfer/attention/variants.cuh:89, mask &= (kv_idx + qo_len + window_left >= kv_len + qo_idx)). CK's make_generic_attention_mask_coordinates_from_lr_window with is_top_left=false saturates a negative right bound to y_total - 1, giving x = x_total — no right constraint — and y = 1 + window_left + qo_len - kv_len, which is that inequality exactly. 0 is the causal convention and is unchanged.

What changed

include/flashinfer/rocm/attention/aiter/aiter_loader.h — VariantKey::causal → needs_mask, in the field, operator== and the hash. BatchPrefillVariantKey is an alias and follows.
csrc/rocm/aiter_loader.cc — build_so_name keys the mask segment off needs_mask; the three dlopen-failure hints no longer print a causal= that would contradict the _mask variant they ask the user to bootstrap.
single_prefill.cuh, batch_prefill.cuh — the three key sites and three args sites. The single-prefill shim also normalizes window_left >= kv_len to -1, as AITER's own torch path does, so a window wider than the sequence stays on _nmask and off a cold ~35-minute CK build. The batch shims cannot: plan() knows only a maximum kv_len. The static_cast<int32_t> on that comparison is load-bearing — params.kv_len is uint32_t, so without it -1 >= kv_len promotes unsigned and strips every window.
flashinfer/prefill_rocm.py — one _aiter_needs_mask() helper, called at all eight sites. Four are the resolved_from_auto probe branches; the other four are the else branches that bootstrap directly, which an explicit backend="aiter" takes. Missing that second group is what would leave a windowed non-causal call failing the C++ dlopen rather than merely being slow.
flashinfer/arch_caps.py, README.md — both prefill rows now advertise sliding window, as batch_decode already did. README is regenerated by scripts/gen_arch_support_matrix.py.

Verification against the installed 0.1.20 wheel

Read out of the container, not a source checkout — ~/devel/aiter tracks master and has a different C ABI.

$ nm -DC mha_varlen_fwd_bf16_nlogits_nbias_mask_lse_ndropout_skip_nqscale.so \
    | grep -oE '(Simplified)?GenericAttentionMask<[^>]*>' | sort | uniq -c
    120 SimplifiedGenericAttentionMask<true>
$ nm -DC ..._nmask_... | ...
    120 SimplifiedGenericAttentionMask<false>

Zero causal-specialised GenericAttentionMask<true,false> instances, so the single _mask binary serves causal, left-only, right-only and two-sided bands alike; its IsOutOfBound enforces i_x < x_start || i_x >= x_end independently. Also confirmed on 0.1.20: the use_mask rule above, the three-branch if (is_causal) / else if (both == -1) / else mapping whose else emits "b:left,right,sink", and that sink_size still sits between window_size_right and mask_type in all four arg structs, matching our vendored headers.

Test plan

New tests/rocm/test_prefill_sliding_window.py: causal × window_left ∈ {-1, 0, 31, 127} × {fp16, bf16} × {fa2, aiter} across all three entry points, plus the paged wrapper at a native page size (1024) and one that degrades to flat-gather (16). gfx942: 138 passed, 0 skipped, 0 failed.
A/B: reintroducing only the args-site defect (key left alone, JIT cache cleared) fails the non-causal row at 95.5% of elements, max abs diff 1.09; restoring it passes. The log confirms mha_fwd_fp16_nbias_mask_... loaded in both halves, so the defect was isolated to mask_type. The causal rows are unaffected by construction — causal ? X : Y and needs_mask ? X : Y are identical when causal is true.
Full pytest -n auto --reruns 2 -m "not slow" tests/rocm/ on gfx942 from a cleared JIT cache: 29072 passed, 3621 skipped, 0 failed on ROCm 7.15.26333 / torch 2.12.0+rocm10.0.0 / amd-aiter 0.1.20 — 172 more than the pre-change collection, which is this PR's new file. Measured at 34e08ecb0, before the rebase onto refactor(rocm): remove the CUDA/HIP dual-platform abstraction #338.
Both arches re-running at fa691dcea — refactor(rocm): remove the CUDA/HIP dual-platform abstraction #338 relocated the Python layer under flashinfer/rocm/, so the gfx942 number above no longer describes the pushed tree and is not being carried forward as a pass.
pre-commit run -a.

Copilot

Pull request overview

Fixes a correctness defect in the ROCm AITER prefill backend where causal=False sliding-window requests incorrectly routed to the unmasked (_nmask) kernel variant, silently producing unwindowed results. The change aligns FlashInfer’s AITER variant selection and mask arguments with AITER’s “anything masked” split, and adds targeted regression tests plus capability-matrix documentation updates.

Changes:

Route non-causal windowed prefill calls through AITER’s masked (_mask) variant by keying on needs_mask = causal || window_left >= 0, and pass masking args accordingly.
Add ROCm regression tests covering causal × window_left × dtype × backend across single, ragged, and paged prefill entry points.
Update arch capability notes and README support matrix to advertise sliding-window support for AITER prefill.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
`include/flashinfer/rocm/attention/aiter/aiter_loader.h`	Rename `VariantKey::causal` to `needs_mask` and update equality/hash to key variant selection on “any masking”.
`csrc/rocm/aiter_loader.cc`	Switch `_mask/_nmask` segment selection to `needs_mask` and update dlopen-failure hint text accordingly.
`include/flashinfer/rocm/attention/aiter/single_prefill.cuh`	Normalize overly-wide windows, compute `needs_mask`, and ensure masked variant + args are used for non-causal windowed calls.
`include/flashinfer/rocm/attention/aiter/batch_prefill.cuh`	Compute `needs_mask` and ensure both flat-gather and native-paged args select masking when needed.
`flashinfer/prefill_rocm.py`	Add `_aiter_needs_mask()` and thread `needs_mask` through AITER bootstrap/probe paths without changing public signatures.
`tests/rocm/test_prefill_sliding_window.py`	New regression suite validating window semantics (including non-causal windowed) across backends and wrappers.
`flashinfer/arch_caps.py`	Update AITER prefill capability notes to include sliding window.
`README.md`	Update the generated support matrix rows to advertise sliding-window support for AITER prefill.

Suppressed comments (2)

csrc/rocm/aiter_loader.cc:195

The dlopen failure hint prints causal_or_windowed=..., but users must pass the actual AITER keyword argument name (is_causal) when calling aiter.ops.mha.mha_varlen_fwd. Using a non-existent parameter name makes the hint harder to follow.

    return "  Hint: trigger AITER's lazy JIT build by importing aiter.ops.mha and "
           "calling mha_varlen_fwd with matching (dtype=" +
           std::string(key.dtype == VariantKey::Dtype::kFp16 ? "fp16" : "bf16") +
           ", causal_or_windowed=" + (key.needs_mask ? "true" : "false") +
           ", has_lse=" + (key.has_lse ? "true" : "false") + ")." + kAbiPinNote;

csrc/rocm/aiter_loader.cc:207

The dlopen failure hint prints causal_or_windowed=..., but mha_batch_prefill_func() takes the keyword causal. Matching the real AITER argument name makes the hint directly actionable, while the added note clarifies that this value is used to select the _mask variant.

        return "  Hint: trigger AITER's lazy JIT build by calling "
               "aiter.ops.mha.mha_batch_prefill_func() once with matching (dtype=" +
               std::string(key.dtype == VariantKey::Dtype::kFp16 ? "fp16" : "bf16") +
               ", causal_or_windowed=" + (key.needs_mask ? "true" : "false") +
               ", has_lse=" + (key.has_lse ? "true" : "false") +

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

A prefill with causal=False and window_left >= 0 returned the unwindowed result -- 0.66-0.69 abs error against an fp32 windowed reference, on the flat-gather and native-paged routes, on gfx942 and gfx950. Nothing raised and the numbers were plausible. Our variant key drove the .so mask axis off causal alone, so the request loaded a _nmask binary -- built with SimplifiedGenericAttentionMask<false>, which has no masking logic compiled in. AITER splits that .so on "is anything masked", not on causality: no_mask = (not causal) and window_size_left == -1 and window_size_right == -1 use_mask = not no_mask -- aiter/ops/mha.py:766-769, 0.1.20 So VariantKey::causal becomes needs_mask, fed causal || window_left >= 0, and the three args sites set mask_type from it. window_size_right keeps the real causal: 0 blocks future tokens, while -1 saturates to the full extent and leaves the left-bound-only band FlashInfer defines at include/flashinfer/attention/variants.cuh:89. Verified on the installed amd-aiter 0.1.20 wheel, not a source checkout: $ nm -DC mha_varlen_fwd_bf16_nlogits_nbias_mask_lse_ndropout_skip_nqscale.so \ | grep -oE '(Simplified)?GenericAttentionMask<[^>]*>' | sort | uniq -c 120 SimplifiedGenericAttentionMask<true> $ nm -DC ..._nmask_... | ... 120 SimplifiedGenericAttentionMask<false> Zero causal-specialised GenericAttentionMask<true,false> instances, so the one _mask binary serves causal, left-only, right-only and two-sided bands; its IsOutOfBound enforces i_x < x_start || i_x >= x_end independently. Adding a separate has_window field would have been dead weight: build_so_name only ever ORs it with causal, so it could never select a different .so while breaking every existing probe caller. Reusing the field keeps the Python signatures untouched -- the same OR goes into the existing parameter at all eight call sites, including the four else-branch bootstraps that an explicit backend="aiter" takes, which would otherwise miss the dlopen. The single-prefill shim also normalizes window_left >= kv_len to -1, as AITER's own torch path does (mha_fwd_kernels.cu:253), so a window wider than the sequence stays on _nmask and off a cold ~35-minute CK build. The batch shims cannot: plan() knows only a maximum kv_len.

The combination that broke was untested everywhere in the tree: upstream's prefill window tests all pass causal=True, and tests/rocm/test_sliding_window.py is decode-only. That is why a silently-unwindowed result shipped. Covers causal x window_left {-1, 0, 31, 127} x {fp16, bf16} x {fa2, aiter} across all three entry points, and the paged wrapper at both a native page size (1024, mha_batch_prefill) and one that degrades to flat-gather (16) -- different args sites in batch_prefill.cuh. test_aiter_and_fa2_agree_on_a_non_causal_window first asserts the windowed and unwindowed references differ by more than 0.1, so the test cannot pass by both backends ignoring the window -- which is exactly the failure it guards. gfx942: 138 passed, 0 skipped, 0 failed.

batch_decode already claimed it; both prefill rows now honour a window for causal and non-causal alike, so the capability notes should say so. README.md is regenerated from arch_caps.py by scripts/gen_arch_support_matrix.py.

Two of these are coverage defects that only show up on the box you did not run on: - The batch tests gated AITER on the arch string alone. is_aiter_supported() does not consult the capability table, so on gfx950 + ROCm 7.2.x -- where batch_prefill/aiter is known_bad -- the wrapper constructor raises ArchCapabilityError and 48 cases ERROR instead of skipping. _skip_unless_aiter now takes the op and asks capability_available(). - test_paged_prefill_window assumed page_size=1024 takes the native-paged kernel. If _aiter_native_paging_available() degrades, both page sizes take flat-gather, batch_prefill.cuh's native args site goes untested, and the suite stays green while a comment claims otherwise. Asserted directly. _aiter_needs_mask's kv_len loses its default. Omitting it at the single-prefill site is not a safe over-approximation: Python would bootstrap _mask while the shim dlopens _nmask, and AITER pre-ships zero mha_fwd*.so, so that is a load failure rather than a wasted build. The batch sites pass None explicitly. The dlopen hint named causal_or_windowed=, which is not an AITER keyword; it now reports is_causal= and notes that a window selects the same variant. Remaining stale "causal" comments in prefill.py and backends.md updated. Rebased onto #338, which moved the Python layer to flashinfer/rocm/; the test imports move with it.

demandal25 · 2026-09-01T00:57:14Z

Suppressed comments, review 5072018324 (2) — both are the causal_or_windowed= hint, fixed in fa691dcea.

The second is right further than it states: mha_batch_prefill_func() takes causal, while mha_fwd/mha_varlen_fwd take is_causal. fa691dcea used is_causal= for all three hints, so the batch-prefill one is still not the real keyword. Correcting it in the next push — both arch suites are currently reading this worktree and editing aiter_loader.cc mid-run would give mixed-version results.

Note fa691dcea is a rebase onto #338, which moved the Python layer under flashinfer/rocm/; the gfx942 full-suite number in the description was measured at 34e08ecb0 and is not carried forward.

Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

The batch-prefill hint said is_causal=, but that entry point takes causal; only mha_fwd and mha_varlen_fwd use is_causal. Measured against the installed 0.1.20 wheel rather than the review comment: mha_fwd aiter/ops/mha.py:205 is_causal: bool, mha_varlen_fwd aiter/ops/mha.py:929 is_causal: bool, mha_batch_prefill_func aiter/ops/mha.py:3881 causal=False, A hint that names a parameter the function does not accept is worse than no hint, since it sends someone debugging a missing .so down a TypeError.

Copilot

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Suppressed comments (3)

Previously missed (1) — in code that hasn't changed since the last review.

csrc/rocm/aiter_loader.cc:185

The dlopen failure hint suggests passing has_lse=... to aiter.ops.mha.mha_fwd, but the AITER Python API uses return_softmax_lse (see flashinfer/rocm/prefill.py where the bootstrap calls mha_fwd(..., return_softmax_lse=...)). Using a non-existent kwarg makes the hint harder to follow.

This issue also appears in the following locations of the same file:

line 195
line 210

           ", is_causal=" + (key.needs_mask ? "true" : "false") +
           " (window_size_left>=0 selects the same variant)" +
           ", has_lse=" + (key.has_lse ? "true" : "false") + ")." + kAbiPinNote;

csrc/rocm/aiter_loader.cc:197

The dlopen failure hint suggests passing has_lse=... to aiter.ops.mha.mha_varlen_fwd, but the AITER API uses return_softmax_lse. The current hint kwarg name is likely invalid for users trying to bootstrap the variant.

           ", is_causal=" + (key.needs_mask ? "true" : "false") +
           " (window_size_left>=0 selects the same variant)" +
           ", has_lse=" + (key.has_lse ? "true" : "false") + ")." + kAbiPinNote;

csrc/rocm/aiter_loader.cc:212

The dlopen failure hint for mha_batch_prefill_func prints has_lse=..., but the callable uses return_lse (see flashinfer/rocm/prefill.py bootstrap). Using the real kwarg name in the hint would reduce confusion when users try to trigger the lazy JIT build.

               ", causal=" + (key.needs_mask ? "true" : "false") +
               " (a window selects the same variant)" +
               ", has_lse=" + (key.has_lse ? "true" : "false") +

demandal25 · 2026-09-01T01:33:22Z

Suppressed comments, review 5073008867 (3) — all valid, fixed in #342 after this merged.

has_lse is not a parameter of any of the three entry points. Read off the installed 0.1.20 wheel: mha_fwd and mha_varlen_fwd take return_softmax_lse (mha.py:209, :933), mha_batch_prefill_func takes return_lse (:3885). #342 also fixes dtype=, which the review did not reach and is not a keyword either — it comes from the q/k/v tensors.

My earlier disposition comment covered only review 5072018324; reviews 5072933371/5072952604/5073008867 landed after later pushes and I did not re-sweep before merge.

## Summary The three dlopen-failure hints in `aiter_loader.cc` tell the reader to call an AITER entry point "with matching `(dtype=..., is_causal=..., has_lse=...)`". `has_lse` is not a parameter of any of them, so someone debugging a missing `.so` follows the hint and gets a `TypeError` rather than the JIT build they wanted. Raised as suppressed comments on #339 after it merged. Verified against the installed amd-aiter 0.1.20 wheel rather than the review text: ``` mha_fwd aiter/ops/mha.py:209 return_softmax_lse: bool, mha_varlen_fwd aiter/ops/mha.py:933 return_softmax_lse: bool, mha_batch_prefill_func aiter/ops/mha.py:3885 return_lse=False, ``` `dtype` is the same defect one step further, which the review did not reach: it is not a keyword either — it comes from the q/k/v tensors. Relabelled `q dtype` so everything inside the parentheses is a real keyword and the next reader is not misled by a mixed list. This is the second round of the same class on this file; #339 corrected `causal` vs `is_causal`. The cause is that the hint reads as prose while being a literal call recipe. ## What changed - **`csrc/rocm/aiter_loader.cc`** — `has_lse=` becomes `return_softmax_lse=` for the `mha_fwd` and `mha_varlen_fwd` hints and `return_lse=` for `mha_batch_prefill_func`; `dtype=` becomes `q dtype=` in all three. Error-message text only, no behaviour change. ## Test plan - [x] Each keyword read off the installed 0.1.20 wheel in the ROCm 10.0 container, not from a source checkout. - [x] `pre-commit run -a`. - No test covers these strings — they are the text of a `std::runtime_error` raised only when a variant `.so` is absent.

Copilot AI lite review requested due to automatic review settings August 31, 2026 22:28

Copilot started reviewing on behalf of demandal25 August 31, 2026 22:29 View session

Copilot AI reviewed Aug 31, 2026

View reviewed changes

Comment thread csrc/rocm/aiter_loader.cc Outdated

demandal25 added 4 commits August 31, 2026 20:32

demandal25 force-pushed the aiter-prefill-window-left branch from 34e08ec to fa691dc Compare September 1, 2026 00:56

Copilot AI review requested due to automatic review settings September 1, 2026 00:56

Copilot started reviewing on behalf of demandal25 September 1, 2026 00:57 View session

Merge branch 'amd-integration' into aiter-prefill-window-left

e112a92

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 1, 2026 01:00

Copilot started reviewing on behalf of demandal25 September 1, 2026 01:01 View session

Copilot AI reviewed Sep 1, 2026

View reviewed changes

Copilot AI review requested due to automatic review settings September 1, 2026 01:10

Copilot started reviewing on behalf of demandal25 September 1, 2026 01:11 View session

demandal25 merged commit 7614b2e into amd-integration Sep 1, 2026
3 checks passed

demandal25 deleted the aiter-prefill-window-left branch September 1, 2026 01:11

Copilot AI reviewed Sep 1, 2026

View reviewed changes

demandal25 mentioned this pull request Sep 1, 2026

Name the real AITER kwargs in the dlopen hints #342

Merged

2 tasks

Navigation Menu

Uh oh!

Route non-causal windowed prefill to AITER's masked kernel - #339

Route non-causal windowed prefill to AITER's masked kernel#339
demandal25 merged 6 commits into
amd-integrationfrom
aiter-prefill-window-left

demandal25 commented Aug 31, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

demandal25 commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

demandal25 commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

demandal25 commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

The defect

Design: reuse the key field rather than add one

Why window_size_right stays -1

What changed

Verification against the installed 0.1.20 wheel

Test plan

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

demandal25 commented Sep 1, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

demandal25 commented Sep 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

demandal25 commented Aug 31, 2026 •

edited

Loading

Why `window_size_right` stays `-1`