Expanded Attention reference decomposition: SizeFromDimension SIZE_MAX underflow at Add node (surfaced by ONNX 1.22 integration)

[titaiwangms]

Summary

The ONNX function-expanded reference decomposition of the Attention op fails at runtime with a TensorShape::SizeFromDimension underflow (a shape dimension computed as -1, surfaced as the unsigned value 18446744073709551615 = SIZE_MAX). It is currently masked by a skip added in PR #28754 and should be removed once the underlying bug is fixed.

This was unmasked by the ONNX 1.22 / opset-27 integration. Note: the ONNX-domain Attention op itself (opset 23/24) was not changed by the opset-27 bump — the bump simply enabled running the expanded reference for these node tests.

Failing tests

• test_attention_4d_softcap_neginf_mask_expanded_cpu
• test_attention_4d_softcap_neginf_mask_poison_expanded_cpu

Both are ONNX backend node-model tests run via onnx_backend_test_series.py on the CPU EP.

Observed error

Non-zero status code returned while running Add node. Name:''
onnxruntime/core/framework/tensor_shape.cc:77
  int64_t onnxruntime::TensorShape::SizeFromDimension(size_t) const
  dimension <= num_dims was false.
  Invalid dimension of 18446744073709551615 for SizeFromDimension. Tensor has 2 dimensions.

First seen on the webgpu / build-and-test (arm64, arm64, Release) leg (macOS-arm64 build, onnxruntime_USE_WEBGPU=ON), run 27243073719.

Root cause (current understanding)

