diff --git a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md b/.agents/skills/cuda-attention-kernel-patterns/SKILL.md index 5c074a50fc294..6d916a73e8095 100644 --- a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md +++ b/.agents/skills/cuda-attention-kernel-patterns/SKILL.md @@ -34,8 +34,47 @@ Unified Unfused → RunUnfusedAttention() (handles both MHA and GQA via reshape-Q trick) ``` +### Eligibility anchors (symbols are stable; line numbers as of cc34d0b914) + +| Stage | Decision symbol | Hard caps | Dispatch gate | +|---|---|---|---| +| Flash | `flash::is_supported` (`flash_api.cc:414`) | fp16/bf16 only, SM≥8.0, `head_size%8==0`, **`head_size<=256`** | `attention.cc:1385` `flash_eligible` (fp32 excluded at :1387) | +| MEA | `has_memory_efficient_attention` (`memory_efficient_attention.h:68`) | `(head_size&7)==0` **and `head_size<=kEfficientAttentionMaxHeadSize` (1024)**; **NO shared-memory feasibility check** (see #28388 + the head_size=512 caveat below) | `attention.cc:1415` `mea_eligible`; bias-stride `%4` at :1436 | +| Unfused | (none — catch-all) | all dtypes/shapes | `attention.cc:1485` `RunUnfusedAttention` | + +**`head_size=512` IS routed to MEA, but its MEA kernel is NOT portably launchable — so it +is not a robust test probe.** +By the predicate, 512 > 256 fails Flash and 512 ≤ 1024 with `512 & 7 == 0` passes the MEA +predicate, **so dispatch selects MEA** — but the MEA eligibility check +(`memory_efficient_attention.h:68-73`) gates only on SM + +`head&7==0` + `head<=1024`, with **no shared-memory check**. For `head_size=512` FP16 the +CUTLASS MEA `SharedStorage` exceeds the dynamic-smem opt-in cap on capacity-limited arches +(sm86 ~99KB, sm80 ~163KB, sm90 ~227KB — **non-monotonic**, no clean SM-version guard). +`fmha_launch_template.h` calls `cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, ...)` +but **ignores its return value and launches anyway**, so on sm86 the kernel dies at launch +with `CUDA failure 1: invalid argument` — there is **no fallback to unfused** (live bug +#28388; its fix PR #28383 was never merged). So `head_size=512`'s MEA kernel launches only +on large-smem arches like sm90/H100. + +**To force the MEA path portably in a test, use `ORT_DISABLE_FLASH_ATTENTION=1` with a small +`head_size` (e.g. 64) whose SharedStorage fits every target arch — NOT `head_size=512`.** +Also guard with `SKIP_IF_MEA_NOT_COMPILED` (see §7) so a MEA-OFF build SKIPs rather than +false-greens via the (correct) unfused fallback. + **Flash eligibility**: fp16/bf16 only, SM≥8.0 (Ampere+), `head_size == v_head_size`, `head_size <= 256`, no `output_qk`, `attn_mask == nullptr`. Uses `mha_fwd` / `mha_fwd_kvcache`. +> **QUICK_BUILD caveat (false hypothesis trap).** *(General principle — build flags can +> silently reroute kernel dispatch — lives in the `ort-build` skill, "Agent tips". The +> attention-specific instance:)* With `onnxruntime_QUICK_BUILD=ON` +> (`-DORT_QUICK_BUILD`), Flash is compiled for **head_dim 128 only**: +> `flash_api.h:147` `is_supported` returns false for `head_size != 128`, and +> `static_switch.h:80` `HEADDIM_SWITCH` only instantiates `kHeadDim=128`. So under +> QUICK_BUILD nearly every shape routes to **MEA**, not FlashAttention-2. If a +> `head_size!=128` test "fails only on some SM", suspect **MEA** (CUTLASS, +> arch-independent), NOT a Flash/FA2 hardware bug. `head_size=512` **is routed to MEA** in all +> MEA-enabled builds (Flash caps at 256), but its MEA kernel **fails to launch** on +> small-smem GPUs — see the `head_size=512` caveat above (#28388). + **MEA eligibility**: SM50+/53+/80+ by dtype, `head_size <= 1024` and divisible by 8 (enforced by `has_memory_efficient_attention`), no `output_qk`. GQA additionally requires `head_size == v_head_size` (for `LaunchUngroup`); decode also requires it (for `LaunchConcatNewToPastKV`). Bias stride must satisfy `total_sequence_length % 4 == 0`. GQA with FP32 is excluded (LaunchUngroup only has fp16/bf16 instantiations). Supports `softcap + attn_mask` — CUTLASS applies softcap before bias in kernel tiles, matching ONNX spec ordering (onnx/onnx#7867, supersedes the now-closed onnx/onnx#7865 issue). **Unified Unfused Attention**: Always available as the final fallback. Handles both MHA (`num_heads == kv_num_heads`, group=1) and GQA (`num_heads != kv_num_heads`, group>1) via a reshape-Q trick with stride-based cuBLAS batched GEMM (no K/V head replication). Uses FP32 QK scratch for precision. Supports all features: @@ -142,15 +181,45 @@ for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; Always call `CUDA_CALL(cudaGetLastError())` after kernel launches in standalone helper functions. This is the established pattern in the file (see `ConcatPastToPresent`, `PastPresentBufferShare`). -## 6. Fully-Masked Batches - -All-false bool masks or `seqlens_k=0` produce NaN in CUTLASS MEA. - -**Additive-bias path** (bool mask converted to bias): Fixed by capping `mask_filter_value` to `-1e+30f` (see section 2). CUTLASS then naturally computes uniform softmax → mean(V). - -**Nonpad path** (`seqlens_k=0`): CUTLASS skips all K/V positions → `s_prime=0` → NaN. Fixed by `ZeroOutputForFullyMaskedBatches` kernel which zeros output for batches where `seqlens_k[b] == 0`. Note: this produces zeros, not mean(V) — a cross-EP consistency TODO exists. - -**CPU/Unfused behavior**: `mask_filter_value = lowest()` (not `-inf`). All masked values are equal → `softmax(equal) = 1/N` → output = mean(V). This is the spec reference. +## 6. Fully-Masked Rows and Batches + +All-false bool masks, an all-`-inf` `attn_mask` row, or a causal/nonpad frontier +with no allowed key produce NaN in CUTLASS MEA (the uniform/empty softmax degenerates: +`s_prime=0` → `1/s_prime=inf` → `0 × inf = NaN`). Per onnx/onnx#8068 (Bug-2), a +**fully-masked query row** — one with no key allowed by the composed causal + nonpad + +mask constraints — must output a **zero row** (`Y = 0`), **not** mean-of-V. + +**This `Y = 0` behavior is now consistent on BOTH EPs** (the earlier mean(V)-vs-zero +cross-EP divergence is RESOLVED — there is no longer an open TODO here): + +- **CUDA**: `ZeroFullyMaskedRowsKernel` (in `attention_mask_impl.cu`) runs after the + MEA/CUTLASS output and zeros each fully-masked row with a **select** (not multiply, + so `0 @ V = 0` even when V is poisoned). It detects a fully-masked row with an exact + per-key predicate (within the causal/nonpad frontier AND the additive-bias slot is + above the mask sentinel), matching the onnx#8068 `isneginf`-of-row-max reference. A + finite (even very negative) user bias is not the sentinel, so its key stays unmasked + and the row is left untouched. +- **CPU**: `core/providers/cpu/llm/attention.cc` applies the same Bug-2 guard — after + softmax it zeros any row whose composed frontier admitted no unmasked key. + +**Additive-bias path** (bool mask converted to bias): `mask_filter_value` is capped to +`-1e+30f` (see section 2) so CUTLASS does not overflow to NaN; a row that is nonetheless +fully masked is then zeroed by the per-row guard above. + +**Whole-batch empty (`seqlens_k[b] == 0`)**: the structural case where an entire batch +has zero valid keys is additionally handled by `ZeroOutputForFullyMaskedBatches`, which +zeros that batch's output. (The per-row guard covers the finer-grained case where only +some query rows are fully masked.) + +**`qk_matmul_output_mode` (mode 3 / post-softmax debug output)**: for a fully-masked row +the mode-3 snapshot is **mandated to be `0`** (zero row), consistent with `Y = 0`, per the +onnx#8068 SIG decision (this superseded the earlier "unspecified" proposal). The CPU +post-softmax snapshot is taken **after** the row-zeroing guard — matching the onnx +reference and the v23/v24 function bodies, where the guard runs before the mode-3 +capture — so the debug tensor reflects the same zero row as the output. Note this +mode-3=0 behavior is served by the **CPU** path: CUDA `qk_matmul_output_mode` beyond +`kNone`/`kQK` (i.e. `kPostSoftCap`/`kPostMaskBias`/`kPostSoftMax`) returns +`NOT_IMPLEMENTED` (`attention.cc`), so an agent must not assume CUDA produces mode-3=0. ## 7. Test Runner Targeting @@ -177,13 +246,21 @@ ScopedEnvironmentVariables scoped_env({ Enable verbose logging to confirm: `LOGS_DEFAULT(VERBOSE) << "ONNX Attention: using ..."`. +**`SKIP_IF_MEA_NOT_COMPILED`** is a local gtest macro (defined in +`test/providers/cpu/llm/attention_op_test.cc`) that `GTEST_SKIP`s — rather than silently +passes — when `USE_MEMORY_EFFICIENT_ATTENTION` is OFF, so an MEA-targeted test cannot +false-green via the (correct) unfused fallback. Use it in any test that must prove the MEA +path ran (see the `ort-test` skill → "Verify which path/kernel actually executed"). + ## 8. Cross-EP Consistency CPU is the spec reference implementation. CUDA outputs should match CPU for all valid inputs. - CPU uses `mask_filter_value = std::numeric_limits::lowest()` (finite, not `-inf`) - CPU softmax: subtract-max-first → works correctly with extreme finite values -- CPU handles fully-masked batches naturally (uniform softmax → mean(V)) +- CPU zeros fully-masked query rows (onnx#8068 Bug-2 guard) — output `Y = 0`, matching + CUDA's `ZeroFullyMaskedRowsKernel`. (Earlier docs claimed CPU produced mean(V) here; + that divergence is resolved — both EPs now emit a zero row.) Run tests with `disable_cpu=false` to always validate against CPU. The C++ test framework (`RunTest4D`) supports `disable_cpu`, `disable_cuda`, `disable_dml` flags. @@ -198,7 +275,7 @@ Run tests with `disable_cpu=false` to always validate against CPU. The C++ test | `core/providers/cuda/llm/attention_mask_impl.h` | Declarations for ONNX mask/bias kernels | | `core/providers/cpu/llm/attention.cc` | CPU reference implementation (ONNX domain) | | `core/providers/cpu/llm/attention_helper.h` | ONNX parameter validation and shape computation | -| `test/providers/cpu/llm/attention_op_test.cc` | C++ attention tests (all EPs) | +| `test/providers/cpu/llm/attention_op_test.cc` | C++ tests for the **ONNX-domain** `Attention` op — suite `AttentionTest.*`, runs in `onnxruntime_provider_test` (all EPs). NOT to be confused with the contrib `test/contrib_ops/attention_op_test.cc` (`ContribOpAttentionTest.*`); see `ort-test` skill. | | `test/python/transformers/test_onnx_attention/test_mha.py` | Python parity tests | | `test/python/transformers/test_onnx_attention/common.py` | Python test utilities and reference `attention_ref()` | @@ -230,38 +307,108 @@ The contrib `QkvToContext` function (used by contrib MHA, NOT by ONNX Attention) The ONNX spec defines two causal alignment modes based on where query positions sit in the full attention matrix: -- **Upper-left**: `q_i` attends to `kv[0..i]`. Query positions start at 0 in the full matrix. -- **Lower-right**: `q_i` attends to `kv[kv_len - q_len + i..kv_len - 1]`. Query positions are at the end. +- **Upper-left** (a.k.a. *top-left*): `q_i` attends to `kv[0..i]`. Query positions start at 0 in the full matrix. +- **Bottom-right** (a.k.a. *lower-right*): `q_i` attends to `kv[0 .. kv_len - q_len + i]` — i.e. keys `j` with `j <= i + offset`, where `offset = kv_len - q_len` (clamped `>= 0`). The causal diagonal is anchored at the end of the key axis. This is the term onnx/onnx#8068 uses; kernel flags spell it `CausalFromBottomRight`. -**ONNX spec rule**: `is_causal=1` always means upper-left in the full matrix. When `past_key` provides context, `past_sequence_length` shifts the query start position forward — the resulting `[S_q × total_kv]` sub-matrix effectively has lower-right alignment. +**ONNX spec rule**: causal alignment depends on how the KV context is supplied. + +- **Internal cache / no cache** (`past_key`, or plain self-attention): `is_causal=1` + is upper-left in the full matrix. When `past_key` provides context, + `past_sequence_length` shifts the query start position forward — the resulting + `[S_q × total_kv]` sub-matrix is effectively bottom-right. +- **External / static cache** (`nonpad_kv_seqlen`, no `past_key`, opset 24): per + onnx/onnx#8068, `is_causal=1` uses **bottom-right** (offset-aware) alignment — + query in-block index `i` attends key `j` iff `j <= i + offset[b]`, where + `offset[b] = nonpad_kv_seqlen[b] - q_sequence_length` (clamped to `>= 0`). ### Per-kernel behavior | Kernel | Alignment | Mechanism | |--------|-----------|-----------| -| **Flash** | Lower-right only | `is_causal` flag → `seqlen_k - seqlen_q` offset in kernel. No top-left option. | +| **Flash** | Bottom-right only | `is_causal` flag → `seqlen_k - seqlen_q` offset in kernel. No upper-left option. | | **MEA (CUTLASS)** | Both | `causal_from_top_left` flag in `MemoryEfficientAttentionParams`. `true` → `CausalFromTopLeft` (offset=0). `false` → `CausalFromBottomRight` (offset = num_keys - num_queries). | -| **Unfused** | Both | `past_kv_length` param. `0` → upper-left. `total_kv - S_q` → lower-right. | +| **Unfused** | Both | `past_kv_length` param. `0` → upper-left. `total_kv - S_q` → bottom-right. | ### Dispatch logic in attention.cc ```cpp -// Flash cannot do upper-left → guarded by causal_cross_no_past +// Pure cross-attention with NO external cache (S_q != S_kv, no past, no nonpad): +// this is the upper-left case Flash cannot express. bool causal_cross_no_past = parameters.is_causal && parameters.q_sequence_length != parameters.total_sequence_length && parameters.past_sequence_length == 0; -// Flash: skip when causal_cross_no_past (no top-left support) -// MEA: NOT skipped — handles it via causal_from_top_left = (past_sequence_length == 0) -// Unfused: always correct via past_kv_length = parameters.past_sequence_length +// Flash: eligible UNLESS (causal_cross_no_past && nonpad_kv_seqlen == nullptr). +// - No external cache -> upper-left required -> skip Flash (no upper-left support). +// - External cache (nonpad_kv_seqlen != nullptr) -> required frontier IS bottom-right +// (onnx#8068), so Flash IS eligible and produces it natively via seqlens_k. +// MEA: external cache -> causal_from_top_left = false (bottom-right, offset = num_keys - +// num_queries == nonpad_kv_seqlen[b] - q_len per batch); otherwise causal_from_top_left +// = (past_sequence_length == 0). +// Unfused: always correct via past_kv_length (0 -> upper-left; total_kv - S_q -> bottom-right). ``` ### When S_q == S_kv -Upper-left and lower-right produce **identical** results when `S_q == S_kv` (the offset is 0 either way). The alignment distinction only matters for cross-attention shapes (`S_q != S_kv`). +Upper-left and bottom-right produce **identical** results when `S_q == S_kv` (the offset is 0 either way). The alignment distinction only matters for cross-attention shapes (`S_q != S_kv`). ### TensorScatter decode (opset 24 external KV cache) -TensorScatter manages KV cache externally — `past_key` is nullptr but K/V already contain the full sequence. Per the ONNX spec, `is_causal` with `S_q != S_kv` and no `past_key` means upper-left (q[0] sees only kv[0]), which is **not meaningful for decode**. - -**Correct pattern**: TensorScatter decode must use `is_causal=0` and rely on `nonpad_kv_seqlen` to bound the active KV range. Models using `is_causal=1` with TensorScatter decode have a spec-invalid combination. +TensorScatter manages KV cache externally — `past_key` is nullptr but K/V already +contain the full sequence, with `nonpad_kv_seqlen[b]` giving each batch's valid +(non-padded) key count. Per onnx/onnx#8068, `is_causal=1` with an external/static KV +cache (no `past_key`) uses **bottom-right** (offset-aware) alignment: query in-block +index `i` attends key `j` iff `j <= i + offset[b]`, where +`offset[b] = nonpad_kv_seqlen[b] - q_sequence_length` (clamped to `>= 0`). For decode +(`q_sequence_length == 1`) the single query row therefore attends all +`nonpad_kv_seqlen[b]` valid keys — the meaningful, spec-correct result (not the +degenerate "q[0] sees only kv[0]" of upper-left). + +**Correct pattern**: `is_causal=1` with TensorScatter + `nonpad_kv_seqlen` (no `past_key`) +is **valid and supported** for both decode and continued-prefill — it yields bottom-right +causal attention bounded by the per-batch valid-key count. (`is_causal=0` is also valid +where a model wants no causal masking.) The earlier `is_causal=1` NOT_IMPLEMENTED reject +was **removed** in the onnx#8068 alignment work; the only still-invalid combination is +`nonpad_kv_seqlen` together with `past_key` (mutually exclusive internal-vs-external +cache, enforced at validation in `attention_helper.h`). + +## 12. Signed Offsets in CUTLASS FMHA (uint wrap hazard) + +This is a specific instance of the general **signed-vs-unsigned wrap** bug class — see +`AGENTS.md` → "Signed vs unsigned on negative-capable differences" for the principle. Below +are the **attention-specific** fix sites in `cutlass_fmha/kernel_forward.h`. See §11 for what +the offset *means* (bottom-right alignment); this section is purely the signed-arithmetic +hazard. + +Any FMHA offset computed as a difference of counts — canonically +`causal_diagonal_offset = num_keys - num_queries` (`CausalFromBottomRight`) — is +**negative** whenever `num_keys < num_queries` (cross-attention / KV-trimmed / +`nonpad_kv_seqlen[b] < q_len`, onnx#8068 / ORT #28904). It **must** be stored and +compared as `int32_t`; a `uint32_t` wraps the negative value to ~4.29e9 (`0xFFFFFFFE`), +the causal-mask guard `min(iter_key_start + kKeysPerBlock, num_keys) >= query_start + offset` +becomes permanently false, the per-element causal mask is **silently skipped**, and boundary +query rows over-attend one extra key. + +### Fix sites in `cutlass_fmha/kernel_forward.h` (symbols are stable; lines as of cc34d0b914) + +| Symbol / guard | Line | What it must do | +|---|---|---| +| `int32_t causal_diagonal_offset` (field decl) | ~206 | Stay **`int32_t`** so the negative offset is preserved (rationale comment ~202-205). | +| `causal_diagonal_offset = num_keys - num_queries;` | ~354 | Set point for `CausalFromBottomRight`; may be negative (comment ~353). | +| `int32_t(query_start + causal_diagonal_offset + kQueriesPerBlock)` | ~366 | First (AttentionKernel) `num_keys` clamp. The inner sum **does** wrap to `0xFFFFFFFE`-style values in unsigned arithmetic when the offset is negative, but casting the **whole sum** to `int32_t` recovers the correct value by two's-complement modular arithmetic, and the result is consumed **arithmetically** (as a `fast_min` operand), so the wrap is harmless. Contrast the ~924 guard, where the value feeds a **relational** comparison — there the unsigned wrap flips the comparison result, so the operand `query_start` must be cast to `int32_t` **before** the compare. | +| "Mask out last if causal" guard: `static_cast(query_start) + p.causal_diagonal_offset` | ~924-926 | `query_start` is `uint32_t` (~707) — cast it to `int32_t` so the comparison is signed (rationale ~919-923). | +| Sliding-window guard ("L957"): `static_cast(query_start) + p.causal_diagonal_offset ...` | ~962-963 | Same cast hardening (rationale ~956-961). | + +### Rules when editing `kernel_forward.h` (or any FMHA kernel) + +- Keep `causal_diagonal_offset` **`int32_t`**. +- `query_start` in the iteration kernels is **`uint32_t`** — `static_cast(query_start)` + before adding the offset in ANY relational guard. +- The same hazard is **dormant but real** at the `window_size > 0` guard: harden it the + same way even though opset-24 Attention currently pins `window=-1` (a future + sliding-window / KV-trim caller could combine `window_size>0` with a negative offset). +- Tests that exercise this need a **negative** offset: `num_keys < num_queries`. Force the + MEA path **portably** with `ORT_DISABLE_FLASH_ATTENTION=1` + a small `head_size` (e.g. 64) + — **not** `head_size=512`, whose MEA launch is arch-fragile on small-smem GPUs (#28388, + see §1). The regression tests live in `test/providers/cpu/llm/attention_op_test.cc` + (`Attention_Causal_NonPadKVSeqLen_MEA_*`), guarded by `SKIP_IF_MEA_NOT_COMPILED`. diff --git a/.agents/skills/cuda-cutlass-fmha-incremental-rebuild/SKILL.md b/.agents/skills/cuda-cutlass-fmha-incremental-rebuild/SKILL.md new file mode 100644 index 0000000000000..ba1c97662e7ef --- /dev/null +++ b/.agents/skills/cuda-cutlass-fmha-incremental-rebuild/SKILL.md @@ -0,0 +1,111 @@ +--- +name: cuda-cutlass-fmha-incremental-rebuild +description: > + Use when rebuilding ONNX Runtime CUDA after editing CUTLASS fused-MHA headers + (onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/*.h such as kernel_forward.h or + fmha_launch_template.h), or when a header edit "passed" an incremental build but + test behavior did not change. Explains the nvcc depfile gotcha that produces stale + Memory-Efficient-Attention (MEA) kernels and binaries, and how to force a correct + recompile. Also covers disk-space frugality on shared GPU dev boxes. +--- + +# Incremental rebuilds silently use STALE CUTLASS fused-MHA kernels + +> The **general** false-green principles (stale binary, wrong-artifact mtime) are summarised +> in the `ort-test` skill's "False-green taxonomy". This skill is the CUDA/CUTLASS-specific +> detail. + +## The gotcha (verification-integrity bug) + +`nvcc`-generated depfiles do **not** track the CUTLASS fused-MHA headers under +`onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/` (e.g. `kernel_forward.h`, +`fmha_launch_template.h`). These headers are `#include`d by the `fmha_sm*.cu` +translation units, but the build system does not record that dependency. + +Consequence: after you edit one of those headers, an **incremental** `build.sh`: + +- does **not** recompile `fmha_sm*.cu`, +- reports `[100%] Built target ...` and exits 0, +- leaves the recompiled artifacts — the `fmha_sm*.cu.o` objects and the + `libonnxruntime_providers_cuda.so` they link into — **unchanged** (same `mtime` as + the pre-edit build). + +(Do **not** use the gtest test-exe mtime as the stale symptom: in the shared-provider +build the exe `dlopen`s the `.so` and is **not** relinked, so its mtime stays old even +after a *correct* rebuild — see "How to confirm" below. The reliable diagnostic signal +is the `fmha_sm*.cu.o` / `.so` mtime.) + +So your "successful" rebuild is running the **old** kernel. Tests that should now +pass (or fail) reflect the previous code, not your edit. This silently invalidates +any FAIL→PASS / PASS→FAIL verification. + +## The fix — force recompile the .cu units + +Before rebuilding after editing any `cutlass_fmha/*.h` header: + +```bash +touch onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/*.cu +``` + +Then run the normal build command. This forces the `fmha_sm*.cu` translation units +(and downstream binaries) to recompile against your header change. + +## How to confirm the rebuild was real (don't trust "[100%] Built") + +Confirm that the artifact which actually **links** the recompiled `fmha_sm*.cu.o` +is newer than your header edit. + +⚠️ **Do NOT just check the test EXE mtime — it can falsely flag a good build as +stale.** In the shared-provider build configuration (the default here), the CUDA +execution provider is a **shared module**: the recompiled `fmha_sm*.cu.o` link into +`libonnxruntime_providers_cuda.so`, and the `onnxruntime_provider_test` executable +**dlopens** that `.so` — it is **not relinked**. So after a *correct* rebuild the +test exe `mtime` stays **old** while the `.so` advances. Checking the exe alone +would wrongly conclude the build was stale. + +Check the right artifact for your link mode: + +- **Shared-provider build (default):** the `.so` that links the recompiled `.o` — + `build///libonnxruntime_providers_cuda.so` +- **Statically-linked provider:** the test exe itself (`onnxruntime_provider_test`) + +Safest check — `stat` both the recompiled object and the `.so`, and confirm BOTH +are newer than the header edit: + +```bash +stat -c '%y %n' onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h +# in your build dir, e.g. build/Debug_quickbuild/Debug/: +stat -c '%y %n' libonnxruntime_providers_cuda.so +# and the actual recompiled object (path varies by build dir): +find . -name 'fmha_sm80.cu.o' -exec stat -c '%y %n' {} + +``` + +If the `.so` (and the `fmha_sm*.cu.o`) timestamps are older than (or equal to) the +header edit, the build was stale — `touch` the `.cu` files and rebuild. The most +reliable signal of all is behavioral: a test that was failing now passes (a stale +binary cannot flip its result). + +## Related: pick the right test binary + +This is the **CUDA/CUTLASS instance of false-green mode 1** (zero-match / wrong binary) — +see the `ort-test` skill's "False-green taxonomy" for the general principle. In short: +attention/MEA/Flash boundary gtests (e.g. `FlashStructuralEmptyRows*`, +`Attention_Causal_NonPadKVSeqLen_MEA_*`) live in **`onnxruntime_provider_test`**, which CI +runs; `onnxruntime_test_all` does not contain them and gives a false green. Verify the +MEA/Flash boundary fix against `onnxruntime_provider_test`. + +## Related: disk frugality on shared GPU dev boxes + +Full ORT CUDA builds are large (test binaries ~1 GB each; a build dir can reach +tens of GB). On a shared box, `/home` filling to 100% makes builds fail in +non-obvious places — e.g. `git submodule sync` reporting `No space left on device` +or a `config.lock` error, not an obvious "disk full" at the compile step. + +Before a big rebuild, check free space and clean only clearly-stale, regenerable +build directories (old dated experiment dirs). Never delete another agent's active +build dir or anything ambiguous: + +```bash +df -h /home +du -sh build/* | sort -h +``` diff --git a/.agents/skills/ort-build/SKILL.md b/.agents/skills/ort-build/SKILL.md index a25d7638ced57..a11e381c583fc 100644 --- a/.agents/skills/ort-build/SKILL.md +++ b/.agents/skills/ort-build/SKILL.md @@ -66,6 +66,7 @@ You do **not** need `--update` when only modifying existing `.cc`/`.h` files — | `--use_cuda` | Enable CUDA EP. Requires `--cuda_home`/`--cudnn_home` or `CUDA_HOME`/`CUDNN_HOME` env vars. On Windows, only `cuda_home`/`CUDA_HOME` is validated. | | `--target T` | Build a specific CMake target (requires `--build`; e.g., `onnxruntime_common`, `onnxruntime_test_all`) | | `--use_webgpu` | Enable WebGPU EP. To run its tests locally on Linux without a GPU, see the `webgpu-local-testing` skill. | +| `--cmake_extra_defines onnxruntime_QUICK_BUILD=ON` | Faster CUDA build: instantiates a reduced kernel set. **Side effect:** Flash is compiled for head_dim 128 only, so most attention shapes fall back to **MEA** (changes which attention kernel is compiled/dispatched). Don't use it to characterize Flash-vs-arch behavior. | | `--build_dir` | Build output directory | ## Build output path @@ -79,6 +80,15 @@ It may be customized with `--build_dir`. ## Agent tips - **Activate a Python virtual environment** before building. See "Python > Virtual environment" in `AGENTS.md`. +- **Build flags can silently reroute which kernel/code path executes.** A build option can + change *which* kernel is compiled, and therefore which code path actually runs — so a CI + failure can live in a different code path than your local build exercises. Before + hypothesizing a hardware- or algorithm-specific cause (e.g. "this GPU arch miscomputes"), + first identify **which kernel actually ran** for the failing configuration (see the + `ort-test` skill → "Verify which path/kernel actually executed"). Concrete instance: + `onnxruntime_QUICK_BUILD=ON` compiles FlashAttention for head_dim 128 only, so most + attention shapes silently dispatch to Memory-Efficient Attention instead of Flash — + details in the `cuda-attention-kernel-patterns` skill. - **Prefer `python tools/ci_build/build.py` directly** over `build.bat`/`build.sh` when redirecting output. The `.bat` wrapper runs in `cmd.exe`, which breaks PowerShell redirection. - **Redirect output to a file** (e.g., `> build_log.txt 2>&1`). Build output is large and will overflow terminal buffers. - **Run builds in the background** — a full build can take tens of minutes to over an hour. Poll the log for `"Build complete"` or errors. diff --git a/.agents/skills/ort-test/SKILL.md b/.agents/skills/ort-test/SKILL.md index 507c5d3b5ae9a..1f78c225b026e 100644 --- a/.agents/skills/ort-test/SKILL.md +++ b/.agents/skills/ort-test/SKILL.md @@ -16,6 +16,20 @@ ONNX Runtime uses **Google Test** for C++ and **unittest** (preferred) / **pytes | `onnxruntime_test_all` | Core framework, graph, optimizer, session tests | | `onnxruntime_provider_test` | Operator/kernel tests (Conv, MatMul, etc.) across execution providers | +### Two `attention_op_test.cc` files — don't confuse them + +There are two same-named files testing **different operators**. Both build into +`onnxruntime_provider_test`: + +| Path | Operator | gtest suite | +|---|---|---| +| `test/providers/cpu/llm/attention_op_test.cc` | **ONNX-domain** `Attention` (opset 23/24) | `AttentionTest.*` | +| `test/contrib_ops/attention_op_test.cc` | **contrib** MultiHeadAttention / GroupQueryAttention | `ContribOpAttentionTest.*` | + +The MEA negative-offset regression tests (`Attention_Causal_NonPadKVSeqLen_MEA_*`, +e.g. `..._MEA_NegOffset_ForceFlashDisabled_FP16_CUDA`) live in the **providers/cpu/llm** file — +the ONNX-domain op. + Use `--gtest_filter` to select specific tests: ```bash @@ -76,7 +90,62 @@ Python test naming convention: `test___[when_ Virtual environment" in `AGENTS.md`. +- **Beware false-green results** — a green run does not always prove anything. See the + "False-green taxonomy" section below for the four ways a test can pass without testing + your change. - **Redirect test output to a file** (e.g., `> test_output.txt 2>&1`) — output can be large. - For C++ tests, verify the build directory exists and a prior build completed before running. - Use `--gtest_filter` to run a targeted subset when the full suite takes too long. - **Running WebGPU tests locally on Linux without a GPU** — WebGPU op tests build into `onnxruntime_provider_test` and can run against a software Vulkan adapter (Mesa lavapipe). See the `webgpu-local-testing` skill. + +## False-green taxonomy — ways a test can "pass" without proving anything + +A green result is not always a real pass. Watch for all five modes: + +1. **Zero-match filter.** A `--gtest_filter` that matches no tests still exits 0 (green). + Confirm the `[==========] N tests ran` line is non-zero — a zero-match run prints + `0 tests from 0 test suites`. Many operator/kernel gtests run only in + **`onnxruntime_provider_test`** (CI runs this), NOT `onnxruntime_test_all`; the wrong + binary matches nothing and looks green. +2. **Stale binary from an incremental build.** If the build did not actually recompile your + change (e.g. a header not tracked by the compiler's depfile), the "passing" run executes + the OLD code. A test that was failing cannot truly flip to passing without a real + rebuild — treat an unexpected FAIL→PASS with suspicion and confirm the linked artifact's + mtime advanced. CUDA/CUTLASS instance (nvcc depfiles don't track `cutlass_fmha/*.h`): see + the `cuda-cutlass-fmha-incremental-rebuild` skill. +3. **Checking the wrong artifact's freshness.** With a dlopen'd shared provider (e.g. + `libonnxruntime_providers_cuda.so`), the test executable is NOT relinked when the provider + recompiles — its mtime stays old while the `.so` advances. Verify the artifact that + actually links your change, not the test exe. Detail: `cuda-cutlass-fmha-incremental-rebuild` + skill. +4. **A correct fallback path masks the intended path.** A value-only assertion can pass via a + *different, correct* code path without ever exercising the one you meant to test (e.g. a + test meant for MEA silently handled by the unfused fallback). Assert/verify **which path + ran**, not just the output value — see "Verify which path/kernel actually executed" below. +5. **Arch-portability false-green (verified on only one GPU arch).** A CUDA kernel that + launches on a large-dynamic-smem arch (e.g. sm90/H100, ~227KB) can **fail to launch** on a + smaller opt-in cap (sm86/89 ~99KB, sm80 ~163KB) with `CUDA failure 1: invalid argument` — + and a path with no fallback (e.g. ORT's MEA) turns that into a hard error, not a silent + degrade. So a green run on your local GPU can mask a launch failure on CI's arch. Verify + arch-portability, or pick a config whose shared-memory footprint fits **every** target arch + (e.g. a small `head_size`). Concrete instance: CUTLASS MEA `head_size=512` FP16 exceeds + sm86's smem opt-in cap and dies at launch — live bug #28388 (the + `cuda-attention-kernel-patterns` skill §1 has the dispatch detail). + +## Verify which path/kernel actually executed + +Value equality alone does not prove the intended code path ran — a correct fallback can +produce the right answer (false-green mode 4 above). When a test targets a specific +kernel/path, confirm it actually dispatched there instead of trusting the output: + +- Enable verbose logging and check the dispatch log line. ORT attention logs one of these + exact strings (`core/providers/cuda/llm/attention.cc`): + - `ONNX Attention: using Flash Attention` (:1400) + - `ONNX Attention: using Memory Efficient Attention` (:1451) + - `Attention: using unified unfused path` (:1482) — note: **no `ONNX ` prefix** and it + reads "unified unfused path", not "Unfused". +- Or force the path via the relevant env var / build config AND add a compile-time guard so + the test **SKIPs** (not silently passes) when the target path is unavailable — e.g. + `SKIP_IF_MEA_NOT_COMPILED`. + +Operator-specific routing/forcing details: `cuda-attention-kernel-patterns` skill §1/§7. diff --git a/AGENTS.md b/AGENTS.md index 18c5a8dd88a1d..44e7257d5d536 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,6 +66,15 @@ Use `reserve()` not `resize()`. Do not use `absl::` directly — use the ORT typ - Prefer `gsl::span` over `const std::vector&` for input parameters - Prefer `std::string_view` by value over `const std::string&` - `SafeInt` (from `core/common/safeint.h`) for memory size arithmetic +- **Signed vs unsigned on negative-capable differences.** Any expression of the form `a - b` + that can be negative (an offset or remaining-budget computed from counts, e.g. + `num_keys - num_queries`) must be stored and compared using a **signed** type + (`int32_t`/`int64_t`), and any unsigned operand must be `static_cast` to signed *before* + the subtraction/comparison. An unsigned result silently wraps to a huge value (`~4.29e9` + for `uint32_t`), which can permanently satisfy or skip a relational guard with **no crash + and no warning** — a correct-looking-but-wrong result. Concrete ORT instance + the exact + fix sites: CUTLASS FMHA `causal_diagonal_offset`, see the `cuda-attention-kernel-patterns` + skill §12. - Don't use `else` after `return` - Avoid `long` (ambiguous width) — use `int64_t` for dimensions, `size_t` for counts - `using namespace` allowed in limited scope but never at global scope in headers diff --git a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h index fc4c6c610cc60..305186235fe55 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h +++ b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h @@ -75,6 +75,8 @@ struct RightPaddingBatchHook { } if (p.custom_mask_type == AttentionKernel::CausalFromBottomRight) { + // May be negative when num_keys < num_queries (nonpad external KV cache, onnx#8068 / ORT #28904). + // causal_diagonal_offset is int32_t so the negative value is preserved (no unsigned wrap). p.causal_diagonal_offset = p.num_keys - p.num_queries; } if (p.custom_mask_type == AttentionKernel::CausalFromTopLeft || diff --git a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h index 88a22fd422af0..5a02ef65933c0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h +++ b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h @@ -199,7 +199,11 @@ struct AttentionKernel { int32_t* seqstart_k_ptr = nullptr; int32_t* seqlen_k_ptr = nullptr; - uint32_t causal_diagonal_offset = 0; + // Signed: CausalFromBottomRight sets this to (num_keys - num_queries), which is NEGATIVE in the + // external-KV-cache regime where nonpad_kv_seqlen < q_sequence_length (onnx#8068 / ORT #28904). + // A uint32_t here wraps the negative value to ~4.29e9 and breaks the causal-mask guard below + // (see the signed comparison at the "Mask out last if causal" block), so it MUST stay int32_t. + int32_t causal_diagonal_offset = 0; // Output tensors output_t* output_ptr = nullptr; // [num_queries, num_heads, head_dim_value] @@ -345,6 +349,8 @@ struct AttentionKernel { // Custom masking if (custom_mask_type == CausalFromBottomRight) { + // May be negative when num_keys < num_queries (nonpad external KV cache, onnx#8068 / ORT #28904). + // causal_diagonal_offset is int32_t so the negative value is preserved (no unsigned wrap). causal_diagonal_offset = num_keys - num_queries; } // We use num_keys_absolute to index into the rng_state @@ -910,9 +916,14 @@ struct AttentionKernel { // first masked element is x = y + offset -> query_start + offset There is // intersection (and we need to mask) if min(iter_key_start + // kKeysPerBlock, num_keys)) >= query_start + offset + // NOTE: query_start is uint32_t and causal_diagonal_offset can be negative + // (CausalFromBottomRight with num_keys < num_queries, onnx#8068 / ORT #28904). Cast query_start + // to int32_t so this comparison is performed in signed arithmetic; otherwise the RHS wraps + // to ~4.29e9, the guard is always false, and the per-element causal mask is skipped + // entirely (boundary query rows then attend one extra key). if (p.custom_mask_type && cutlass::fast_min(iter_key_start + kKeysPerBlock, p.num_keys) >= - (query_start + p.causal_diagonal_offset)) { + (static_cast(query_start) + p.causal_diagonal_offset)) { auto query_start = blockIdx.x * kQueriesPerBlock; auto lane_offset = MM0::AccumLambdaIterator::get_lane_offset( my_lane_id, my_warp_id, iteratorC_tile_offset); @@ -942,8 +953,14 @@ struct AttentionKernel { // query_start + kQueriesPerBlock - window_size mask if x_fist > // x_lowerleft + // Mirrors the signed-comparison hardening of the "Mask out last if causal" guard above + // (onnx#8068 / ORT #28904): query_start is uint32_t, so cast it to int32_t to keep this + // relational comparison in signed arithmetic. Dormant today (window_size > 0 never co-occurs + // with a negative causal_diagonal_offset — opset-24 Attention hard-codes window=-1, GQA/Paged + // keep offset >= 0), but this prevents an unsigned-wrap landmine if a future sliding-window / + // KV-trimming caller combines window_size > 0 with a negative offset. if (p.window_size > 0 && - (query_start + p.causal_diagonal_offset + + (static_cast(query_start) + p.causal_diagonal_offset + cutlass::fast_min( static_cast(kQueriesPerBlock), static_cast(p.num_queries)) - p.window_size >= diff --git a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu b/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu index a0c9d4666cae3..77ffc12d4b379 100644 --- a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu +++ b/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu @@ -99,6 +99,7 @@ __global__ void UnfusedSoftmaxKernel( const int past_kv_length, const float scale, const float softcap, + const float masked_bias_value, T* __restrict__ softmax_out) { // Grid: (N_q * S_q, B, 1). Block: (TPB, 1, 1). const int q_in_head = blockIdx.x % q_sequence_length; @@ -151,8 +152,22 @@ __global__ void UnfusedSoftmaxKernel( __shared__ float s_max; __shared__ float s_inv_sum; - // Pass 1: compute max of masked values. + // A real additive mask sentinel is only in play when masked_bias_value is negative + // (composed is_causal + attn_mask). It is 0 for contrib GQA/MHA which pass no bias. + const bool mask_sentinel_active = has_bias && masked_bias_value < 0.0f; + + // Pass 1: compute max of masked values, and detect whether any in-range key is retained. + // A key is "retained" iff its composed additive bias is strictly greater than the masked + // sentinel (i.e. it was not masked to -inf). When no real mask sentinel is active, every + // in-range key is trivially retained. This exact per-key test matches the MEA + // ZeroFullyMaskedRows kernel and the onnx/onnx#8068 reference (isneginf of the additive-bias + // row max), and replaces the prior fractional score-collapse heuristic that wrongly zeroed + // rows carrying a finite-but-very-negative user bias. The compare is exact because masked + // slots are ASSIGNED the sentinel verbatim (bool-false / causal-overlay write the same + // constant; the sentinel is finite, never NaN) while a legit user bias is passed through + // un-clamped, so only a true masked key trips the (bias_val > masked_bias_value) test. float thread_max = -CUDART_INF_F; + int thread_retained = 0; for (int i = threadIdx.x; i < total_kv_length; i += TPB) { if (i < start || i >= end) continue; float x = qk_in[row_offset + i] * scale; @@ -160,7 +175,11 @@ __global__ void UnfusedSoftmaxKernel( x = softcap * tanhf(x / softcap); } if (has_bias) { - x += ToFloat(attn_bias[bias_row_offset + i]); + const float bias_val = ToFloat(attn_bias[bias_row_offset + i]); + x += bias_val; + if (!mask_sentinel_active || bias_val > masked_bias_value) thread_retained = 1; + } else { + thread_retained = 1; } if (x > thread_max) thread_max = x; } @@ -171,10 +190,19 @@ __global__ void UnfusedSoftmaxKernel( block_max = BlockReduce(tmp_storage).Reduce(thread_max, cub::Max()); #endif if (threadIdx.x == 0) s_max = block_max; - __syncthreads(); + // __syncthreads_or() both publishes s_max (it is a full block barrier) and computes the + // block-wide OR of the per-thread "key retained" predicate in a single pass -- a boolean + // reduction that needs no shared memory and no cub float reduce. The result is identical to the + // prior BlockReduce(max) over a 0/1 float flag, only cheaper (tianleiwu review, onnx#28958). + const int any_retained = __syncthreads_or(thread_retained); // If the row is fully masked, emit zeros (match existing mask-of-zeros behavior). - if (s_max == -CUDART_INF_F) { + // A row is fully masked iff no in-range key is retained -- either: + // (1) causal/window/seqlens leave no in-range key (s_max stays -inf), or + // (2) composed is_causal + attn_mask masked every in-range key to the additive sentinel. + // Both collapse to any_retained == 0. onnx/onnx#8068 (Bug-2) requires a zero row here + // (otherwise softmax would emit a spurious uniform mean-of-V). + if (s_max == -CUDART_INF_F || any_retained == 0) { for (int i = threadIdx.x; i < total_kv_length; i += TPB) { softmax_out[row_offset + i] = T(0.f); } @@ -245,6 +273,7 @@ void LaunchUnfusedSoftmax( params.past_kv_length, params.scale, params.softcap, + params.masked_bias_value, softmax_out); } diff --git a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.h b/onnxruntime/contrib_ops/cuda/bert/unfused_attention.h index 8fb3a18ac7570..86d5be087e8f0 100644 --- a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/unfused_attention.h @@ -77,6 +77,14 @@ struct UnfusedAttentionParams { // Per-batch K lengths (optional). When non-null, positions k >= seqlens_k[b] // are masked out (useful for right-padded packed batches). const int* seqlens_k = nullptr; + + // Composed is_causal + attn_mask fully-masked-row -> 0 guard (onnx/onnx#8068, Bug-2). + // When a query row's every in-range key is masked by a finite additive-bias sentinel, + // the row would otherwise softmax to a uniform mean-of-V instead of zero. Set this to the + // (negative) sentinel value used for masked keys to enable the guard: a row whose softmax + // max stays at/below the sentinel is emitted as zeros. Leave at 0 (default) to disable; + // disabled = legacy mean-of-V behavior (used by contrib GQA, which passes no bias). + float masked_bias_value = 0.0f; }; // Returns required scratch size in bytes. Caller must allocate diff --git a/onnxruntime/core/providers/cpu/llm/attention.cc b/onnxruntime/core/providers/cpu/llm/attention.cc index cd146dd8a35d4..1b0f37feab9fb 100644 --- a/onnxruntime/core/providers/cpu/llm/attention.cc +++ b/onnxruntime/core/providers/cpu/llm/attention.cc @@ -6,6 +6,7 @@ #include "core/providers/cpu/llm/attention_softmax.h" #include "core/common/common.h" +#include "core/common/inlined_containers.h" #include "core/common/safeint.h" #include "core/mlas/inc/mlas.h" #include "core/platform/threadpool.h" @@ -14,6 +15,7 @@ #include "core/providers/cpu/math/gemm.h" #include +#include #include using onnxruntime::attention_helper::AttentionParameters; @@ -346,9 +348,14 @@ void AttentionBase::ComputeAttentionProbs(T* attention_probs, T* mask_data = nullptr; bool delete_mask_data = false; bool causal = parameters.is_causal && parameters.q_sequence_length > 1; + // When nonpad_kv_seqlen is present the causal frontier is offset-aware + // (bottom-right) and per-batch, so it cannot be baked into the batch-shared mask + // buffer here; it is applied per-batch in the main loop below. Skip the top-left + // overlay in that case (see ONNX opset-24 / onnx#8068). + const bool shared_causal_overlay = causal && !parameters.has_nonpad_kv_seqlen; if (mask_index == nullptr) { - // No external mask: allocate only if causal behavior needed. - if (causal) { + // No external mask: allocate only if a batch-shared causal overlay is needed. + if (shared_causal_overlay) { size_t mask_bytes = SafeInt(parameters.q_sequence_length) * parameters.total_sequence_length * sizeof(T); void* raw = allocator->Alloc(mask_bytes); memset(raw, 0, mask_bytes); // start all allowed @@ -362,7 +369,7 @@ void AttentionBase::ComputeAttentionProbs(T* attention_probs, } } else { const bool is_bool_mask = mask_index->IsDataType(); - const bool need_copy = is_bool_mask || causal; // copy if we must convert or overlay causal pattern + const bool need_copy = is_bool_mask || shared_causal_overlay; // copy if we must convert or overlay causal pattern if (need_copy) { size_t mask_bytes = SafeInt(mask_index->Shape().Size()) * sizeof(T); mask_data = static_cast(allocator->Alloc(mask_bytes)); @@ -372,7 +379,7 @@ void AttentionBase::ComputeAttentionProbs(T* attention_probs, } else { make_copy(mask_data, mask_index->Data(), SafeInt(mask_index->Shape().Size())); } - if (causal) { + if (shared_causal_overlay) { // Overlay causal -inf above diagonal for every broadcast slice int slices = mask_batch_size * mask_num_heads; for (int slice = 0; slice < slices; ++slice) { @@ -532,10 +539,28 @@ void AttentionBase::ComputeAttentionProbs(T* attention_probs, // `onnxruntime/core/providers/cpu/llm/attention.h`). The CPU softmax uses this finite sentinel // (not IEEE -inf) because MLAS' softmax kernel expects only finite inputs; the value is small // enough relative to any softcap-saturated score that the corresponding softmax weight is 0. + // + // When is_causal is also set, the causal frontier on this external/static KV-cache path is + // bottom-right (offset-aware) per ONNX opset-24 / onnx#8068: query s attends key t iff + // `t <= s + offset`, where `offset = nonpad_kv_seqlen[b] - q_sequence_length` (clamped to >= 0). + // The first masked key for row s is therefore `min(valid_kv_len, s + offset + 1)`. This folds + // the bottom-right causal mask together with the per-batch valid-key bound. The `past_key` + // decode path (no nonpad_kv_seqlen) and the initial full-prefill (offset == 0) are unaffected: + // they keep using the batch-shared top-left overlay built above and never enter this block. + // Per-row first-masked-key index from the nonpad/causal frontier. Retained so the + // fully-masked-row guard below can reuse the exact same frontier that was just masked. + InlinedVector first_masked_kv_per_row; if (parameters.has_nonpad_kv_seqlen) { - int valid_kv_len = static_cast(parameters.nonpad_kv_seqlen_data[batch_i]); + const int valid_kv_len = static_cast(parameters.nonpad_kv_seqlen_data[batch_i]); + first_masked_kv_per_row.resize(static_cast(parameters.q_sequence_length)); for (int s = 0; s < parameters.q_sequence_length; ++s) { - std::fill(output + s * parameters.total_sequence_length + valid_kv_len, + int first_masked_kv = valid_kv_len; + if (causal) { + const int causal_frontier = s + valid_kv_len - parameters.q_sequence_length + 1; + first_masked_kv = std::min(first_masked_kv, std::max(causal_frontier, 0)); + } + first_masked_kv_per_row[static_cast(s)] = first_masked_kv; + std::fill(output + s * parameters.total_sequence_length + first_masked_kv, output + (s + 1) * parameters.total_sequence_length, mask_filter_value()); } @@ -547,8 +572,71 @@ void AttentionBase::ComputeAttentionProbs(T* attention_probs, memcpy(out_qk, output, SafeInt(probs_matrix_size) * sizeof(T)); } + // Bug-2 (ONNX opset-24 / onnx#8068): a fully-masked query row — one with no key allowed by the + // composed mask/causal/nonpad constraints — must produce a zero output row, not the mean-of-V + // that softmax over an all-sentinel row would otherwise yield (the CPU sentinel + // `mask_filter_value()` is finite, so softmax does not naturally collapse such a row to 0). + // + // The row is classified with an EXACT structural predicate — never a magnitude threshold on the + // combined QK+bias score: a key is "masked" iff its additive mask-bias slot is either the + // internal sentinel (`mask_filter_value()`, a finite large negative) OR an IEEE negative + // infinity supplied directly in a user float `attn_mask`. A row with zero unmasked keys is fully + // masked. A finite (even very negative, e.g. `-40000`) user bias is neither the sentinel nor + // `-inf`, so its key stays unmasked and the row is never zeroed. This matches the onnx#8068 + // reference (`isneginf` of the additive-bias row max), and agrees with the CUDA guard for the + // two cases that matter here — a user `-inf` and the internal sentinel. (It is NOT bit-for-bit + // identical to CUDA for every finite bias: the CUDA guard masks any `bias <= masked_bias_value`, + // whose value is capped at `kCutlassSafeMaskFilterValue` (-1e30), so CUDA also masks finite + // biases in `(lowest(), -1e30]`; this CPU predicate, like the reference `isneginf`, does not.) + const bool apply_fully_masked_guard = (mask_data != nullptr) || parameters.has_nonpad_kv_seqlen; + InlinedVector fully_masked_row; + if (apply_fully_masked_guard) { + fully_masked_row.resize(static_cast(parameters.q_sequence_length), false); + const T* mask_slice = (mask_data != nullptr) ? mask_data + mask_data_offset : nullptr; + const T sentinel = mask_filter_value(); + // Portable conversion to float so the IEEE -inf test works for both float and MLFloat16. + const auto to_float = [](const T v) -> float { + if constexpr (std::is_same::value) { + return v.ToFloat(); + } else { + return static_cast(v); + } + }; + for (int s = 0; s < parameters.q_sequence_length; ++s) { + const int frontier = parameters.has_nonpad_kv_seqlen + ? first_masked_kv_per_row[static_cast(s)] + : parameters.total_sequence_length; + bool has_unmasked_key = false; + for (int t = 0; t < frontier; ++t) { + if (mask_slice == nullptr) { + has_unmasked_key = true; + break; + } + const T mask_value = mask_slice[s * parameters.total_sequence_length + t]; + const float mask_value_f = to_float(mask_value); + const bool key_masked = + (mask_value == sentinel) || (std::isinf(mask_value_f) && mask_value_f < 0.0f); + if (!key_masked) { + has_unmasked_key = true; + break; + } + } + fully_masked_row[static_cast(s)] = !has_unmasked_key; + } + } + ComputeAttentionSoftmaxInplace(output, parameters.q_sequence_length, parameters.total_sequence_length, nullptr, allocator); + if (apply_fully_masked_guard) { + for (int s = 0; s < parameters.q_sequence_length; ++s) { + if (fully_masked_row[static_cast(s)]) { + std::fill(output + s * parameters.total_sequence_length, + output + (s + 1) * parameters.total_sequence_length, + static_cast(0.0f)); + } + } + } + // Snapshot kPostSoftMax (post-softmax). if (out_qk != nullptr && parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kPostSoftMax) { diff --git a/onnxruntime/core/providers/cuda/llm/attention.cc b/onnxruntime/core/providers/cuda/llm/attention.cc index 9e097ea07b880..edafbfb3ede65 100644 --- a/onnxruntime/core/providers/cuda/llm/attention.cc +++ b/onnxruntime/core/providers/cuda/llm/attention.cc @@ -26,6 +26,17 @@ namespace cuda { namespace llm_attention_detail { +// CUTLASS online softmax multiplies attention scores by kLog2e (≈1.4427); for float/bf16, +// |lowest() × kLog2e| > FLT_MAX overflows to -inf and causes s_prime=0 → NaN for fully-masked +// rows. The additive mask sentinel is therefore capped to kCutlassSafeMaskFilterValue. This is +// also the exact value the fully-masked-row guard compares against to decide whether a key was +// masked. See kCutlassSafeMaskFilterValue in memory_efficient_attention.h for details. +template +float MaskedBiasSentinel() { + return std::max(static_cast(std::numeric_limits::lowest()), + ::onnxruntime::contrib::cuda::kCutlassSafeMaskFilterValue); +} + template bool HasOutput(const NodeType& node, size_t output_index) { if constexpr (requires(const NodeType& candidate) { candidate.OutputCount(); candidate.OutputExists(output_index); }) { @@ -158,12 +169,7 @@ Status Attention::ConvertAttnMaskToBias( int64_t num_elements = attn_mask->Shape().Size(); converted_mask_buffer = GetScratchBuffer( num_elements * sizeof(NativeCudaT), GetComputeStream(context)); - // CUTLASS online softmax multiplies attention scores by kLog2e (≈1.4427). - // For float/bf16, |lowest() × kLog2e| > FLT_MAX, overflowing to -inf and - // causing s_prime=0 → NaN for fully-masked batches. Cap to prevent this. - // See kCutlassSafeMaskFilterValue in memory_efficient_attention.h for details. - float mask_filter_value = std::max(static_cast(std::numeric_limits::lowest()), - ::onnxruntime::contrib::cuda::kCutlassSafeMaskFilterValue); + float mask_filter_value = llm_attention_detail::MaskedBiasSentinel(); ORT_RETURN_IF_ERROR(LaunchConvertBoolMaskToAttentionBias( attn_mask->Data(), reinterpret_cast(converted_mask_buffer.get()), @@ -797,15 +803,15 @@ Status Attention::RunMemoryEfficientAttention( p.qk_head_size = parameters.head_size; p.v_head_size = parameters.v_head_size; p.causal = parameters.is_causal; - // ONNX spec: is_causal means upper-left alignment in the full attention matrix. - // When past_sequence_length == 0 and S_q != S_kv (cross-attention without KV cache), - // queries start at absolute position 0, so causal mask is upper-left. - // When past_sequence_length > 0 (decode with KV cache), queries start at position - // past_seq, so causal mask is effectively lower-right on the [S_q x total_kv] sub-matrix. - // NOTE: For external KV cache (TensorScatter), nonpad_kv_seqlen provides per-batch - // actual lengths and seqlens_k handles the masking — the causal_from_top_left flag - // is only consulted when params.causal is true, so it's correct here. - p.causal_from_top_left = (parameters.past_sequence_length == 0); + // External KV cache (nonpad_kv_seqlen) requires BOTTOM-RIGHT causal alignment per + // onnx/onnx#8068: query in-block index i attends key j iff j <= i + offset[b], with + // offset[b] = nonpad_kv_seqlen[b] - q_sequence_length. The CUTLASS kernel derives this + // automatically when causal_from_top_left is false: it sets + // causal_diagonal_offset = num_keys - num_queries, and num_keys is the per-batch + // seqlens_k value (== nonpad_kv_seqlen[b]) — so the offset equals offset[b] per batch. + // Do NOT key this off past_sequence_length (which is 0 for the external-cache path and + // would force the incorrect top-left frontier). + p.causal_from_top_left = false; p.scale = parameters.scale; p.softcap = parameters.softcap; p.seqlen_k_ptr = seqlens_k_buffer.get(); @@ -834,10 +840,6 @@ Status Attention::RunMemoryEfficientAttention( // On the MEA (CUTLASS) path (used for both MHA and GQA when nonpad_kv_seqlen is provided), // zero out output for fully-masked batches to prevent NaN. // CUTLASS epilogue computes 1/s_prime where s_prime=0 for seqlens_k=0, producing NaN. - // TODO(titaiwang): ZeroOutputForFullyMaskedBatches outputs zeros for fully-masked - // batches (seqlens_k=0), which diverges from CPU/Unfused behavior (uniform mean of V). - // For cross-EP consistency, replace with LaunchMeanOfVForFullyMaskedBatches that - // computes mean(V[b,n,:,h]) for each masked batch. See issue #27516. { using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; int64_t elements_per_batch = static_cast(parameters.q_sequence_length) * @@ -850,6 +852,28 @@ Status Attention::RunMemoryEfficientAttention( cuda_stream, device_prop.maxThreadsPerBlock)); } + + // Fully-masked-row -> 0 guard (onnx#8068, Bug-2). A query row is fully masked when no key + // remains in its causal/seqlens frontier with a finite (non-sentinel) composed bias. Two + // causes, both handled by ZeroFullyMaskedRowsKernel: (a) STRUCTURAL — is_causal bottom-right + // with nonpad_kv_seqlen[b] < q_sequence_length leaves early rows with an empty frontier (no + // attn_mask needed); (b) every causally-allowed key is masked by the finite attn_bias + // sentinel. Both would otherwise yield NaN/mean-of-V. Fire whenever is_causal can produce an + // empty frontier OR a real attn_mask is present (matches the reference isneginf(max(attn_bias)) + // which folds the causal -inf into attn_bias). + if (parameters.is_causal || attn_bias_data != nullptr) { + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + float masked_bias_value = llm_attention_detail::MaskedBiasSentinel(); + ORT_RETURN_IF_ERROR(LaunchZeroFullyMaskedRows( + reinterpret_cast(out_data), + reinterpret_cast(attn_bias_data), + seqlens_k_buffer.get(), + parameters.batch_size, parameters.q_num_heads, parameters.q_sequence_length, + parameters.total_sequence_length, parameters.v_head_size, + parameters.is_causal, /*causal_from_top_left=*/false, + broadcast_bias_dim_0, broadcast_bias_dim_1, masked_bias_value, + cuda_stream, device_prop.maxThreadsPerBlock)); + } } // Standard MEA path: float attention bias, bool mask (converted to bias), or no mask. // Bool masks are converted to additive attention bias (true→0, false→mask_filter_value). @@ -901,9 +925,27 @@ Status Attention::RunMemoryEfficientAttention( p.workspace = nullptr; } onnxruntime::contrib::cuda::run_memory_efficient_attention(p); - } - // --- Transpose output BSNH → BNSH if input was 4D (BNSH) --- + // Fully-masked-row -> 0 guard (onnx#8068, Bug-2). No seqlens_k here, so the key bound is + // total_sequence_length and the causal frontier can never be empty (offset is 0 for top-left, + // or num_keys - q_seq >= 0 when past is present), so the only fully-masked cause is an + // attn_mask that masks every key. Gate on attn_bias presence only -- unlike the nonpad branch, + // is_causal alone cannot produce a structurally-empty row here, so firing on it would just be a + // guaranteed no-op launch. + if (attn_bias_data != nullptr) { + using CudaT = typename onnxruntime::cuda::OrtToCudaType::type; + float masked_bias_value = llm_attention_detail::MaskedBiasSentinel(); + ORT_RETURN_IF_ERROR(LaunchZeroFullyMaskedRows( + reinterpret_cast(out_data), + reinterpret_cast(attn_bias_data), + /*seqlens_k=*/nullptr, + parameters.batch_size, parameters.q_num_heads, parameters.q_sequence_length, + parameters.total_sequence_length, parameters.v_head_size, + parameters.is_causal, p.causal_from_top_left, + broadcast_bias_dim_0, broadcast_bias_dim_1, masked_bias_value, + cuda_stream, device_prop.maxThreadsPerBlock)); + } + } if (!is_bsnh && out_bsnh_buffer != nullptr) { ORT_RETURN_IF_ERROR(TransposeBSNHtoBNSH( parameters.batch_size, parameters.q_sequence_length, @@ -1189,6 +1231,12 @@ Status Attention::RunUnfusedAttention( p.scale = parameters.scale; p.softcap = parameters.softcap; p.seqlens_k = seqlens_k_ptr; + // Composed is_causal + attn_mask fully-masked-row -> 0 guard (onnx#8068, Bug-2). The bool + // mask is converted to a finite additive sentinel by ConvertAttnMaskToBias; pass that same + // sentinel so the softmax kernel zeros a row whose every in-range key is masked instead of + // producing mean-of-V. No mask -> 0 (disabled); a float -inf mask is handled by the kernel's + // -inf branch and does not depend on this value. + p.masked_bias_value = (attn_mask != nullptr) ? llm_attention_detail::MaskedBiasSentinel() : 0.0f; NativeCudaT* output_qk_data = (output_qk != nullptr) ? reinterpret_cast(output_qk->MutableData()) @@ -1307,37 +1355,25 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { // acceptable per the ONNX spec. // Flash Attention uses lower-right (bottom-right) causal alignment with no option for - // upper-left. The ONNX spec requires upper-left alignment when there is no past context: - // query[0] attends only to key[0]. The difference only manifests when S_q != S_kv - // (cross-attention shape) with no past. Skip Flash for this case; MEA handles it correctly - // via the causal_from_top_left flag, and Unified Unfused uses past_kv_length=0. - // Defined here for visibility — only Flash needs this guard (MEA/Unfused handle upper-left natively). + // upper-left. The ONNX spec requires upper-left alignment when there is no past context + // AND no external cache: query[0] attends only to key[0]. The difference only manifests + // when S_q != S_kv (cross-attention shape) with no past. Skip Flash for that pure + // cross-attention case; MEA handles it via the causal_from_top_left flag and Unified + // Unfused uses past_kv_length=0. (When an external cache is present — nonpad_kv_seqlen — + // the required frontier IS bottom-right, so Flash is eligible; see below.) const bool causal_cross_no_past = parameters.is_causal && parameters.q_sequence_length != parameters.total_sequence_length && parameters.past_sequence_length == 0; - // Reject causal + TensorScatter decode (S_q < S_kv without past_key). - // Per ONNX spec, is_causal without past_key means upper-left alignment: q[i] attends - // only to kv[0..i]. For decode with external cache (S_q=1, S_kv=cache_size), this means - // q[0] sees only kv[0] — not meaningful for autoregressive generation. - // - // Why is_causal=0 is correct for external cache decode: - // - With S_q=1, there's only one query position at the end of the sequence - // - All KV positions are in the "past" relative to this query — nothing to mask - // - nonpad_kv_seqlen already bounds attention to valid cache positions - // - // For external cache prompt (S_q == S_kv), is_causal=1 works correctly (square matrix, - // upper-left == lower-right). For chunked prefill (S_q > 1 but S_q < S_kv), use an - // explicit attn_mask instead of is_causal. - if (causal_cross_no_past && nonpad_kv_seqlen != nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, - "Causal attention with TensorScatter (nonpad_kv_seqlen) and S_q != S_kv without " - "past_key is not supported. Per ONNX spec, is_causal without past_key produces " - "upper-left alignment where q[i] only attends to kv[0..i], which for decode (S_q=1) " - "means q[0] sees only kv[0]. Use is_causal=0 for TensorScatter decode; the KV bounds " - "are already enforced by nonpad_kv_seqlen without needing a causal mask. For chunked " - "prefill with external cache, use an explicit attn_mask instead."); - } + // is_causal=1 + nonpad_kv_seqlen (external KV cache) without past_key defines a + // bottom-right causal frontier per onnx/onnx#8068: query in-block index i attends key j + // iff j <= i + offset[b], where offset[b] = nonpad_kv_seqlen[b] - q_sequence_length. + // This shape is computed (not rejected): Flash's mha_fwd_kvcache with seqlens_k already + // produces this frontier natively (preferred fast path), and the MEA fallback builds the + // same bottom-right alignment via causal_from_top_left=false (the CUTLASS kernel sets + // causal_diagonal_offset = num_keys - num_queries == offset[b] per batch). Pure + // cross-attention without an external cache (causal_cross_no_past && nonpad_kv_seqlen == + // nullptr) keeps upper-left alignment and is handled by MEA/Unfused below. #if USE_FLASH_ATTENTION { @@ -1349,7 +1385,10 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { parameters.q_num_heads, parameters.kv_num_heads) && parameters.head_size == parameters.v_head_size && !has_output_qk && - !causal_cross_no_past && + // Upper-left causal cross-attention (no external cache) is excluded — Flash only does + // bottom-right. With an external cache (nonpad_kv_seqlen), the required frontier IS + // bottom-right, so Flash handles it via seqlens_k (onnx#8068). + (!causal_cross_no_past || nonpad_kv_seqlen != nullptr) && // Flash does not support attention masks — reject when attn_mask is present. attn_mask == nullptr; diff --git a/onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu b/onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu index 2ba7f2e1a9836..a45811d478db3 100644 --- a/onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu +++ b/onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu @@ -141,6 +141,132 @@ template Status LaunchZeroOutputForFullyMaskedBatches<__half>( template Status LaunchZeroOutputForFullyMaskedBatches<__nv_bfloat16>( __nv_bfloat16*, const int*, int, int64_t, cudaStream_t, int); +// Zero output rows fully masked by the intersection of the causal frontier and an +// explicit additive attention bias (composed is_causal + attn_mask), per onnx#8068 +// "fully-masked-row -> 0" (Bug-2). The MEA/CUTLASS path uses a finite mask sentinel, +// so a row with no allowed key softmaxes to mean-of-V; this overwrites it with zeros. +// One (batch, head, query) row per loop iteration via a grid-stride loop, so every row +// is covered even when batch_size * num_heads * q_sequence_length exceeds +// gridDim.x * blockDim.x (the grid is capped at kMaxGridDimX). Output is BSNH. +template +__global__ void ZeroFullyMaskedRowsKernel( + T* __restrict__ output, + const T* __restrict__ attn_bias, + const int* __restrict__ seqlens_k, + const int batch_size, + const int num_heads, + const int q_sequence_length, + const int total_sequence_length, + const int v_head_size, + const bool is_causal, + const bool causal_from_top_left, + const int bias_dim0, // attn_bias extent along batch: broadcast ? 1 : batch_size + const int bias_dim1, // attn_bias extent along head: broadcast ? 1 : num_heads + const float masked_bias_value) { + const int64_t total_rows = + static_cast(batch_size) * num_heads * q_sequence_length; + for (int64_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + row < total_rows; + row += static_cast(gridDim.x) * blockDim.x) { + // Decode flattened index into (batch b, head n, query i) for BNS ordering. + const int i = static_cast(row % q_sequence_length); + const int n = static_cast((row / q_sequence_length) % num_heads); + const int b = static_cast(row / (static_cast(q_sequence_length) * num_heads)); + + const int num_keys = (seqlens_k != nullptr) ? seqlens_k[b] : total_sequence_length; + + // Causally-allowed keys are j in [0, key_upper). Bottom-right anchors the diagonal at + // offset = num_keys - q_sequence_length (== nonpad_kv_seqlen[b] - q_seq); top-left uses 0. + int key_upper; + if (is_causal) { + const int offset = causal_from_top_left ? 0 : (num_keys - q_sequence_length); + const int causal_last = i + offset; // inclusive last causally-allowed key index + key_upper = min(num_keys, causal_last + 1); + } else { + key_upper = num_keys; + } + + bool any_allowed = false; + if (key_upper > 0) { + if (attn_bias == nullptr) { + any_allowed = true; // no explicit mask -> at least one causally-allowed key + } else { + const int b_idx = (bias_dim0 == 1) ? 0 : b; + const int n_idx = (bias_dim1 == 1) ? 0 : n; + const int64_t base = + ((static_cast(b_idx) * bias_dim1 + n_idx) * q_sequence_length + i) * + total_sequence_length; + for (int j = 0; j < key_upper; ++j) { + if (static_cast(attn_bias[base + j]) > masked_bias_value) { + any_allowed = true; + break; + } + } + } + } + + if (!any_allowed) { + // Output is BSNH: zero the [b, i, n, :] row (select, not multiply). + const int64_t out_base = + ((static_cast(b) * q_sequence_length + i) * num_heads + n) * v_head_size; + for (int d = 0; d < v_head_size; ++d) { + output[out_base + d] = T(0.0f); + } + } + } +} + +template +Status LaunchZeroFullyMaskedRows( + T* output, + const T* attn_bias, + const int* seqlens_k, + int batch_size, + int num_heads, + int q_sequence_length, + int total_sequence_length, + int v_head_size, + bool is_causal, + bool causal_from_top_left, + bool broadcast_bias_dim_0, + bool broadcast_bias_dim_1, + float masked_bias_value, + cudaStream_t stream, + int max_threads_per_block) { + const int64_t total_rows = + static_cast(batch_size) * num_heads * q_sequence_length; + if (total_rows == 0) { + return Status::OK(); + } + + const int bias_dim0 = broadcast_bias_dim_0 ? 1 : batch_size; + const int bias_dim1 = broadcast_bias_dim_1 ? 1 : num_heads; + + int threads = static_cast(std::min(static_cast(max_threads_per_block), total_rows)); + int64_t blocks = (total_rows + threads - 1) / threads; + // Cap the grid at kMaxGridDimX; the kernel grid-strides over total_rows, so any rows + // beyond grid_size * threads are still covered (matches ZeroOutputForFullyMaskedBatches). + constexpr int64_t kMaxGridDimX = 65535; + unsigned int grid_size = static_cast(std::min(blocks, kMaxGridDimX)); + + ZeroFullyMaskedRowsKernel<<>>( + output, attn_bias, seqlens_k, batch_size, num_heads, q_sequence_length, + total_sequence_length, v_head_size, is_causal, causal_from_top_left, + bias_dim0, bias_dim1, masked_bias_value); + + return CUDA_CALL(cudaGetLastError()); +} + +template Status LaunchZeroFullyMaskedRows( + float*, const float*, const int*, int, int, int, int, int, bool, bool, bool, bool, float, + cudaStream_t, int); +template Status LaunchZeroFullyMaskedRows<__half>( + __half*, const __half*, const int*, int, int, int, int, int, bool, bool, bool, bool, float, + cudaStream_t, int); +template Status LaunchZeroFullyMaskedRows<__nv_bfloat16>( + __nv_bfloat16*, const __nv_bfloat16*, const int*, int, int, int, int, int, bool, bool, bool, bool, + float, cudaStream_t, int); + // Simple kernel to fill an int32 buffer with a constant value on device. // Used for CUDA-graph-capturable seqlens_k initialization (no host memory). __global__ void FillInt32Kernel(int* __restrict__ output, const int value, const int count) { diff --git a/onnxruntime/core/providers/cuda/llm/attention_mask_impl.h b/onnxruntime/core/providers/cuda/llm/attention_mask_impl.h index d2cb4dbbd25ae..c31f4a7bcee29 100644 --- a/onnxruntime/core/providers/cuda/llm/attention_mask_impl.h +++ b/onnxruntime/core/providers/cuda/llm/attention_mask_impl.h @@ -45,6 +45,40 @@ Status LaunchZeroOutputForFullyMaskedBatches( cudaStream_t stream, int max_threads_per_block); +// Zero output rows that are fully masked by the intersection of the causal frontier +// and an explicit attention bias (composed is_causal + attn_mask), per onnx/onnx#8068 +// "fully-masked-row -> 0" (Bug-2). The MEA/CUTLASS path applies a finite mask sentinel +// (kCutlassSafeMaskFilterValue) rather than -inf, so a query row with no allowed key +// softmaxes to a uniform mean-of-V instead of zero. For each (batch, head, query) this +// kernel reconstructs the per-batch bottom-right/top-left causal frontier and checks the +// additive bias over the causally-allowed keys; if every such key is masked +// (bias <= masked_bias_value) the corresponding output row is overwritten with zeros +// (select-not-multiply). Output is BSNH: [batch_size, q_sequence_length, num_heads, v_head_size]. +// +// seqlens_k: per-batch valid key count (external cache); pass nullptr to use +// total_sequence_length for every batch. +// is_causal / causal_from_top_left: causal frontier selection matching the MEA params. +// broadcast_bias_dim_0 / broadcast_bias_dim_1: attn_bias broadcast over batch / head. +// masked_bias_value: the additive-bias sentinel used for masked keys (a key counts as +// masked when bias <= this value); use the same value passed to the mask conversion. +template +Status LaunchZeroFullyMaskedRows( + T* output, + const T* attn_bias, + const int* seqlens_k, + int batch_size, + int num_heads, + int q_sequence_length, + int total_sequence_length, + int v_head_size, + bool is_causal, + bool causal_from_top_left, + bool broadcast_bias_dim_0, + bool broadcast_bias_dim_1, + float masked_bias_value, + cudaStream_t stream, + int max_threads_per_block); + // Fill an int32 buffer with a constant value entirely on device. // CUDA-graph-capturable alternative to host vector + cudaMemcpyAsync. Status LaunchFillInt32(int* output, int value, int count, cudaStream_t stream, int max_threads_per_block); diff --git a/onnxruntime/test/onnx/TestCase.cc b/onnxruntime/test/onnx/TestCase.cc index 2640e0731f429..4757c064449b0 100644 --- a/onnxruntime/test/onnx/TestCase.cc +++ b/onnxruntime/test/onnx/TestCase.cc @@ -1012,6 +1012,39 @@ std::unique_ptr> GetBrokenTests(const std::string& provider "Skipped until cmake/external/onnx >= v1.22 (includes onnx/onnx#7867)"}, {"attention_4d_diff_heads_mask4d_padded_kv_expanded", "Skipped until cmake/external/onnx >= v1.22 (includes onnx/onnx#7867)"}, + // #28904: preemptive skip for the onnx#8068 Attention tests (bottom-right + // is_causal + nonpad_kv_seqlen, and composed is_causal+attn_mask + // fully-masked-row->0). The bundled cmake/external/onnx is v1.21.0 and does + // not yet contain these tests, so these entries are no-ops today; they + // prevent the new tests from failing the C++ onnx_test_runner the moment the + // onnx pin is bumped to a release including #8068, before the CPU/CUDA + // kernels land. This mirrors the jsonc filter block in + // onnx_backend_test_series_filters.jsonc (the Python series uses the jsonc; + // this C++ runner uses GetBrokenTests()). De-skip per path as each kernel + // lands, and remove entirely once the onnx pin includes #8068. + // TODO(#28904): remove this block when the CPU/CUDA kernels land and cmake/external/onnx is bumped to a release containing onnx/onnx#8068. + {"attention_4d_gqa_causal_nonpad_decode", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_4d_gqa_causal_nonpad_decode_expanded", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_4d_gqa_causal_nonpad_decode_fp16", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_4d_gqa_causal_nonpad_decode_fp16_expanded", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_4d_causal_nonpad_continued_prefill", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_4d_causal_nonpad_continued_prefill_expanded", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_causal_boolmask_nan_robustness", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_causal_boolmask_nan_robustness_expanded", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_23_boolmask_fullymasked_row_nan_robustness", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + {"attention_23_boolmask_fullymasked_row_nan_robustness_expanded", + "Skipped until #28904 kernels land and cmake/external/onnx includes onnx/onnx#8068"}, + // NOTE: the #8068 mode-3 qk_matmul_output skips are CUDA-only (CPU implements mode 3 and + // would pass) — they live in the `provider_name == "cuda"` block below, not here. See #28994. {"loop13_seq", "Creation of empty sequences is currently not supported in the test runner"}, {"sequence_insert_at_front", "shape mismatch, expect {4} got {3}"}, {"cast_FLOAT_to_BFLOAT16", "expect uint16 got bfloat16"}, @@ -1184,6 +1217,25 @@ std::unique_ptr> GetBrokenTests(const std::string& provider #else broken_tests->insert({"bidaf", "this test should be recovered when multi-gpu pipeline deprecates NV12", {"opset9"}}); #endif + // #28994: preemptive skip for the #8068 mode-3 qk_matmul_output conformance tests. CUDA returns + // NOT_IMPLEMENTED for qk_matmul_output modes beyond kQK (cuda/llm/attention.cc:1477), so these + // hard-fail the moment the onnx pin is bumped past onnx/onnx#8068. CPU IMPLEMENTS mode 3 and + // passes, so these are scoped to CUDA only (not the provider-agnostic block above) to preserve + // CPU mode-3 conformance. The bundled cmake/external/onnx is v1.21.0 and does not yet contain + // these tests, so the entries are no-ops today. Exact names sourced from the onnx#8068 generated + // corpus. De-skip when mode>kQK is supported in Attention-cuda (#27712). See #28994 Section A. + broken_tests->insert({"attention_23_fullymasked_qk_matmul_output_mode3_zero", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); + broken_tests->insert({"attention_23_fullymasked_qk_matmul_output_mode3_zero_expanded", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); + broken_tests->insert({"attention_24_fullymasked_qk_matmul_output_mode3_zero", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); + broken_tests->insert({"attention_24_fullymasked_qk_matmul_output_mode3_zero_expanded", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); + broken_tests->insert({"attention_24_qk_matmul_output_mode3_softmax_precision", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); + broken_tests->insert({"attention_24_qk_matmul_output_mode3_softmax_precision_expanded", + "CUDA NOT_IMPLEMENTED for qk_matmul_output mode>kQK; skip until onnx#8068 pin + #27712 (#28994)"}); } if (provider_name == "nnapi") { diff --git a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc b/onnxruntime/test/providers/cpu/llm/attention_op_test.cc index 54c6955114abb..cf76ca0fa00f8 100644 --- a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc +++ b/onnxruntime/test/providers/cpu/llm/attention_op_test.cc @@ -518,17 +518,116 @@ TEST(AttentionTest, Attention4DAttnMaskBoolAllFalse) { ASSERT_EQ(m.size(), q_sequence_length * kv_sequence_length); ASSERT_EQ(y.size(), batch_size * q_num_heads * q_sequence_length * v_head_size); + // Bug-2 (ONNX opset-23/24 errata, onnx#8068): every position is masked, so every query row + // is fully masked and must produce a zero output row (not the incidental mean-of-V that the + // finite mask sentinel used to yield). See the fully-masked-row guard in attention.cc. + std::fill(y.begin(), y.end(), 0.0f); + RunTest4D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length, q, k, v, std::vector(), m, std::vector(), std::vector(), -1, -1, std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), -1, TensorType::kFloat, // is_causal, qk_matmul_output_mode, scale, softcap, softmax_precision, tensor_type y, std::vector(), std::vector(), std::vector(), - // Note: all-false bool mask (every position masked) is a degenerate case. It works because - // mask_filter_value (~-3.4e38) is so extreme that float precision loses QK differences, - // producing uniform softmax weights matching CPU behavior. + // All-false bool mask (every position masked): per the opset-23/24 errata a fully-masked + // query row produces a zero output row. false, false, true // disable_cpu, disable_cuda, disable_dml ); } +// Bug-1 (onnx#28958, tianleiwu review): a user-supplied float `attn_mask` that uses IEEE -inf to +// mask a position must be recognized as masked by the CPU fully-masked-row guard, exactly like the +// internal finite sentinel and like the CUDA EP / the onnx#8068 reference (`isneginf`). A row whose +// every key is a user -inf is fully masked and must yield a zero output row (and a zero mode-3 +// qk_matmul_output row), not a NaN row. A finite key (e.g. mask bias 0) on a row keeps that row +// finite, proving the guard does not over-zero. +TEST(AttentionTest, Attention4DAttnMaskFloatNegInfFullyMasked) { + int batch_size = 1; // Q.shape[0] + int q_num_heads = 1; // Q.shape[1] + int q_sequence_length = 2; // Q.shape[2] + int head_size = 4; // Q.shape[3] + int kv_sequence_length = 2; // K.shape[2] and V.shape[2] + int kv_num_heads = 1; // K.shape[1] and V.shape[1] + int v_head_size = 4; // V.shape[3] + int past_sequence_length = 0; // past_key.shape[2] and past_value.shape[2] + + // Q and K are arbitrary: row 0 is fully masked (zeroed) and row 1 attends exactly one key, so the + // softmax over a single retained key is 1.0 regardless of the scores -> the goldens are exact. + std::vector q = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f}; + std::vector k = {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f}; + // V key 0 = {1,2,3,4}, V key 1 = {5,6,7,8}. + std::vector v = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + + const float ninf = -std::numeric_limits::infinity(); + // Float additive mask {q_sequence_length, total_sequence_length} = {2, 2}: + // row 0: both keys -inf -> fully masked -> zero row (this is the merge-blocker case: before the + // fix the CPU guard missed user -inf and emitted a NaN row). + // row 1: key 0 finite (0), key 1 -inf -> attends only key 0 -> Y = V[key0], softmax = [1, 0]. + std::vector attn_mask = {ninf, ninf, 0.0f, ninf}; + + // Y {batch, q_num_heads, q_sequence_length, v_head_size} = {1,1,2,4}. + std::vector y = {0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + // Mode-3 qk_matmul_output (post-softmax) {1,1,2,2}: masked row 0 zeroed; row 1 = [1, 0]. + std::vector qk_matmul_output = {0.0f, 0.0f, 1.0f, 0.0f}; + + ASSERT_EQ(q.size(), static_cast(batch_size * q_num_heads * q_sequence_length * head_size)); + ASSERT_EQ(k.size(), static_cast(batch_size * kv_num_heads * kv_sequence_length * head_size)); + ASSERT_EQ(v.size(), static_cast(batch_size * kv_num_heads * kv_sequence_length * v_head_size)); + ASSERT_EQ(attn_mask.size(), static_cast(q_sequence_length * kv_sequence_length)); + ASSERT_EQ(y.size(), static_cast(batch_size * q_num_heads * q_sequence_length * v_head_size)); + + RunTest4D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length, + q, k, v, attn_mask, {}, std::vector(), std::vector(), + -1, 3, std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), -1, TensorType::kFloat, // is_causal, qk_matmul_output_mode=3, scale, softcap, softmax_precision, tensor_type + y, std::vector(), std::vector(), qk_matmul_output, + // CPU-only: this pins the CPU float -inf predicate fix (onnx#28958). CUDA/DML masking is + // covered separately and uses a different masked-bias guard. + false, true, true // disable_cpu, disable_cuda, disable_dml + ); +} + +// FP16 twin of Attention4DAttnMaskFloatNegInfFullyMasked: exercises the MLFloat16 branch of the +// CPU predicate's to_float conversion with a user fp16 -inf attn_mask. Same structure/goldens. +TEST(AttentionTest, Attention4DAttnMaskFloat16NegInfFullyMasked) { + int batch_size = 1; // Q.shape[0] + int q_num_heads = 1; // Q.shape[1] + int q_sequence_length = 2; // Q.shape[2] + int head_size = 4; // Q.shape[3] + int kv_sequence_length = 2; // K.shape[2] and V.shape[2] + int kv_num_heads = 1; // K.shape[1] and V.shape[1] + int v_head_size = 4; // V.shape[3] + int past_sequence_length = 0; // past_key.shape[2] and past_value.shape[2] + + // Q and K are arbitrary: row 0 is fully masked (zeroed) and row 1 attends exactly one key, so the + // softmax over a single retained key is 1.0 regardless of the scores -> the goldens are exact. + std::vector q = {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f, 0.7f, 0.8f}; + std::vector k = {0.9f, 0.8f, 0.7f, 0.6f, 0.5f, 0.4f, 0.3f, 0.2f}; + // V key 0 = {1,2,3,4}, V key 1 = {5,6,7,8}. + std::vector v = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + + const float ninf = -std::numeric_limits::infinity(); + // FP16 additive mask {2, 2}: row 0 both keys -inf (fully masked -> zero row); row 1 key 0 finite, + // key 1 -inf (attends only key 0 -> Y = V[key0], softmax = [1, 0]). + std::vector attn_mask = {ninf, ninf, 0.0f, ninf}; + + // Y {1,1,2,4}: masked row 0 zeroed; row 1 = V[key0]. + std::vector y = {0.0f, 0.0f, 0.0f, 0.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + // Mode-3 qk_matmul_output (post-softmax) {1,1,2,2}: masked row 0 zeroed; row 1 = [1, 0]. + std::vector qk_matmul_output = {0.0f, 0.0f, 1.0f, 0.0f}; + + ASSERT_EQ(q.size(), static_cast(batch_size * q_num_heads * q_sequence_length * head_size)); + ASSERT_EQ(k.size(), static_cast(batch_size * kv_num_heads * kv_sequence_length * head_size)); + ASSERT_EQ(v.size(), static_cast(batch_size * kv_num_heads * kv_sequence_length * v_head_size)); + ASSERT_EQ(attn_mask.size(), static_cast(q_sequence_length * kv_sequence_length)); + ASSERT_EQ(y.size(), static_cast(batch_size * q_num_heads * q_sequence_length * v_head_size)); + + RunTest4D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length, + q, k, v, attn_mask, {}, std::vector(), std::vector(), + -1, 3, std::numeric_limits::quiet_NaN(), std::numeric_limits::quiet_NaN(), -1, TensorType::kFloat16, // is_causal, qk_matmul_output_mode=3, scale, softcap, softmax_precision, tensor_type + y, std::vector(), std::vector(), qk_matmul_output, + // CPU-only: pins the MLFloat16 branch of the CPU -inf predicate fix (onnx#28958). + false, true, true // disable_cpu, disable_cuda, disable_dml + ); +} + // Regression guard: all-false bool mask in decode mode (past_sequence_length > 0). // Guards against a bug where fully-masked batches produce NaN or incorrect output. // Expected behavior: uniform softmax over all KV values produces Y = mean-of-V. @@ -593,17 +692,14 @@ TEST(AttentionTest, Attention4DAttnMaskBoolAllFalseDecodeWithPast) { 0.7f, 0.7f, 0.7f, 0.7f, 0.7f, 0.7f, 0.7f, 0.7f, 0.8f, 0.8f, 0.8f, 0.8f, 0.8f, 0.8f, 0.8f, 0.8f}; - // With all-false mask, softmax produces uniform weights: 1/4 per position. - // Standard concat places new token at past_sequence_length, so present_value = - // [past_v[0], past_v[1], past_v[2], new_v]. Output = mean of all V rows: - // head 0: mean(0.1, 0.2, 0.3, 0.4) = 0.25 - // head 1: mean(0.5, 0.6, 0.7, 0.8) = 0.65 - // These values match upstream/main behavior (unfused standard concat path). + // With all-false mask every key is disallowed, so the single decode query row is fully masked. + // Per the opset-23/24 errata (onnx#8068) a fully-masked row produces a zero output row, not the + // incidental mean-of-V that the finite mask sentinel used to yield. std::vector y = { // head 0 - 0.25f, 0.25f, 0.25f, 0.25f, 0.25f, 0.25f, 0.25f, 0.25f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, // head 1 - 0.65f, 0.65f, 0.65f, 0.65f, 0.65f, 0.65f, 0.65f, 0.65f}; + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; ASSERT_EQ(q.size(), batch_size * q_num_heads * q_sequence_length * head_size); ASSERT_EQ(k.size(), batch_size * kv_num_heads * kv_sequence_length * head_size); @@ -2000,8 +2096,7 @@ TEST(AttentionTest, Attention_NonPadKVSeqLen_AllMasked) { std::vector q = {1.0f, 1.0f}; std::vector k(8, 1.0f); - // All V positions are "invalid" — result is uniform over 4 highly-negative-scored positions, - // so softmax is uniform 0.25 each. + // All KV positions are invalid (nonpad=0), so the single query row is fully masked. std::vector v = {10.0f, 20.0f, 30.0f, 40.0f, 50.0f, 60.0f, 70.0f, 80.0f}; test.AddInput("Q", q_shape, q); @@ -2012,9 +2107,9 @@ TEST(AttentionTest, Attention_NonPadKVSeqLen_AllMasked) { test.AddOptionalInputEdge(); test.AddInput("nonpad_kv_seqlen", {1}, {0}); - // With all positions masked to -inf, softmax produces uniform weights and - // the result is the mean of all V rows: [(10+30+50+70)/4, (20+40+60+80)/4] = [40, 50]. - std::vector expected_y = {40.0f, 50.0f}; + // Per the opset-23/24 errata (onnx#8068) a fully-masked query row produces a zero output row, + // not the incidental mean-of-V that the finite mask sentinel used to yield. + std::vector expected_y = {0.0f, 0.0f}; test.AddOutput("Y", {1, 1, 1, 2}, expected_y, false, 0, 1e-3f); test.AddOptionalOutputEdge(); test.AddOptionalOutputEdge(); @@ -2188,6 +2283,202 @@ TEST(AttentionTest, Attention_NonPadKVSeqLen_NoneMasked) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// Bottom-right (offset-aware) is_causal + nonpad_kv_seqlen and composed +// is_causal + attn_mask fully-masked-row -> 0 (ONNX opset-24 / onnx#8068). +// +// These goldens are EP-independent (analytic): constant Q/K make every attended +// score equal, so the output of each query row is the mean of V over exactly the +// set of keys the bottom-right/composed mask allows. The CUDA developer reuses +// these golden values and adds CUDA-dispatch assertions (Flash-selected / not +// NOT_IMPLEMENTED) for the same shapes. +// --------------------------------------------------------------------------- + +// External/static-cache decode (S_q=1) with is_causal + per-batch valid lengths. +// With bottom-right alignment the single decode query attends keys 0..nonpad[b]-1. +// Mirrors onnx test_attention_4d_gqa_causal_nonpad_decode. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_Decode_BottomRight) { + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + // batch=2, heads=1, q_seq=1, kv_seq=4, head_size=2. Batch 0 has 4 valid keys, batch 1 has 2. + std::vector q_shape = {2, 1, 1, 2}; + std::vector k_shape = {2, 1, 4, 2}; + std::vector v_shape = {2, 1, 4, 2}; + + std::vector q(2 * 1 * 1 * 2, 1.0f); + std::vector k(2 * 1 * 4 * 2, 1.0f); + std::vector v = { + 1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f, // batch 0 + 10.0f, 10.0f, 20.0f, 20.0f, 30.0f, 30.0f, 40.0f, 40.0f // batch 1 + }; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {2}, {4, 2}); + + // Batch 0: mean(V[0..3]) = [2.5, 2.5]. Batch 1: mean(V[0..1]) = [15, 15]. + std::vector expected_y = {2.5f, 2.5f, 15.0f, 15.0f}; + test.AddOutput("Y", {2, 1, 1, 2}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Continued / chunked prefill (S_q=2) into a partially-filled static cache. +// nonpad=[4], S_q=2 -> offset = 4 - 2 = 2: query 0 attends keys {0,1,2}, query 1 +// attends {0,1,2,3}. The old top-left alignment would mask everything past the +// diagonal ({0} and {0,1}), so this case fails pre-fix and passes post-fix. +// Mirrors onnx test_attention_4d_causal_nonpad_continued_prefill. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_ContinuedPrefill_BottomRight) { + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + // batch=1, heads=1, q_seq=2, kv_seq=4, head_size=2. + std::vector q_shape = {1, 1, 2, 2}; + std::vector k_shape = {1, 1, 4, 2}; + std::vector v_shape = {1, 1, 4, 2}; + + std::vector q(1 * 1 * 2 * 2, 1.0f); + std::vector k(1 * 1 * 4 * 2, 1.0f); + std::vector v = {1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f}; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {4}); + + // Query 0: mean(V[0..2]) = [2, 2]. Query 1: mean(V[0..3]) = [2.5, 2.5]. + std::vector expected_y = {2.0f, 2.0f, 2.5f, 2.5f}; + test.AddOutput("Y", {1, 1, 2, 2}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Composed is_causal + boolean attn_mask, fully-masked-row -> 0 (Bug-2). +// Causal frontier (offset 0): query 0 attends {0}, query 1 attends {0,1}. The bool +// mask [[T,F],[F,F]] intersects: query 0 keeps key 0 (output = V[0]); query 1 has no +// allowed key (fully masked) and must output an all-zero row, not mean-of-V / NaN. +// Mirrors onnx test_attention_causal_boolmask_nan_robustness. +TEST(AttentionTest, Attention_Causal_BoolMask_FullyMaskedRow_Zero) { + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + // batch=1, heads=1, q_seq=2, kv_seq=2, head_size=2. + std::vector q_shape = {1, 1, 2, 2}; + std::vector k_shape = {1, 1, 2, 2}; + std::vector v_shape = {1, 1, 2, 2}; + + std::vector q(1 * 1 * 2 * 2, 1.0f); + std::vector k(1 * 1 * 2 * 2, 1.0f); + std::vector v = {5.0f, 7.0f, 9.0f, 11.0f}; // V[0]=[5,7], V[1]=[9,11] + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + // 2D bool mask (broadcast over batch/heads): row 0 allows key 0, row 1 allows none. + test.AddInput("attn_mask", {2, 2}, {true, false, false, false}); + + // Query 0 attends only key 0 -> V[0] = [5, 7]. Query 1 is fully masked -> [0, 0]. + std::vector expected_y = {5.0f, 7.0f, 0.0f, 0.0f}; + test.AddOutput("Y", {1, 1, 2, 2}, expected_y, false, 0, 1e-4f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// M-1 regression guard (onnx#8068): a row whose additive bias is FINITE but very negative must NOT be +// treated as fully masked. The old detector used a 0.5x sentinel magnitude threshold; in fp16 the +// sentinel is MLFloat16::MinValue (-65504), so a finite user bias below 0.5*sentinel (-32752) — e.g. +// -40000 — was wrongly classified as fully masked and zeroed. The exact structural predicate keys off +// sentinel-equality on the bias slot, so a finite bias keeps its key unmasked: equal bias on both keys +// cancels in softmax and the row attends uniformly (mean-of-V), it is never zeroed. +TEST(AttentionTest, Attention_FiniteNegativeBias_RowNotZeroed_Fp16) { + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + + // batch=1, heads=1, q_seq=1, kv_seq=2, head_size=2. + std::vector q_shape = {1, 1, 1, 2}; + std::vector k_shape = {1, 1, 2, 2}; + std::vector v_shape = {1, 1, 2, 2}; + + std::vector q(1 * 1 * 1 * 2, 1.0f); + std::vector k(1 * 1 * 2 * 2, 1.0f); + std::vector v = {5.0f, 7.0f, 9.0f, 11.0f}; // V[0]=[5,7], V[1]=[9,11] + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + // Finite, very negative additive bias on BOTH keys (below the old fp16 0.5x threshold of -32752, + // but NOT the -65504 sentinel). Equal bias on both keys -> softmax is uniform. + std::vector attn_mask = {-40000.0f, -40000.0f}; + test.AddInput("attn_mask", {1, 1, 1, 2}, ToFloat16(attn_mask)); + + // Row is NOT fully masked: attends both keys equally -> mean(V) = [(5+9)/2, (7+11)/2] = [7, 9], + // not [0, 0]. + std::vector expected_y = {7.0f, 9.0f}; + test.AddOutput("Y", {1, 1, 1, 2}, ToFloat16(expected_y), false, 0, 3e-3f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// M-1 structural-empty guard (onnx#8068): the fully-masked-row guard must fire purely from the +// causal/nonpad frontier, with NO attn_mask present. With is_causal and nonpad_kv_seqlen < q_seqlen, +// the bottom-right offset (nonpad - q_seqlen) is negative, so early query rows have zero allowed keys +// (frontier collapses to 0). Those rows must output 0, not mean-of-V. q_seqlen=4, nonpad=2 -> rows 0,1 +// are structurally empty; row 2 attends key 0; row 3 attends keys 0,1. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_StructuralEmptyRow_Zero) { + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + // batch=1, heads=1, q_seq=4, kv_seq=4, head_size=2. + std::vector q_shape = {1, 1, 4, 2}; + std::vector k_shape = {1, 1, 4, 2}; + std::vector v_shape = {1, 1, 4, 2}; + + std::vector q(1 * 1 * 4 * 2, 1.0f); + std::vector k(1 * 1 * 4 * 2, 1.0f); + // Only the first two keys are valid (nonpad=2); V[0]=[5,7], V[1]=[9,11], rest unused. + std::vector v = {5.0f, 7.0f, 9.0f, 11.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask: none (structural-only masking) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // offset = nonpad - q_seqlen = -2. Row i attends keys j <= i + offset AND j < nonpad: + // i=0,1 -> no allowed keys -> [0, 0] (structural-empty -> zeroed) + // i=2 -> key 0 -> V[0] = [5, 7] + // i=3 -> keys 0,1 -> mean = [(5+9)/2, (7+11)/2] = [7, 9] + std::vector expected_y = {0.0f, 0.0f, 0.0f, 0.0f, 5.0f, 7.0f, 7.0f, 9.0f}; + test.AddOutput("Y", {1, 1, 4, 2}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); + test.AddOptionalOutputEdge(); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCpuExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + // Validation: negative nonpad_kv_seqlen should be rejected. TEST(AttentionTest, Attention_NonPadKVSeqLen_NegativeValue) { OpTester test("Attention", 24, onnxruntime::kOnnxDomain); @@ -3315,5 +3606,804 @@ TEST(AttentionTest, Attention_QkMatmulOutputMode_PostMaskBias_WithSoftcapAndNonp test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// CUDA-dispatch assertions for the onnx#8068 bottom-right / fully-masked-row +// semantics (issue microsoft/onnxruntime#28904, CUDA half). +// +// These mirror the EP-independent goldens authored above for the CPU EP and run +// them on the CUDA EP. Their purpose is dispatch-level: assert that the +// is_causal=1 + nonpad_kv_seqlen (no past_key) shape is ACCEPTED on CUDA (it used +// to return NOT_IMPLEMENTED), that the MEA bottom-right fallback and the unfused +// fully-masked-row -> 0 guard compute the correct result, and that the no-mask +// fp16 decode shape routes through the Flash fast path (head_size=64 so Flash's +// is_supported() gate passes). Constant Q/K keep every attended score equal, so +// each output row is the analytic mean of V over the allowed keys regardless of +// head_size or EP. Tests are skipped when no CUDA device is present. +// --------------------------------------------------------------------------- + +// External/static-cache decode (S_q=1), is_causal + per-batch valid lengths, fp32. +// No mask -> Flash-ineligible for fp32, so this exercises the MEA bottom-right path +// (causal_from_top_left=false). Reuses the goldens of +// Attention_Causal_NonPadKVSeqLen_Decode_BottomRight. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_Decode_BottomRight_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {2, 1, 1, 2}; + std::vector k_shape = {2, 1, 4, 2}; + std::vector v_shape = {2, 1, 4, 2}; + + std::vector q(2 * 1 * 1 * 2, 1.0f); + std::vector k(2 * 1 * 4 * 2, 1.0f); + std::vector v = { + 1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f, // batch 0 + 10.0f, 10.0f, 20.0f, 20.0f, 30.0f, 30.0f, 40.0f, 40.0f // batch 1 + }; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {2}, {4, 2}); + + std::vector expected_y = {2.5f, 2.5f, 15.0f, 15.0f}; + test.AddOutput("Y", {2, 1, 1, 2}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Continued / chunked prefill (S_q=2) into a partially-filled static cache, fp32. +// offset = nonpad - S_q = 4 - 2 = 2. Reuses the goldens of +// Attention_Causal_NonPadKVSeqLen_ContinuedPrefill_BottomRight. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_ContinuedPrefill_BottomRight_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 2, 2}; + std::vector k_shape = {1, 1, 4, 2}; + std::vector v_shape = {1, 1, 4, 2}; + + std::vector q(1 * 1 * 2 * 2, 1.0f); + std::vector k(1 * 1 * 4 * 2, 1.0f); + std::vector v = {1.0f, 1.0f, 2.0f, 2.0f, 3.0f, 3.0f, 4.0f, 4.0f}; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {4}); + + std::vector expected_y = {2.0f, 2.0f, 2.5f, 2.5f}; + test.AddOutput("Y", {1, 1, 2, 2}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Composed is_causal + boolean attn_mask, fully-masked-row -> 0 (Bug-2), fp32. +// total_sequence_length=2 (not divisible by 4) makes MEA ineligible, so this +// exercises the unified-unfused fully-masked-row guard on CUDA. Reuses the goldens +// of Attention_Causal_BoolMask_FullyMaskedRow_Zero. +TEST(AttentionTest, Attention_Causal_BoolMask_FullyMaskedRow_Zero_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 2, 2}; + std::vector k_shape = {1, 1, 2, 2}; + std::vector v_shape = {1, 1, 2, 2}; + + std::vector q(1 * 1 * 2 * 2, 1.0f); + std::vector k(1 * 1 * 2 * 2, 1.0f); + std::vector v = {5.0f, 7.0f, 9.0f, 11.0f}; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddInput("attn_mask", {2, 2}, {true, false, false, false}); + + std::vector expected_y = {5.0f, 7.0f, 0.0f, 0.0f}; + test.AddOutput("Y", {1, 1, 2, 2}, expected_y, false, 0, 1e-4f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// fp16 external-cache decode with a Flash-eligible head_size (64). No attn_mask, so +// the is_causal=1 + nonpad_kv_seqlen shape routes to the Flash fast path (which does +// bottom-right + per-batch seqlens_k masking natively). Constant Q/K make each output +// row the analytic mean of V over the allowed keys, independent of head_size. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_Decode_FlashFP16_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + // batch=2, heads=1, q_seq=1, kv_seq=4, head_size=64. Batch 0: 4 valid keys, batch 1: 2. + std::vector q_shape = {2, 1, 1, kHeadSize}; + std::vector k_shape = {2, 1, 4, kHeadSize}; + std::vector v_shape = {2, 1, 4, kHeadSize}; + + std::vector q(2 * 1 * 1 * kHeadSize, 1.0f); + std::vector k(2 * 1 * 4 * kHeadSize, 1.0f); + // Each (batch, key) row of V is a constant value so the mean is exact in fp16. + const float v_row_values[2][4] = {{1.0f, 2.0f, 3.0f, 4.0f}, {10.0f, 20.0f, 30.0f, 40.0f}}; + std::vector v; + v.reserve(2 * 4 * kHeadSize); + for (int b = 0; b < 2; ++b) { + for (int s = 0; s < 4; ++s) { + for (int h = 0; h < kHeadSize; ++h) { + v.push_back(v_row_values[b][s]); + } + } + } + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {2}, {4, 2}); + + // Batch 0: mean(1,2,3,4)=2.5. Batch 1: mean(10,20)=15. Broadcast across head_size. + std::vector expected_y_f; + expected_y_f.reserve(2 * kHeadSize); + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(2.5f); + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(15.0f); + test.AddOutput("Y", {2, 1, 1, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Direct MEA (CUTLASS) fp16 fully-masked-row golden. head_size=8 (MEA-eligible, Flash-ineligible) +// + total_sequence_length=4 (MEA-eligible) + nonpad_kv_seqlen routes to the MEA bottom-right +// path, and the bool attn_mask drives LaunchZeroFullyMaskedRows directly. Batch 0 is fully +// masked (-> zero row); batch 1 retains all keys (-> mean of V). This covers the MEA +// ZeroFullyMaskedRows partial+full path that the unfused tests only reach indirectly. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_FP16_FullyMaskedRow_Zero_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + constexpr int kHeadSize = 8; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {2, 1, 1, kHeadSize}; + std::vector k_shape = {2, 1, 4, kHeadSize}; + std::vector v_shape = {2, 1, 4, kHeadSize}; + + std::vector q(2 * 1 * 1 * kHeadSize, 1.0f); + std::vector k(2 * 1 * 4 * kHeadSize, 1.0f); + // Batch 1 V rows are constant per key so the retained-row mean is exact in fp16. + const float v_row_values[2][4] = {{0.0f, 0.0f, 0.0f, 0.0f}, {1.0f, 2.0f, 3.0f, 4.0f}}; + std::vector v; + v.reserve(2 * 4 * kHeadSize); + for (int b = 0; b < 2; ++b) + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[b][s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + // attn_mask (bool, broadcast over heads): batch 0 fully masked, batch 1 fully allowed. + test.AddInput("attn_mask", {2, 1, 1, 4}, + {false, false, false, false, true, true, true, true}); + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {2}, {4, 4}); + + // Batch 0: fully masked -> zero row. Batch 1: mean(1,2,3,4)=2.5 across head_size. + std::vector expected_y_f; + expected_y_f.reserve(2 * kHeadSize); + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(0.0f); + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(2.5f); + test.AddOutput("Y", {2, 1, 1, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// M-1 regression guard: a FINITE but very-negative fp16 additive bias must NOT be treated as a +// fully-masked row. head_size=4 + total_sequence_length=2 route to the unfused kernel (the path +// the prior 0.5x score-collapse heuristic wrongly zeroed). Every key carries bias=-40000 +// (finite, strictly greater than the -65504 fp16 sentinel), so all keys are retained and each +// row softmaxes to the uniform mean of V rather than zero. +TEST(AttentionTest, Attention_Unfused_FiniteNegativeBias_NotZeroed_FP16_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + constexpr int kHeadSize = 4; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(0)); + + std::vector q_shape = {1, 1, 2, kHeadSize}; + std::vector k_shape = {1, 1, 2, kHeadSize}; + std::vector v_shape = {1, 1, 2, kHeadSize}; + + std::vector q(1 * 1 * 2 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 2 * kHeadSize, 1.0f); + // Key 0 rows = 5, key 1 rows = 9 -> uniform-softmax mean = 7. + std::vector v = {5.0f, 5.0f, 5.0f, 5.0f, 9.0f, 9.0f, 9.0f, 9.0f}; + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + // Additive fp16 bias of -40000 on every (query,key): finite but below the old 0.5x threshold + // (0.5 * -65504 = -32752). The old heuristic wrongly zeroed; the exact predicate keeps the row. + test.AddInput("attn_mask", {2, 2}, ToFloat16(std::vector(4, -40000.0f))); + + // Uniform softmax over the 2 equally-biased keys -> mean(5,9)=7 for both query rows. + std::vector expected_y(1 * 1 * 2 * kHeadSize, 7.0f); + test.AddOutput("Y", {1, 1, 2, kHeadSize}, ToFloat16(expected_y), false, 0, 0.05f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Structural fully-masked rows with NO attn_mask: is_causal bottom-right + nonpad_kv_seqlen < q +// leaves the early query rows with an empty causal frontier (offset = nonpad - q_seq < 0), so +// they must output zero per onnx#8068 even though no attn_mask is present. fp32 head_size=4 + +// total_sequence_length=4 routes to the UNFUSED path (fp32 disables Flash at attention.cc:114, and +// head_size=4 fails MEA's (head_size&7)==0 eligibility at memory_efficient_attention.h:68, so it +// falls through to the unified unfused path at attention.cc:1463/1482 — verified via verbose +// dispatch logging), exercising the ungated (is_causal) fully-masked guard on that path. +// q_seq=4, nonpad=2: rows 0,1 see zero allowed keys -> 0; row 2 sees key 0; row 3 sees +// keys 0,1. Constant Q/K make retained rows the uniform mean of the allowed V rows. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_StructuralEmptyRows_Zero_CUDA) { + if (!HasCudaEnvironment(0)) { + GTEST_SKIP() << "CUDA device not available"; + } + constexpr int kHeadSize = 4; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + // Key 0 rows = 1, key 1 rows = 3 (keys 2,3 are padding beyond nonpad=2 and never attended). + std::vector v = {1.0f, 1.0f, 1.0f, 1.0f, 3.0f, 3.0f, 3.0f, 3.0f, + 7.0f, 7.0f, 7.0f, 7.0f, 9.0f, 9.0f, 9.0f, 9.0f}; + + test.AddInput("Q", q_shape, q); + test.AddInput("K", k_shape, k); + test.AddInput("V", v_shape, v); + test.AddOptionalInputEdge(); // attn_mask (none — masking is purely structural) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // Rows 0,1 structurally empty -> 0. Row 2 -> V[0]=1. Row 3 -> mean(V[0],V[1])=2. + std::vector expected_y = {0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, + 1.0f, 1.0f, 1.0f, 1.0f, + 2.0f, 2.0f, 2.0f, 2.0f}; + test.AddOutput("Y", {1, 1, 4, kHeadSize}, expected_y, false, 0, 1e-4f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Structural fully-masked rows routed through the FLASH path, SINGLE-TILE epilogue (regression for +// #28958). Unlike Attention_Causal_NonPadKVSeqLen_StructuralEmptyRows_Zero_CUDA (fp32, head_size=4 +// -> UNFUSED: fp32 disables Flash at attention.cc:114, and head_size=4 fails MEA's (head_size&7)==0 +// eligibility at memory_efficient_attention.h:68, so it routes to the unified unfused path at +// attention.cc:1463/1482 — verified via verbose dispatch logging), this uses fp16 + head_size=64 +// (== v_head_size) + is_causal=1 + nonpad_kv_seqlen < q_seq + +// NO attn_mask + NO past_key, which satisfies flash_eligible (cuda/llm/attention.cc:1385-1397) and +// routes to RunFlashAttention nonpad Path-1 (mha_fwd_kvcache with seqlens_k). The Flash path has no +// host-side LaunchZeroFullyMaskedRows guard (that runs only on MEA/unfused), so a structurally empty +// leading row depends entirely on the FA2 epilogue zeroing an lse==-inf row instead of emitting NaN. +// total_sequence_length=4 -> num_n_blocks=1 -> num_splits=1, so this exercises ONLY the single-tile +// epilogue (flash_attention/softmax.h:182-186, inv_sum=1 when sum==0 and acc_o=0). The split-KV +// combine is locked by the _SplitKV_ test below. offset = nonpad - q_seq = 2 - 4 = -2: rows 0,1 +// attend zero keys -> 0; row 2 -> V row 0; row 3 -> mean(V row 0, V row 1). Constant Q/K make +// retained rows the uniform mean of the allowed V rows. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_FlashStructuralEmptyRows_Zero_FP16_CUDA) { + if (!HasCudaEnvironment(800)) { + return; // Flash requires SM 8.0+; on older GPUs this silently falls to MEA, losing the FA2 intent. + } + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + // V row 0 = 1, V row 1 = 3 (V rows 2,3 are padding beyond nonpad=2 and never attended). + const float v_row_values[4] = {1.0f, 3.0f, 7.0f, 9.0f}; + std::vector v; + v.reserve(4 * kHeadSize); + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none — masking is purely structural) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // Rows 0,1 structurally empty -> 0. Row 2 -> V[0]=1. Row 3 -> mean(V[0],V[1])=2. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {0.0f, 0.0f, 1.0f, 2.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// BFloat16 counterpart of the Flash single-tile structural-empty-row regression above (#28958). +// BFloat16 Flash requires SM 8.0+. total_sequence_length=4 -> num_splits=1, so this also exercises +// only the single-tile epilogue (flash_attention/softmax.h:182-186). Same shapes/expectations: +// leading offset<0 rows must be 0, not NaN. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_FlashStructuralEmptyRows_Zero_BF16_CUDA) { + if (!HasCudaEnvironment(800)) { + return; // BFloat16 Flash requires SM 8.0+ + } + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + // V row 0 = 1, V row 1 = 3 (V rows 2,3 are padding beyond nonpad=2 and never attended). + const float v_row_values[4] = {1.0f, 3.0f, 7.0f, 9.0f}; + std::vector v; + v.reserve(4 * kHeadSize); + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, FloatsToBFloat16s(q)); + test.AddInput("K", k_shape, FloatsToBFloat16s(k)); + test.AddInput("V", v_shape, FloatsToBFloat16s(v)); + test.AddOptionalInputEdge(); // attn_mask (none — masking is purely structural) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // Rows 0,1 structurally empty -> 0. Row 2 -> V row 0 = 1. Row 3 -> mean(V row 0, V row 1) = 2. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {0.0f, 0.0f, 1.0f, 2.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, FloatsToBFloat16s(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Structural fully-masked rows routed through the FLASH SPLIT-KV combine path (regression for +// #28958). This is the higher NaN-risk path the single-tile tests above do NOT cover: with a large +// total_sequence_length the FA2 num_splits heuristic chooses num_splits > 1, so each masked split +// emits lse == -inf and the cross-split combine must avoid expf(-inf - -inf) -> NaN. fp16 + +// head_size=64 + total_sequence_length=257 forces num_n_blocks = ceil(257/256) = 2 -> num_splits = 2 +// on an A100-class GPU (num_splits_heuristic in flash_attention/flash_api.cc with block_n=256 for +// head_size<=64, 108 SMs). nonpad_kv_seqlen=2 with q_seq=4 (offset = 2 - 4 = -2) leaves rows 0,1 +// structurally empty across every split; the remaining 255 KV positions are padding beyond nonpad +// and never attended. This locks the split-KV combine guards (flash_attention/flash_fwd_kernel.h: +// 1125 lse_max==-inf -> 0, :1135 lse_sum==0 -> lse_logsum=+inf, :1153 scale=expf(-inf)=0 -> O=0). +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_FlashStructuralEmptyRows_Zero_SplitKV_FP16_CUDA) { + if (!HasCudaEnvironment(800)) { + return; // Flash requires SM 8.0+; otherwise this falls to MEA and the split-KV intent is lost. + } + constexpr int kHeadSize = 64; + constexpr int kKvSeqLen = 257; // > block_n (256) -> num_n_blocks=2 -> num_splits=2 (split-KV path) + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, kKvSeqLen, kHeadSize}; + std::vector v_shape = {1, 1, kKvSeqLen, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * kKvSeqLen * kHeadSize, 1.0f); + // V row 0 = 1, V row 1 = 3 (V rows 2..256 are padding beyond nonpad=2 and never attended). + std::vector v; + v.reserve(kKvSeqLen * kHeadSize); + for (int s = 0; s < kKvSeqLen; ++s) { + const float val = (s == 0) ? 1.0f : (s == 1 ? 3.0f : 99.0f); + for (int h = 0; h < kHeadSize; ++h) v.push_back(val); + } + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none — masking is purely structural) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // Rows 0,1 structurally empty -> 0. Row 2 -> V row 0 = 1. Row 3 -> mean(V row 0, V row 1) = 2. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {0.0f, 0.0f, 1.0f, 2.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// ============================================================================ +// Memory-Efficient Attention (MEA / cutlass FMHA) negative-offset regression tests (#28958). +// +// The _FlashStructuralEmptyRows_ tests above use head_size=64, which is Flash-eligible in a normal +// full build (flash_api.cc:415: head_size%8==0 && head_size<=256) and therefore routes to FA2. +// They only exercise the buggy MEA path when QUICK_BUILD=ON disables the hdim-64 flash kernel. +// +// The tests below force the MEA path in EVERY build config (independent of the QUICK_BUILD CI leg) by +// setting ORT_DISABLE_FLASH_ATTENTION=1 at head_size=64: Flash is turned off by the env var while MEA +// stays eligible (head_size%8==0 && head_size<=1024, no attn_mask, no past_key, non-GQA -> attention.cc +// :1415-1429 mea_eligible=true), so dispatch lands on attention.cc:1457 RunMemoryEfficientAttention. +// The only other reachable branch (the correct unfused path at attention.cc:1463) is excluded, so a +// value mismatch can only originate in the MEA kernel -- this structural guarantee is what guards the +// negative causal_diagonal_offset fix (kernel_forward.h:926). ScopedEnvironmentVariables is re-read per +// run because DefaultCudaExecutionProvider() builds a fresh EP each Run and +// AttentionKernelOptions::InitializeOnce (std::call_once per EP) re-parses the env (validated). +// +// WHY head_size=64 and NOT head_size=512: an earlier revision forced MEA flag-free via head_size=512 +// (Flash caps at head_size<=256 in all builds). That cutlass kernel's SharedStorage exceeds the ~99KB +// dynamic-shared-memory opt-in cap on sm86/sm89 CI GPUs; fmha_launch_template.h:241-247 calls +// cudaFuncSetAttribute(cudaFuncAttributeMaxDynamicSharedMemorySize) UNCHECKED -- gated only by +// ORT_ENFORCE(sm>=70), not by actual device capacity -- so the attribute-set silently fails and the +// subsequent launch dies with "CUDA failure 1: invalid argument" (surfacing later at +// attention_mask_impl.cu:134). It passed only on sm90 (227KB opt-in). The opt-in cap is NON-monotonic +// in SM version (sm80=163KB, sm86/sm89=99KB, sm90=227KB), so no clean SM-version skip guard exists. +// head_size=64 needs only a few KB of static smem (no opt-in path at all) and runs on every arch, so +// the env-var mechanism is both build-config- AND architecture-independent. Do not reintroduce +// head_size>256 MEA tests without a runtime cudaDevAttrMaxSharedMemoryPerBlockOptin capacity skip +// (tracking: microsoft/onnxruntime#28388). +// +// For every test below SKIP_IF_MEA_NOT_COMPILED() applies: if MEA is compiled out +// (onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION=OFF) dispatch falls to the correct unfused path and a +// value assertion alone could not prove the MEA offset guard ran -- so we SKIP rather than false-pass. +// +// Negative offset = nonpad_kv_seqlen(2) - q_sequence_length(4) = -2 reproduces the bug: pre-fix the +// uint32_t causal_diagonal_offset wraps, row 2 attends keys 0 AND 1 -> mean(V0,V1)=2 instead of V0=1. +// Post-fix the int32_t offset keeps -2 signed, row 2 -> V0=1. Expected frontier per query row: +// {0, 0, 1, 2} (rows 0,1 structurally empty -> 0; row2 -> V0=1; row3 -> mean(V0,V1)=2). + +// SKIP (not pass) when MEA is not compiled: without it dispatch falls to the unfused path, which +// computes the same frontier correctly, so the value assertions below cannot prove the MEA negative +// causal_diagonal_offset guard ran. Defined here, #undef'd after the last MEA test in this block. +#if USE_MEMORY_EFFICIENT_ATTENTION +#define SKIP_IF_MEA_NOT_COMPILED() ((void)0) +#else +#define SKIP_IF_MEA_NOT_COMPILED() \ + GTEST_SKIP() << "Memory Efficient Attention not compiled (onnxruntime_USE_MEMORY_EFFICIENT_ATTENTION=OFF)" +#endif + +// Negative offset (nonpad_kv_seqlen=2 < q_sequence_length=4 -> offset = -2, the wrap-bug regime). +// head_size=64 + ORT_DISABLE_FLASH_ATTENTION=1 forces RunMemoryEfficientAttention in every build and +// on every arch (see block header). The padding V values 7,9 are distinct nonzero "poison": if a +// boundary row over-attends into the padding (the pre-fix wrap bug) the output shifts off {0,0,1,2}. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_NegOffset_ForceFlashDisabled_FP16_CUDA) { + SKIP_IF_MEA_NOT_COMPILED(); + if (!HasCudaEnvironment(530)) { + return; // MEA fp16 requires SM 5.3+ (memory_efficient_attention.h:has_memory_efficient_attention) + } + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}}; + + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + const float v_row_values[4] = {1.0f, 3.0f, 7.0f, 9.0f}; + std::vector v; + v.reserve(4 * kHeadSize); + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + // Rows 0,1 structurally empty -> 0. Row 2 -> V[0]=1. Row 3 -> mean(V[0],V[1])=2. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {0.0f, 0.0f, 1.0f, 2.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// BFloat16 counterpart of the negative-offset force-MEA test above, completing the {FP16, BF16} matrix. +// BFloat16 MEA requires SM 8.0+. Same negative-offset frontier {0,0,1,2}; exercises the same int32_t +// causal_diagonal_offset guard. Tolerance is 0.03 here vs 0.02 for fp16 because bf16's narrower 7-bit +// mantissa (vs fp16's 10-bit) gives a coarser representation of the averaged outputs. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_NegOffset_ForceFlashDisabled_BF16_CUDA) { + SKIP_IF_MEA_NOT_COMPILED(); + if (!HasCudaEnvironment(800)) { + return; // BFloat16 MEA requires SM 8.0+ + } + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}}; + + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + const float v_row_values[4] = {1.0f, 3.0f, 7.0f, 9.0f}; + std::vector v; + v.reserve(4 * kHeadSize); + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, FloatsToBFloat16s(q)); + test.AddInput("K", k_shape, FloatsToBFloat16s(k)); + test.AddInput("V", v_shape, FloatsToBFloat16s(v)); + test.AddOptionalInputEdge(); // attn_mask (none) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {2}); + + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {0.0f, 0.0f, 1.0f, 2.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, FloatsToBFloat16s(expected_y_f), false, 0, 0.03f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Defense-in-depth: zero offset (nonpad_kv_seqlen == q_sequence_length -> offset = 0). Asserts the +// int32_t change leaves offset >= 0 paths bit-identical (no value ever wrapped) and did not regress the +// standard bottom-right == top-left causal case through the MEA path. No structurally empty rows: every +// query attends a non-empty key prefix. head_size=64 + ORT_DISABLE_FLASH_ATTENTION=1 forces MEA. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_ZeroOffset_ForceFlashDisabled_FP16_CUDA) { + SKIP_IF_MEA_NOT_COMPILED(); + if (!HasCudaEnvironment(530)) { + return; // MEA fp16 requires SM 5.3+ + } + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}}; + + constexpr int kHeadSize = 64; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, 4, kHeadSize}; + std::vector v_shape = {1, 1, 4, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * 4 * kHeadSize, 1.0f); + // All 4 keys are non-pad (nonpad == q_seq). V rows = {1,3,5,7}. + const float v_row_values[4] = {1.0f, 3.0f, 5.0f, 7.0f}; + std::vector v; + v.reserve(4 * kHeadSize); + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {4}); + + // offset = 4 - 4 = 0. Row i attends keys 0..i: row0 -> V0=1; row1 -> mean(V0,V1)=2; + // row2 -> mean(V0,V1,V2)=3; row3 -> mean(V0..V3)=4. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {1.0f, 2.0f, 3.0f, 4.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Defense-in-depth: positive offset (nonpad_kv_seqlen > q_sequence_length -> offset > 0, the +// external-KV-cache-longer-than-query frontier). offset = nonpad(6) - q_seq(4) = +2. Confirms the +// int32_t change leaves the positive-offset bottom-right alignment unchanged through MEA. head_size=64 +// + ORT_DISABLE_FLASH_ATTENTION=1 forces MEA. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_PositiveOffset_ForceFlashDisabled_FP16_CUDA) { + SKIP_IF_MEA_NOT_COMPILED(); + if (!HasCudaEnvironment(530)) { + return; // MEA fp16 requires SM 5.3+ + } + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}}; + + constexpr int kHeadSize = 64; + constexpr int kKvSeqLen = 6; // all 6 KV positions non-pad; offset = 6 - 4 = +2 + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {1, 1, 4, kHeadSize}; + std::vector k_shape = {1, 1, kKvSeqLen, kHeadSize}; + std::vector v_shape = {1, 1, kKvSeqLen, kHeadSize}; + + std::vector q(1 * 1 * 4 * kHeadSize, 1.0f); + std::vector k(1 * 1 * kKvSeqLen * kHeadSize, 1.0f); + // V rows = {1,3,5,7,9,11}. + const float v_row_values[kKvSeqLen] = {1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f}; + std::vector v; + v.reserve(kKvSeqLen * kHeadSize); + for (int s = 0; s < kKvSeqLen; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {1}, {kKvSeqLen}); + + // Row i attends keys 0..(i+2): row0 -> mean(V0,V1,V2)=3; row1 -> mean(V0..V3)=4; + // row2 -> mean(V0..V4)=5; row3 -> mean(V0..V5)=6. + std::vector expected_y_f; + expected_y_f.reserve(4 * kHeadSize); + const float expected_row[4] = {3.0f, 4.0f, 5.0f, 6.0f}; + for (int s = 0; s < 4; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_row[s]); + test.AddOutput("Y", {1, 1, 4, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +// Per-batch coverage of the production RightPaddingBatchHook: batch=2 with mixed nonpad_kv_seqlen +// {2, 6} and a shared kv_seq=6. Each batch entry computes its OWN causal_diagonal_offset = +// nonpad_kv_seqlen[b] - q_sequence_length, so batch 0 gets a NEGATIVE offset (2 - 4 = -2, the wrap +// bug regime) while batch 1 gets a POSITIVE offset (6 - 4 = +2). This exercises the per-batch offset +// path that the batch=1 tests cannot reach. head_size=64 + ORT_DISABLE_FLASH_ATTENTION=1 forces MEA. +// Batch 0 frontier {0, 0, 1, 2} (rows 0,1 structurally empty; row2 -> V0; row3 -> mean(V0,V1)); +// keys 2..5 are padding (poison V = 7,9,13,15) and must never leak. Batch 1 has all 6 keys valid: +// row i attends keys 0..(i+2) -> {mean(V0..V2), mean(V0..V3), mean(V0..V4), mean(V0..V5)} = {3,4,5,6}. +TEST(AttentionTest, Attention_Causal_NonPadKVSeqLen_MEA_MixedBatchOffsets_ForceFlashDisabled_FP16_CUDA) { + SKIP_IF_MEA_NOT_COMPILED(); + if (!HasCudaEnvironment(530)) { + return; // MEA fp16 requires SM 5.3+ + } + ScopedEnvironmentVariables scoped_env_vars{ + EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}}; + + constexpr int kHeadSize = 64; + constexpr int kBatch = 2; + constexpr int kQSeq = 4; + constexpr int kKvSeq = 6; + OpTester test("Attention", 24, onnxruntime::kOnnxDomain); + test.AddAttribute("is_causal", static_cast(1)); + + std::vector q_shape = {kBatch, 1, kQSeq, kHeadSize}; + std::vector k_shape = {kBatch, 1, kKvSeq, kHeadSize}; + std::vector v_shape = {kBatch, 1, kKvSeq, kHeadSize}; + + std::vector q(kBatch * kQSeq * kHeadSize, 1.0f); + std::vector k(kBatch * kKvSeq * kHeadSize, 1.0f); + // Batch 0: V rows {1,3,7,9,13,15} -- only rows 0,1 are non-pad (nonpad=2); 7,9,13,15 are poison. + // Batch 1: V rows {1,3,5,7,9,11} -- all 6 non-pad (nonpad=6). + const float v_row_values[kBatch][kKvSeq] = {{1.0f, 3.0f, 7.0f, 9.0f, 13.0f, 15.0f}, + {1.0f, 3.0f, 5.0f, 7.0f, 9.0f, 11.0f}}; + std::vector v; + v.reserve(kBatch * kKvSeq * kHeadSize); + for (int b = 0; b < kBatch; ++b) + for (int s = 0; s < kKvSeq; ++s) + for (int h = 0; h < kHeadSize; ++h) v.push_back(v_row_values[b][s]); + + test.AddInput("Q", q_shape, ToFloat16(q)); + test.AddInput("K", k_shape, ToFloat16(k)); + test.AddInput("V", v_shape, ToFloat16(v)); + test.AddOptionalInputEdge(); // attn_mask (none) + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // past_value + test.AddInput("nonpad_kv_seqlen", {kBatch}, {2, 6}); + + // Batch 0 (offset -2): {0,0,1,2}. Batch 1 (offset +2): {3,4,5,6}. + const float expected_rows[kBatch][kQSeq] = {{0.0f, 0.0f, 1.0f, 2.0f}, {3.0f, 4.0f, 5.0f, 6.0f}}; + std::vector expected_y_f; + expected_y_f.reserve(kBatch * kQSeq * kHeadSize); + for (int b = 0; b < kBatch; ++b) + for (int s = 0; s < kQSeq; ++s) + for (int h = 0; h < kHeadSize; ++h) expected_y_f.push_back(expected_rows[b][s]); + test.AddOutput("Y", {kBatch, 1, kQSeq, kHeadSize}, ToFloat16(expected_y_f), false, 0, 0.02f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + +#undef SKIP_IF_MEA_NOT_COMPILED + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py index 6b3f6d1c3ff34..eb2c5eb2dc599 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/test_tensorscatter_attention.py @@ -981,70 +981,46 @@ def test_nonpad_with_bool_mask_cuda_fp16( numpy.testing.assert_allclose(present_v, ref_present_v, rtol=rtol["fp16"], atol=atol["fp16"]) -class TestCausalTensorScatterRejected(unittest.TestCase): - """Test that is_causal=1 + TensorScatter decode (S_q != S_kv, no past) is rejected. - - Per ONNX spec, is_causal without past_key means upper-left alignment: q[i] attends - only to kv[0..i]. For decode with external cache (S_q=1, S_kv=cache_size), this means - q[0] sees only kv[0] — not meaningful for autoregressive generation. - - The dispatch guard should return NOT_IMPLEMENTED for this combination. - Models should use is_causal=0 for TensorScatter decode. +class TestCausalTensorScatterBottomRight(unittest.TestCase): + """Test that is_causal=1 + TensorScatter decode (S_q != S_kv, no past) is SUPPORTED. + + Per onnx/onnx#8068, is_causal with an external KV cache (nonpad_kv_seqlen) and no + past_key uses BOTTOM-RIGHT alignment: query in-block index i attends key j iff + j <= i + offset[b], where offset[b] = nonpad_kv_seqlen[b] - S_q. For decode + (S_q=1, nonpad=5) the offset is 4, so the single query row attends keys 0..4 — all + valid cache positions — a meaningful, correct decode result. + + This combination previously returned NOT_IMPLEMENTED under the pre-#8068 upper-left + assumption (where q[0] would have seen only kv[0]); onnxruntime#28904 removed that + dispatch guard and now computes the bottom-right frontier via Flash (seqlens_k) or the + CUTLASS memory-efficient fallback (causal_diagonal_offset = num_keys - num_queries). + The is_causal + nonpad_kv_seqlen + past_key combination remains rejected upstream + (ORT_ENFORCE in attention_helper.h). Deeper S_q>1 / nonpad0 (onnx#8068). The "^" prefix also matches the _expanded (function-body) and EP-suffixed variants. + "^test_attention_4d_gqa_causal_nonpad_decode", // #28904 bottom-right decode + "^test_attention_4d_gqa_causal_nonpad_decode_fp16", // #28904 bottom-right decode (fp16) + "^test_attention_4d_causal_nonpad_continued_prefill", // #28904 bottom-right continued-prefill + "^test_attention_causal_boolmask_nan_robustness", // #28904 fully-masked-row -> 0 (Bug-2, opset-24) + "^test_attention_23_boolmask_fullymasked_row_nan_robustness", // #28904 fully-masked-row -> 0 (Bug-2, opset-23 errata) + // #28994: preemptive skip — onnx pinned at v1.21.0 lacks these #8068 mode-3 tests; the + // regex is a no-op today and prevents CI breakage the moment the onnx pin is bumped past + // #8068. CUDA returns NOT_IMPLEMENTED for qk_matmul_output modes beyond kQK + // (cuda/llm/attention.cc:1477). De-skip when mode>kQK is supported (#27712). See #28994 Section A. + // Matches: test_attention_{23,24}_fullymasked_qk_matmul_output_mode3_zero[_expanded] and + // test_attention_24_qk_matmul_output_mode3_softmax_precision[_expanded]. + "^test_attention_(23|24)_.*qk_matmul_output_mode3", // #28994 CUDA NOT_IMPLEMENTED mode>kQK (attention.cc:1477); #27712 to support // TODO: support qk_matmul_output modes beyond kQK in Attention-cuda (see issue #27712) // Tests combining qk_matmul with softcap need unfused-path qk_matmul support (deferred). "^test_attention_3d_with_past_and_present_qk_matmul_softcap_cuda", // qk_matmul modes beyond kQK not supported