Fix Attention is_causal bottom-right alignment for external KV cache (onnx#8068, #28904)

[titaiwangms]

Summary

Aligns the opset-24 ONNX-domain Attention kernels (CPU + CUDA) with the ONNX errata onnx/onnx#8068 (tracking RFC onnx/onnx#8054) for the external static KV-cache path — keyed by nonpad_kv_seqlen (input #6), no past_key.

Addresses #28904.

What changed

1. Bottom-right is_causal alignment. Per batch, offset[b] = nonpad_kv_seqlen[b] − q_sequence_length; a query at in-block index i attends key j iff j <= i + offset[b]. Applied on CPU and on the CUDA Flash / Memory-Efficient (MEA) / unfused paths. The MEA causal-alignment selector is now offset-aware (no unconditional top-left when an external cache is present); Flash's native bottom-right + per-batch seqlens_k is used where eligible.
2. Fully-masked-row → 0 guard (Bug-2). A query row with no allowed key now outputs a zero row instead of mean-of-V (the finite-sentinel softmax result). Detected with an exact per-key structural predicate (isneginf-equivalent) and zeroed with select (not multiply) before P @ V, so 0 @ V = 0. Added on CPU and the CUDA MEA path. The Flash is_causal + seqlens_k path (offset >= 0) cannot produce a fully-masked row and is intentionally left unguarded. Bool-mask conversion was already select-not-multiply on both EPs (Bug-1 satisfied; no change needed).
3. Reject removal. Removes the CUDA NOT_IMPLEMENTED reject for is_causal + nonpad_kv_seqlen with S_q != total_kv and no past_key — the spec now defines this result, so the op computes it rather than rejecting.

Full-prefill (offset = 0) and past_key decode paths remain bit-identical. Contrib MultiHeadAttention / GroupQueryAttention consume the shared FMHA kernels and are unchanged — only the ONNX-domain Attention dispatch is retargeted.

Test coverage

• C++ AttentionTest gtests: 73/73 pass, including new bottom-right-offset, structural-empty causal row → 0 (CPU + CUDA), and fp16 fully-masked-row goldens.
• Python test_onnx_attention: 277/0 — includes the updated test_tensorscatter_attention.py (stale negative-reject → positive bottom-right acceptance).
• QA final gate: from-scratch Debug build green.

Preemptive onnx#8068 node-test skips (de-skip TODO)

This branch adds the new onnx/onnx#8068 Attention backend node tests to both skip lists so they don't fail before the onnx dependency is bumped:

• onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc
• onnxruntime/test/onnx/TestCase.cc (C++ GetBrokenTests)

These are a no-op on the current onnx pin (v1.21.0). TODO (de-skip): remove both skip lists once cmake/external/onnx is bumped to a release containing onnx#8068.

Deferred follow-up

q_seq > 1 Python bottom-right parity coverage requires upgrading the test_onnx_attention suite's numpy/torch reference functions from total-kv-relative causal (offset = kv_seq − q_seq) to nonpad-relative bottom-right (offset = nonpad_kv_seqlen − q_seq); a naive is_causal=1 flip on the current refs is a no-op or a false failure against the correct kernel. The q_seq > 1 / nonpad < q_seq behavior (including structural-empty rows) is already locked by the C++ gtest goldens. Tracked as follow-up.

References

• Fix Attention is_causal causal-mask alignment + composed is_causal/attn_mask NaN robustness for external (static) KV cache (#8054) onnx/onnx#8068 — spec + reference errata (bottom-right is_causal on the nonpad_kv_seqlen/no-past_key path + composed is_causal + attn_mask NaN robustness). Separately pushed, CI green, awaiting SIG review.
• Attention op: support offset-aware causal masking for KV-cache decode and chunked/mid-cache prefill (nonpad_kv_seqlen) onnx/onnx#8054 — RFC: offset-aware causal masking for KV-cache decode / chunked prefill.

[github-actions]

You can commit the suggested changes from lintrunner.

[titaiwangms]

Review summary — bottom-right is_causal + fully-masked-row→0 (onnx#8068)

Reviewed by a multi-model team (readability, code, correctness, adversarial, deep spec-adherence, cross-module integration).

Verdict: Core math is sound and spec-verified. The bottom-right frontier (j ≤ i + offset, offset = nonpad_kv_seqlen[b] − q_seq) is exact on CPU and CUDA; full-prefill/past-decode paths are bit-identical; the Flash path being left unguarded is provably safe (empty frontier → 0 × 1 = 0, not NaN). Two issues worth fixing before merge.

🔴 Major

1. ZeroFullyMaskedRowsKernel silently drops rows beyond the grid cap — attention_mask_impl.cu
One row per thread, no grid-stride loop, but grid_size is capped at 65535. If batch·heads·q_seq > 65535 × threads, tail rows are never zeroed → MEA mean-of-V / NaN leaks with no error.
Fix: grid-stride loop over row with int64_t indexing (matches the sibling kernels in this file).

2. CPU qk_matmul_output_mode=3 (kPostSoftMax) snapshot is wrongly zeroed — core/providers/cpu/llm/attention.cc
The errata exempts the mode-3 debug output from the fully-masked guard (locked upstream by test_attention_23_fullymasked_qk_matmul_output_mode3_nan), but the code zeroes output before the out_qk memcpy, so the snapshot captures guarded (zeroed) probabilities instead of raw softmax. Y is correct; the debug tensor isn't.
Fix: take the kPostSoftMax memcpy before the apply_fully_masked_guard zeroing loop. (CPU-only observable — CUDA doesn't support mode 3.)

🟡 Minor

• Float-mask predicate isn't isneginf-exact, and CPU/CUDA diverge. CPU uses bias == sentinel, CUDA uses bias > masked_bias_value; for a genuine -inf or finite < -1e30 float mask they classify differently. Ranked Minor because every onnx#8068 fully-masked example uses boolean masks and the predicate is exact for the boolean/structural paths — but worth a comment that it's sentinel-equality (not isneginf), and ideally aligning the two EP comparisons.
• Comment formula contradicts code at cpu/llm/attention.cc (~line 544): comment says the offset is "clamped to ≥ 0", but the code clamps the frontier index (max(s+offset+1, 0)). Code is correct; comment misleads.
• Negative-offset zeroing over-attributes to onnx#8068. The spec calls offset < 0 "out of contract"; the comments cite [Android] Move add header files into AAR to using Gradle #8068 as mandating the zero. Reword to "defensive handling of the out-of-contract negative-offset case".
• Preemptive skip-lists (TestCase.cc, onnx_backend_test_series_filters.jsonc) for the very tests this PR implements risk durable drift once the onnx pin bumps. Keep the de-skip TODO tightly tracked.
• Perf (shared kernel): UnfusedSoftmaxKernel runs the s_any_retained BlockReduce unconditionally, paying sync overhead on the contrib GQA/MHA hot path where mask_sentinel_active is always false. Gate the reduction behind mask_sentinel_active (the seqlens_k==0 boundary is already caught by the s_max==-inf bailout).

⚪ Nits

Naming (masked_bias_value using 0.0f as "disabled", thread_retained as float-encoded bool); stale comment in cuda/llm/attention.cc (~line 884, still says mean-of-V); redundant 4-line comment after the MaskedBiasSentinel extraction; undefined "M-1" tag in test comments; CUDA tests copy-paste goldens while comments say "reuses".

📌 Follow-up (out of PR scope)

.agents/skills/cuda-attention-kernel-patterns/SKILL.md still asserts the removed behavior ("is_causal=1 with TensorScatter decode is spec-invalid"). Worth updating so future tasks don't re-reject this path.

---

Praise: clean first-principles bottom-right derivation (CPU/CUDA independently match), analytic goldens that pin the frontier set, the fp16 -40000 vs -65504 regression test nailing the old heuristic's boundary, and Flash safety asserted-and-verified rather than assumed.

[titaiwangms]

Multi-agent review synthesis — re-verified against the onnx#8068 reference

Updated: I re-ran the review team against the actual onnx#8068 reference (onnx/reference/ops/op_attention.py) and its 15 new backend test families, then reconciled the skip-lists against the PR diff. This supersedes the earlier version of this comment, which had two inaccuracies (noted under "Corrections").

Verdict: The core spec alignment is correct. The three Major items are real but, after grounding in the reference, two are latent (not exercised by any onnx#8068 conformance test) and the third has a narrow hard-failure set (3 tests on CUDA).

Confirmed correct against op_attention.py (solid)

• Bottom-right offset & inclusive frontier: offset = nonpad - q_len (no clamp), j <= i + offset — CPU/CUDA/MEA all reproduce op_attention.py:155,162 exactly. The clamp is on the frontier index, not the offset (correct).
• Fully-masked → Y=0 via select-not-multiply (survives poisoned V), and mode-3 = the guarded softmax with guard-then-snapshot ordering — matches op_attention.py:265-275.
• causal_from_top_left=false is scoped to the nonpad branch only → past_key/cross-attention paths byte-unchanged ("no behavior change" verified).
• fp16 sentinel correctly avoids the lowest()×kLog2e → -inf → s_prime=0 → NaN CUTLASS overflow.

Major (re-characterized after spec grounding)

M1 · Float-mask fully-masked predicate diverges from isneginf — CONFIRMED, but latent/untested
The reference detects a masked key with np.isneginf(np.max(attn_bias)) (op_attention.py:265), and a native float attn_mask is added as-is. ORT uses a finite-sentinel proxy instead:

• CPU (cpu/llm/attention.cc:599): masked ⇔ == lowest() → a user float -inf fully-masked row is missed → softmax over all--inf → NaN (spec wants 0).
• CUDA (attention_mask_impl.cu:200, sentinel max(lowest, -1e30)): masked ⇔ bias <= -1e30 → catches -inf, but a finite fp32/bf16 bias in (-3.4e38, -1e30] is wrongly zeroed.

So the "bit-for-bit identical CPU/CUDA / matches isneginf" claim (attention.cc:583-585) is inaccurate for native float biases. However, no onnx#8068 test exercises this — every fully-masked test uses a bool or pure-structural mask; the only float--inf tests (softcap_neginf_mask[_poison]) mask keys 4–5 only, leaving keys 0–3, so no row is fully masked and the guard never fires. Latent gap, worth a follow-up, not a conformance failure.

M2 · Flash path has no fully-masked-row guard — CONFIRMED, but latent/untested
LaunchZeroFullyMaskedRows is invoked only in the MEA/unfused paths; RunFlashAttention (attention.cc:244-451) has none. For nonpad < S_q (negative offset) the structurally-empty early rows must be Y=0. The onnx negative_offset_structural_empty test is fp32 → routes to MEA (Flash requires non-fp32), so it passes. The Flash gap only bites a user fp16/bf16 model with S_q>1, nonpad < S_q, no mask, Flash-eligible head_size — and is covered by no onnx test and no ORT gtest (the gtest uses head_size=4 → MEA). Whether FA2's lse=-inf rows emit 0 vs NaN is unverified. Suggest a Flash test with nonpad < S_q or an explicit guard/proof.

M3 · Preemptive skip-list is incomplete — CONFIRMED (8 new-behavior families unskipped; 3 hard-fail)
Of the 14 genuinely-new #8068 families, 5 are skipped. The 9 unskipped include causal_with_past_and_present (a regression guard on the unchanged past_key path → passes), leaving 8 new-behavior families unskipped:

• Will hard-fail on CUDA (qk_matmul_output_mode==3 → NOT_IMPLEMENTED, see Support softcap, softmax_precision in Attention(CUDA) #27712): test_attention_23_fullymasked_qk_matmul_output_mode3_zero, test_attention_24_fullymasked_qk_matmul_output_mode3_zero, test_attention_24_qk_matmul_output_mode3_softmax_precision. ← the real pin-bump risk. (CPU passes these — it snapshots kPostSoftMax after the guard.)
• Expected to pass, but no skip and no local gtest mirror: test_attention_4d_softcap_neginf_mask, ..._softcap_neginf_mask_poison, ..._causal_nonpad_negative_offset_structural_empty, ..._causal_nonpad_attn_mask_composition, ..._causal_nonpad_batch_prefill.

Suggest adding a ^test_attention_(23|24)_.*qk_matmul_output_mode3 skip (guarded by #27712) to both TestCase.cc and the jsonc until CUDA mode-3 lands; and adding local gtest mirrors for batch>1 prefill, composition, and the Flash structural-empty case.

Additional (out of this PR's scope, recorded for completeness)

attention_helper.h:141-147 enforces attn_mask.shape[-1] == total_sequence_length, but the spec (op_attention.py:120-128) allows a shorter attn_mask padded with -inf/False. The corresponding test (..._diff_heads_mask4d_padded_kv) is a pre-existing #7867 skip, so this is a known prior limitation, not introduced here.

Minor / nits

• attention.cc:583-585 comment: "bit-for-bit identical / matches isneginf" is true only for bool + structural masking; reword to scope out native float biases (M1 follow-up).
• attention.cc:1243 masked_bias_value = ... : 0.0f overloads 0.0f as a "disabled" marker on a field also used as a real threshold; consider a named sentinel / std::optional<float>.
• TestCase.cc:791 TODO(#28904) bundles an already-satisfied condition (kernels land in this PR); trim to just the onnx-pin bump.

Corrections to the earlier version of this comment

1. M1 previously claimed test_attention_4d_softcap_neginf_mask exercises the CPU float--inf gap — incorrect: that test is a partial mask (no fully-masked row), so the guard never fires.
2. M3 previously listed 6 missing families — the precise count is 8 new-behavior families (the earlier list omitted attention_24_fullymasked_qk_matmul_output_mode3_zero and attention_24_qk_matmul_output_mode3_softmax_precision), of which only the 3 mode-3 families hard-fail (CUDA NOT_IMPLEMENTED).

Note

The "contrib MHA/GQA unchanged" claim is confirmed at the code level (GQA leaves masked_bias_value=0 → guard inert; MHA uses a separate softmax path); build wiring/template instantiations are complete; no other EP implements opset-24 ONNX Attention, so no silent divergence.

_{🤖 Synthesized from a multi-agent review (code, correctness/critical, spec/deep, integration), grounded in the onnx#8068 reference and reconciled against the PR diff before posting.}

[tianleiwu]

Thanks for tightening the bottom-right/nonpad behavior and adding CUDA coverage across Flash, MEA, and unfused. I found one CPU/CUDA consistency issue in the fully-masked-row detector that should be fixed before this lands.

[titaiwangms]

Thanks @tianleiwu for the careful review — all three points were helpful. Addressed below;
the first two are fixed in the latest push, and the third is a clarification (no code change,
though I'm happy to add a comment).

1. CPU float--inf mask predicate (cpu/llm/attention.cc)

Good catch, and thank you for flagging the priority — you're right that the CPU/CUDA
divergence on -inf masking is a pre-merge correctness item, not a follow-up. Fixed.

The CPU fully-masked classifier previously keyed only on the internal finite sentinel
(mask == mask_filter_value<T>()), so a user-supplied IEEE -inf in a float mask was not
recognized as masked — a fully--inf row then went through softmax and produced NaN,
diverging from the CUDA guard. We've widened the CPU predicate to treat a row position as
masked when it is either the internal sentinel or isneginf(mask), while leaving
finite user biases (e.g. -40000) unmasked. This matches the ONNX reference, which
classifies masked rows via isneginf (op_attention.py), and closes the -inf gap you
identified so CPU and CUDA now agree on both the internal sentinel and IEEE -inf. A
fully-masked row now yields 0 on CPU — for both Y and qk_matmul_output — instead of
NaN, consistent with the reference.

(One scoping note: CUDA caps its masked-bias threshold at -1e30
(kCutlassSafeMaskFilterValue), so it also treats finite biases in (lowest(), -1e30] as
masked, whereas CPU follows the reference's isneginf and treats those as live. That
pre-existing finite-bias difference is separate from this fix; this change targets the -inf
case you raised and aligns CPU with the reference.)

We've added a CPU test covering a fully-masked row driven by a user float -inf, plus an
fp16 -inf twin. A finite large-negative bias (e.g. -40000) staying unmasked is covered by
the existing Attention_FiniteNegativeBias_RowNotZeroed_Fp16 test.

2. Unfused block-wide "any key retained" flag (unfused_attention.cu)

Agreed — using a float BlockReduce (with its shared-memory cost) just to compute a boolean
"any key retained" was wasteful. Switched to __syncthreads_or() over the per-thread
predicate, which gives the same block-wide OR with no shared memory and no cub reduction. The
change is correctness-neutral (it computes the identical flag) and the gtests pass. Thanks for
the concrete suggestion.

3. causal_from_top_left = false on the nonpad branch (cuda/llm/attention.cc)

This one is intentional and spec-correct per onnx#8068 — not a decoding regression. I traced
the decode case to make sure, and I think the concern comes from reading the line in
isolation; the surrounding refactor splits what it controls.

The change splits the old unconditional causal_from_top_left = (past_sequence_length == 0)
into two MEA branches:

• External / static-cache branch (the nonpad_kv_seqlen path) → false. On this branch
  there is no past_key (the two are mutually exclusive, enforced at validation), so the old
  (past_sequence_length == 0) reading resolved to top-left — and that was precisely the
  bug: for an external-cache decode (S_q == 1) top-left makes the query attend only key 0.
  false selects bottom-right alignment with offset = nonpad_kv_seqlen − q_sequence_length,
  which is exactly the bottom-right mandate in onnx#8068 (op_attention.py offset semantics).
• Internal-cache / no-cache branch → still (past_sequence_length == 0), unchanged. The
  ordinary decode path (past_sequence_length > 0 → bottom-right) and the no-cache /
  cross-attention path (past_sequence_length == 0 → top-left) keep their previous behavior.
  So there is no change to the internal-cache decoding case you flagged.

In other words, only the external-cache path moved (top-left → bottom-right, which fixes the
non-functional decode), while internal-cache decode is untouched. This is pinned by the
continued-prefill and structural-empty gtests, which assert the bottom-right offsets against
the reference.

Happy to add a short code comment on the two branches making the split explicit if you think
it would help future readers.

Thanks again for the thorough pass.

[tianleiwu]

See below

Posted in error to the wrong PR (this review text belongs to a different change about CUDA reduction integer accumulators). Dismissing; the correct review for #28958 follows.

[tianleiwu]

Re-reviewed the latest head. Thanks for addressing the three items from the previous round:

• CPU float -inf predicate — the std::isinf(mask_value_f) && mask_value_f < 0 branch now classifies a user -inf as masked alongside the finite sentinel; the fp32/fp16 NegInfFullyMasked gtests pin it. Resolved.
• Unfused fully-masked detection — switched to __syncthreads_or(thread_retained), which both publishes s_max and reduces the retained predicate with no extra shared memory / cub reduce. Resolved.
• causal_from_top_left decode change — confirmed it is scoped to the nonpad_kv_seqlen != nullptr branch only; the no-cache / past_key MEA branch keeps (past_sequence_length == 0), so decode/prefill on those paths is byte-unchanged and the external-cache path correctly uses bottom-right. Resolved.

Core bottom-right math (offset = nonpad_kv_seqlen[b] - q_seq, j <= i + offset, frontier clamped at 0) and the select-not-multiply zeroing look correct and are well covered on CPU / MEA / unfused.

Two remaining items (non-blocking) I'd like tracked before the onnx pin bump — both are latent on the current v1.21.0 pin but become live once the #8068 tests run:

1. Flash path has no fully-masked-row guard. LaunchZeroFullyMaskedRows runs only in the MEA and unfused paths. The eligibility change now routes is_causal + nonpad_kv_seqlen (no past_key) to Flash, and for S_q > 1 with nonpad_kv_seqlen[b] < S_q the negative offset leaves early query rows structurally empty. Decode (S_q == 1) is safe (offset >= 0), and the fp32 / head_size == 4 tests route to MEA — so nothing currently exercises Flash here. A user fp16/bf16 model with a Flash-eligible head_size and nonpad < S_q would depend on FA2 emitting 0 (not NaN) for an lse == -inf row, which is unverified. Suggest either an explicit post-Flash guard or a gtest with a Flash-eligible head_size and nonpad < S_q to lock the behavior.

2. Preemptive skip-list omits the mode-3 tests that hard-fail on CUDA. The purpose of the preemptive skip block is to keep CI green at the onnx pin bump, but the qk_matmul_output_mode3 families are deferred to a TODO. CUDA returns NOT_IMPLEMENTED for modes beyond kQK, so test_attention_(23|24)_*qk_matmul_output_mode3* will hard-fail the moment the pin includes #8068. Adding the ^test_attention_(23|24)_.*qk_matmul_output_mode3 skip now (guarded by the #27712 / #28994 reference) to both this jsonc and TestCase.cc would make the block self-consistent and avoid a known future break.

[titaiwangms]

Thanks @tianleiwu — both items addressed, and verified on an A100.

1. Flash fully-masked-row guard. Rather than add a redundant host-side guard, I added regression gtests that drive a structurally-empty (negative-offset, nonpad_kv_seqlen < S_q) leading-row case through the Flash path and assert the leading rows are exactly 0, not NaN:

• …FlashStructuralEmptyRows_Zero_{FP16,BF16}_CUDA — head_size=64, q_seq=4, nonpad=2, is_causal, no mask/past → routes to FA2 (confirmed via the dispatch log); locks the single-tile epilogue (softmax.h:182-186, sum==0 → acc_o=0).
• …FlashStructuralEmptyRows_Zero_SplitKV_FP16_CUDA — total_seq=257 (> block_n=256) forces num_splits=2, so it also locks the split-KV combine (flash_fwd_kernel.h:1125/1135/1153, lse_logsum=+inf → expf(-inf)=0) — the higher-risk path. num_splits=2 holds independently of SM count for these params (verified against num_splits_heuristic in flash_api.cc).

All gate SM>=8.0. On A100 the leading rows are 0 (no NaN) on both epilogues, so FA2's empty-row behavior is now locked by tests; no post-Flash guard is needed.

2. Preemptive skip-list for mode-3. Promoted the qk_matmul_output_mode3 skips from TODO to active, matching the sibling #8068 families:

• onnx_backend_test_series_filters.jsonc: active regex ^test_attention_(23|24)_.*qk_matmul_output_mode3 (global; #28994 / attention.cc:1477 / #27712 refs kept).
• TestCase.cc: the 6 exact names scoped to the provider == "cuda" block, so CPU mode-3 conformance (which CPU implements, =0 / T1) stays covered; #28994/#27712 refs kept.

CI stays green at the onnx pin bump and the block is self-consistent.

[tianleiwu]

Re-reviewed at cc34d0b9. Both of my previously open threads are addressed, and I've resolved them:

• Flash empty-row coverage (cuda/llm/attention.cc:1395): my ask for an explicit guard or a gtest with a Flash-eligible head_size and nonpad < S_q is satisfied by the three new Flash regression tests — Attention_Causal_NonPadKVSeqLen_FlashStructuralEmptyRows_Zero_{FP16,BF16}_CUDA (single-tile epilogue) and ..._SplitKV_FP16_CUDA (cross-split combine). They lock that FA2 emits a zero row (not NaN) for an lse == -inf frontier. Resolved.
• Mode-3 skip self-consistency (onnx_backend_test_series_filters.jsonc): the ^test_attention_(23|24)_.*qk_matmul_output_mode3 regex and the matching TestCase.cc CUDA-scoped entries were added now (not deferred), so the onnx pin bump can't break CUDA CI. Resolved.

The core fix looks correct and is exceptionally well-tested:

• int32_t causal_diagonal_offset (signed) plus the static_cast<int32_t>(query_start) hardening in kernel_forward.h correctly preserves the negative num_keys - num_queries offset; the head_size=512 MEA tests and the mixed-batch {-2, +2} offset test exercise the per-batch RightPaddingBatchHook path directly.
• The fully-masked-row → 0 guard uses an exact structural predicate (sentinel-equality / IEEE -inf) on both EPs rather than the prior magnitude heuristic; the FiniteNegativeBias fp16 regression tests pin that a finite -40000 bias is no longer wrongly zeroed.
• __syncthreads_or for the per-row retained flag (replacing the float reduce) is correct — it doubles as the s_max publish barrier.

Non-blocking suggestion (no need to hold the PR): the new CUDA MEA fully-masked-row tests all drive the nonpad_kv_seqlen (seqlens_k != nullptr) branch of LaunchZeroFullyMaskedRows. The standard MEA path call site (attn_bias != nullptr, seqlens_k == nullptr — e.g. an all-false bool mask on an MEA-eligible shape with total_sequence_length % 4 == 0) doesn't appear to have a dedicated CUDA test; Attention4DAttnMaskBoolAllFalse is CPU/DML-only. A small fp16 all-masked test on that shape would close the seqlens_k == nullptr kernel branch. Approving.

[titaiwangms]

Update: CI fix + regression-test migration + docs

Why the previous CI run was red (and why it wasn't the offset fix): the int32 causal_diagonal_offset fix itself is validated on CI — the original FlashStructuralEmptyRows_* tests and the head_size=64 MEA negative-offset tests pass on the sm86 (A10) QUICK_BUILD legs. The red came entirely from the new head_size=512 MEA regression tests I had added, which fail at kernel launch on sm86/sm89 with CUDA failure 1: invalid argument. Root cause is a separate, pre-existing latent gap: for head_size > 256 the ONNX Attention MEA path is selected (eligibility has no shared-memory check), but the CUTLASS FMHA kernel's dynamic shared-memory request exceeds the sm86/89 opt-in cap (~99 KB), and the cudaFuncSetAttribute(MaxDynamicSharedMemorySize) return is not checked, so the over-budget launch fails with no fallback. It only works where the smem fits (sm90/H100, ~227 KB). This is independent of this PR's fix and is tracked in #28388 (reopened with a repro; its prior fix PR #28383 was never merged).

Test migration: the 5 head_size=512 MEA regression tests are replaced with head_size=64 + ORT_DISABLE_FLASH_ATTENTION=1 force-MEA equivalents (negative / zero / positive / mixed-batch offsets). This preserves the negative-offset MEA coverage — the wrap guard is head_size-independent, since the offset = nonpad_kv_seqlen − q_seq is seqlen-only — while being portable across all architectures (head_size=64 needs only a few KB of static smem). head_size=512 was removed as arch-fragile, with a do-not-reintroduce note pointing at #28388. Verified green under both QUICK_BUILD and full Flash builds (AttentionTest.* 83/83), with all migrated tests verbose-confirmed MEA-routed.

Note: forcing MEA now relies on the env-var mechanism rather than head_size=512's by-construction routing; the tests verbose-confirm MEA was actually selected, and under QUICK_BUILD head_size=64 routes to MEA by construction (Flash is hdim128-only there).

Docs: this PR also bundles the knowledge-capture skill/doc updates from debugging this issue (attention dispatch + CUTLASS FMHA rebuild gotchas, plus a general build-flag / false-green / signed-offset lesson set). Happy to split these into a separate docs-only PR if reviewers would prefer to keep this PR scoped to the kernel fix + tests — just say the word.

[titaiwangms]

Update: CI fix + regression-test migration + docs

Why the previous CI run was red (and why it wasn't the offset fix): the int32 causal_diagonal_offset fix itself is validated on CI — the original FlashStructuralEmptyRows_* tests and the head_size=64 MEA negative-offset tests pass on the sm86 (A10) QUICK_BUILD legs. The red came entirely from the new head_size=512 MEA regression tests I had added, which fail at kernel launch on sm86/sm89 with CUDA failure 1: invalid argument. Root cause is a separate, pre-existing latent gap: for head_size > 256 the ONNX Attention MEA path is selected (eligibility has no shared-memory check), but the CUTLASS FMHA kernel's dynamic shared-memory request exceeds the sm86/89 opt-in cap (~99 KB), and the cudaFuncSetAttribute(MaxDynamicSharedMemorySize) return is not checked, so the over-budget launch fails with no fallback. It only works where the smem fits (sm90/H100, ~227 KB). This is independent of this PR's fix and is tracked in #28388 (reopened with a repro; its prior fix PR #28383 was never merged).

Test migration: the 5 head_size=512 MEA regression tests are replaced with head_size=64 + ORT_DISABLE_FLASH_ATTENTION=1 force-MEA equivalents (negative / zero / positive / mixed-batch offsets). This preserves the negative-offset MEA coverage — the wrap guard is head_size-independent, since the offset = nonpad_kv_seqlen − q_seq is seqlen-only — while being portable across all architectures (head_size=64 needs only a few KB of static smem). head_size=512 was removed as arch-fragile, with a do-not-reintroduce note pointing at #28388. Verified green under both QUICK_BUILD and full Flash builds (AttentionTest.* 83/83), with all migrated tests verbose-confirmed MEA-routed.

Note: forcing MEA now relies on the env-var mechanism rather than head_size=512's by-construction routing; the tests verbose-confirm MEA was actually selected, and under QUICK_BUILD head_size=64 routes to MEA by construction (Flash is hdim128-only there).

Docs: this PR also bundles the knowledge-capture skill/doc updates from debugging this issue (attention dispatch + CUTLASS FMHA rebuild gotchas, plus a general build-flag / false-green / signed-offset lesson set). Happy to split these into a separate docs-only PR if reviewers would prefer to keep this PR scoped to the kernel fix + tests — just say the word.

@tianleiwu For couple of new commits: I fixed some tests that were accidentally using head_size=512 to test MEA kernel which does not officially support 512. Also, docs updated for ai-agents and human reading. As well as the resolve conflicts with main branch. Details are in above message!