• The failure is in the EXPANDED (ONNX function-decomposed reference) Attention graph, not the native fused Attention kernel.
• A shape dimension is computed as -1 somewhere downstream of the bias Add in the expanded decomposition and then passed into SizeFromDimension, where the unsigned underflow trips the dimension <= num_dims enforce in core/framework/tensor_shape.cc:73.
• The _poison variant (identical graph, -inf mask values swapped in) fails identically, so the -inf/softcap mask values are not the trigger — it is the shape flow of the expanded decomposition. Prime suspect: the bias-Add → softmax shape handling path.
• Only surfaced after the ONNX 1.22 / opset-27 bump (PR Integrate ONNX 1.22.0 (opset 27) — issue #28752 #28754) enabled the expanded Attention reference for these tests.

Why it is NOT user-facing today

• The native fused Attention kernel passes on every architecture (test_attention_4d_softcap_neginf_mask and _poison, without _expanded).
• The CPU EP passes these expanded tests on x64 and Linux-arm64, and the ONNX ReferenceEvaluator passes them — the failure only reproduces on the macOS-arm64-webgpu build.
• Production inference uses the fused kernel; the expanded decomposition is a conformance-reference artifact. So there is no known user-facing regression — but the underflow is a genuine latent correctness/robustness bug worth root-causing.

Current mitigation (to be removed when fixed)

Skipped via the global current_failing_tests list in onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc (lines 43–44), added in commit 027820e7e6 on PR #28754:

"^test_attention_4d_softcap_neginf_mask_expanded*",
"^test_attention_4d_softcap_neginf_mask_poison_expanded*",

Action when fixed: remove these two skip entries and confirm the tests pass on macOS-arm64-webgpu.

Repro

1. Build with --use_webgpu on macOS-arm64 Release (or reproduce the expanded reference graph on the affected build).
2. Run onnx_backend_test_series.py filtering to test_attention_4d_softcap_neginf_mask_expanded_cpu.
3. Observe the SizeFromDimension SIZE_MAX error at the Add node.

References

• PR: Integrate ONNX 1.22.0 (opset 27) — issue #28752 #28754 (ONNX 1.22.0rc2 / opset-27 integration)
• Skip commit: 027820e7e6
• Enforce site: onnxruntime/core/framework/tensor_shape.cc:73 (SizeFromDimension)
• CI: run 27243073719, job 80450998304 (webgpu / build-and-test (arm64, arm64, Release))

[titaiwangms]

Investigation pointers (from critical review)

A critical review flagged that this is likely not a platform-independent CPU softmax.cc:177 SizeFromDimension(rank-1) issue. If it were, test_attention_4d_softcap_neginf_mask_expanded_cpu would fail on all platforms — but it passes on x64 and Linux-arm64 and only trips on macOS-arm64 (the webgpu CI leg).

Whoever fixes this should:

1. Investigate the platform/arch-specific shape-flow path that produces the rank-0 / dim = -1 intermediate on macOS-arm64 — do not assume the CPU softmax path; the trigger is arch-specific.
2. Apply a rank >= 1 guard (or equivalent) at the underflow site so the expanded reference path is robust regardless of arch.

The onnx_backend_test_series_filters.jsonc skip is only a temporary CI-unblock, not the fix — remove it once the guard lands.

[titaiwangms]

Root-cause analysis & fix pointers (for the WebGPU EP fixer)

This is a latent WebGPU EP bug, not a CPU softmax.cc issue and not an ONNX regression.

Exact site

onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc, BinaryElementwise::ComputeInternal:

• Underflow (lines 216–218):

  reshaped_lhs_shape[...] = lhs_shape.SizeFromDimension(lhs_shape.NumDimensions() - num_shared_dimension);

  When num_shared_dimension > NumDimensions(), the unsigned subtraction wraps to ~SIZE_MAX, so
  SizeFromDimension(SIZE_MAX) trips the dimension <= num_dims enforce in
  core/framework/tensor_shape.cc:75 → the observed
  Invalid dimension of 18446744073709551615 … Tensor has 2 dimensions on an Add node.

• Over-count (lines 158–168): the loop counts shared trailing dims up to the output rank and
  treats a missing (broadcast) input dim as 1:

  for (size_t i = 1; i < output_shape.NumDimensions(); i++) {
    int64_t dimA = lhs_shape.NumDimensions() >= i ? lhs_shape[lhs_shape.NumDimensions() - i] : 1;
    int64_t dimB = rhs_shape.NumDimensions() >= i ? rhs_shape[rhs_shape.NumDimensions() - i] : 1;
    if (dimA == dimB) { shared_dimension *= dimA; num_shared_dimension++; } else break;
  }

  So a broadcast operand with fewer dims than the matching trailing run (e.g. the broadcast bias
  Add in the expanded Attention graph), with the shared product divisible by 4 (the
  need_reshape gate at line 208 + the %4 gate at line 169), makes num_shared_dimension exceed
  that operand's own rank.

Proposed fix (small, low-risk)

Bound the count loop so a phantom broadcast 1-tail is not merged:

for (size_t i = 1; i < output_shape.NumDimensions(); i++) {
  if (lhs_shape.NumDimensions() < i || rhs_shape.NumDimensions() < i) break;  // add
  ...
}

(Equivalent: clamp num_shared_dimension to min(lhs_rank, rhs_rank) before line 216.) Worst case
falls back to the non-merged broadcast path — correct results, just less vectorized. Add a WebGPU
binary-elementwise broadcast unit test for this shape pattern.

Why it only shows on the webgpu CI leg (not arch-specific)

The Python ONNX backend builds the session from get_available_providers()
(onnxruntime/python/backend/backend.py:148-151). On the macOS-arm64 webgpu wheel that list is
[WebGPUExecutionProvider, CPUExecutionProvider], so the broadcast Add runs on the WebGPU EP.
The x64 / Linux-arm64 legs are CPU-only builds → the Add runs on CPU (no bug) → they pass.

Provenance (new-in-1.22 test, not an rc1→rc2 change)

The triggering tests (test_attention_4d_softcap_neginf_mask[_poison][_expanded]) were added by
upstream onnx/onnx#7867 ("Fix Attention op softcap ordering", merged 2026-04-30), after ONNX
v1.21.0 (2026-03-27) and before 1.22.0rc1 (2026-06-01). The Attention op + function expansion came
via onnx/onnx#6501. Nothing in the rc1→rc2 window changed the Attention op/function/reference. The
bug was simply newly exercised once the ORT test stage started reaching these expanded references.

After the fix

Remove the two skips in onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc (lines
43–44, added in PR #28754 commit 027820e7e6) and confirm the expanded-Attention tests pass on the
webgpu leg.

[titaiwangms]

Upstream provenance (which ONNX change revealed this)

This is not an ONNX regression — the WebGPU underflow is latent/pre-existing in ORT. The "reveal"
is purely that a new backend test started running:

• Tests added by Fix Attention op softcap ordering: apply before mask/bias onnx/onnx#7867 — "Fix Attention op softcap ordering: apply before mask/bias"
  (commit fe30560bdf, merged 2026-04-30). It adds export_attention_softcap_with_neginf_mask[_poison]()
  and the test_attention_4d_softcap_neginf_mask*[_expanded] data dirs.
  Fix Attention op softcap ordering: apply before mask/bias onnx/onnx#7867
• Attention op + function-body expansion originally from Add Attention Op to ONNX Opset 23 onnx/onnx#6501 (opset 23).
  Add Attention Op to ONNX Opset 23 onnx/onnx#6501
• NEW-in-1.22, not an rc1→rc2 change. ONNX v1.21.0 (2026-03-27) has no neginf test; Clean up CPU kernel definition for opset 13 Pad #7867 merged
  2026-04-30, before 1.22.0rc1 (2026-06-01). The rc1→rc2 window changed nothing in the Attention
  op/function/reference (the expanded-Attention model is byte-identical rc1↔rc2; rc2's cherrypicks are
  all Conv/Pool/checker/xcode).
• Actual unmask: an ORT-side change on PR Integrate ONNX 1.22.0 (opset 27) — issue #28752 #28754 (the Python quant suite no longer aborting before
  the OnnxBackendNodeModelTest stage), which let the backend-test stage reach these expanded
  references — at which point the webgpu wheel runs the broadcast Add through the buggy path.