onnxruntime
058ce95b - Fix CPU Attention softcap/attn_mask ordering (onnx#7867, #7913) + consolidate CUDA spec coverage tests from #28371 (#28379)

Commit
83 days ago
Fix CPU Attention softcap/attn_mask ordering (onnx#7867, #7913) + consolidate CUDA spec coverage tests from #28371 (#28379) Fix CPU Attention softcap/attn_mask ordering to match ONNX v24 spec (#7867, #7913) + consolidate CUDA spec coverage tests (#28371) Branch: `fix-cpu-attention-softcap-mask-ordering` (off `microsoft/onnxruntime` `main`) Tip: `bfe33d3d08` > **Supersedes #28370 and #28371.** > > - **#28370** diagnosed the symptom correctly but inverted the fix: it modified the (already spec-correct) CUDA path to match the (buggy) CPU path. This PR takes the opposite — and correct — approach, aligning CPU with both CUDA and the ONNX v24 spec as established by onnx/onnx#7867 and onnx/onnx#7913. > - **#28371** ("Tighten CUDA Attention MEA eligibility on head_size%4 + add GQA test coverage") has been **fully consolidated into this PR** at the user's request. The CUDA test coverage and SKILL.md cleanup that lived on #28371 share a `qk_matmul_output_mode` numbering and softcap/mask-ordering spec with the CPU fix here, so keeping them on separate PRs created merge-ordering and mental-model hazards. Rolling them up presents one coherent fix-tests-docs reviewer story. > > We recommend closing both #28370 and #28371 in favor of this one. --- ## Scope 1. **CPU spec-correct softcap+mask ordering fix** (per onnx/onnx#7867) — the headline bug. The CPU `attention.cc` was applying `attn_mask` *before* `softcap` (which would clip a `-inf` masked logit to a finite value, leaking the masked V slot through softmax). Now applies `scale → softcap → +attn_mask → softmax` in spec order. 2. **`qk_matmul_output_mode` enum value swap** (per onnx/onnx#7913). Breaking change to mode `1` ↔ `2` semantics + matching C++ enumerator rename. CPU op fully implements the new ordering; CUDA op (Flash/MEA/Unfused dispatch) confirmed unchanged in behaviour. 3. **Tianlei blocker tests** (CPU): mode-1+softcap differentiating, softcap+nonpad_kv_seqlen no-leakage. Both failed before this PR's fix. 4. **Comprehensive CUDA Attention spec coverage** (consolidated from #28371): - `TestONNXAttentionGQAAsymmetricHeadSize` (REG) — guards the silent-broken-output regression on GQA + asymmetric Q/V head sizes (microsoft/onnxruntime#28357, fixed by #28358). - `TestONNXAttentionGQAOutputQK` (1c) — pins GQA + raw-QK output (`qk_matmul_output_mode == kQK`) on the unfused fp16 path. - `TestONNXAttentionGQASoftcapFloat32` (1e baseline) + `TestONNXAttentionGQASoftcapFloat32MaskOrdering` (1e ordering) — pins fp32 + softcap on the unfused GQA path, including the post-#7867 softcap+mask ordering using a poison-V leakage detector. 5. **SKILL.md updates** (`.agents/skills/cuda-attention-kernel-patterns/SKILL.md`): - §4 (Softcap Ordering) revert — undoes #28370's inverted statement. - §1 (MEA eligibility paragraph) cleanup — `head_size%8` is enforced by `has_memory_efficient_attention`; `head_size == v_head_size` is required for GQA (LaunchUngroup) in addition to decode (LaunchConcatNewToPastKV). Reflects the post-#28358 host-side gate. 6. **ONNX node-test skip lists** (transient, until ONNX v1.22+ is bundled): - C++ runner: `onnxruntime/test/onnx/TestCase.cc::GetBrokenTests()` — 20 attention entries. - Python runner: `onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc::current_failing_tests` — 7 CPU attention entries. Both block-comments cite onnx/onnx#7867 + onnx/onnx#7913 and the v1.22+ unblock criterion: the fixtures bundled in v1.21.0 were generated *before* the spec PRs landed and therefore encode the now-incorrect ordering. Once `cmake/external/onnx` is bumped to ≥1.22, both skip blocks should be removed in one commit. --- ## Cross-references - Issue: microsoft/onnxruntime#28351 (CPU spec coverage gaps), microsoft/onnxruntime#28357 (asymmetric-V regression). - Closed/replaced: #28370 (inverted fix), #28371 (consolidated here). - Coordinated with: #28365 (qk_matmul_output_mode rename across consumers). - Spec PRs: onnx/onnx#7867, onnx/onnx#7913 (supersedes the now-closed onnx/onnx#7865 issue). --- ## Verification - `AttentionTest.*` — 60/60 pass (gtest, CPU + CUDA). - `TestONNXAttentionCPUSoftcapMaskOrdering` — 4/4 pass (CPU spec fix + Tianlei blockers). - `TestONNXAttentionGQAAsymmetricHeadSize` — 2/2 pass (fp16 + bf16, H100). - `TestONNXAttentionGQAOutputQK` — 1/1 pass (fp16, H100). - `TestONNXAttentionGQASoftcapFloat32` + `…MaskOrdering` — 4/4 pass (fp32, H100). - ONNX node tests: previously-failing 20 (C++) + 7 (Python) attention fixtures now SKIPPED with citation, not failed. - `lintrunner` clean on all touched files. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading