EmbeddingGemma-300M fp16 CPU inference regression: NaN output on Linux aarch64 since 1.24.1

Describe the bug

Starting in 1.24.1 (the release immediately after 1.23.2), running EmbeddingGemma-300M's public fp16 ONNX export (onnx-community/embeddinggemma-300m-ONNX, onnx/model_fp16.onnx) through the CPU execution provider on Linux aarch64 produces all-NaN output for batches that ran fine on 1.23.2. Confirmed via a full release bisection from 1.19.0 through 1.28.0: 1.23.2 is the last good version, 1.24.1 is the first bad one, and it is still present as of the latest release, 1.30.0. It does not reproduce on macOS aarch64 (Apple Silicon) with the same onnxruntime version and inputs.

Root cause (via graph surgery)

I exposed every intermediate tensor in the graph as a model output and ran inference on the real model with real inputs directly on the affected Linux aarch64 node. The model's residual stream grows roughly exponentially across its 24 transformer layers — from a magnitude of ~900 at layer 0 to ~50,000–65,000 by layer 22. fp16's max representable value is 65504. At /model/layers.22/Add_2 (the residual add after the feed-forward block), two already-huge fp16 tensors sum to just over that ceiling and overflow to +Inf, which propagates to NaN by layer 23.

This overflow is inherent to the fp16 export in general — for sequences padded near the tokenizer's 512-token max, even 1.23.2 can overflow. What changed between 1.23.2 and 1.24.1 is that some CPU-EP kernel now produces a sub-1% numeric difference that's enough to flip already-marginal batches (comfortably under 65504 on 1.23.2) to just over it.

To Reproduce

Self-contained repro using synthetic filler text (no real user data — chosen purely to match the character-length distribution that triggers this in our production traffic):

import json
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
import numpy as np
import onnxruntime as ort

REPO = "onnx-community/embeddinggemma-300m-ONNX"
model_path = hf_hub_download(REPO, "onnx/model_fp16.onnx", subfolder=None)
hf_hub_download(REPO, "onnx/model_fp16.onnx_data", subfolder=None)  # external data, same dir
tokenizer_path = hf_hub_download(REPO, "tokenizer.json")

LOREM = ("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor "
         "incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud "
         "exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure "
         "dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. "
         "Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt "
         "mollit anim id est laborum. ")

# Character lengths chosen to match a real failing production batch (mix of short and long,
# one long enough to require truncation to the 512-token model max).
LENGTHS = [279, 264, 261, 186, 87, 167, 297, 644, 926, 77, 625, 1164, 4621, 198, 358, 873,
           1670, 123, 260, 124]
texts = [(LOREM * (n // len(LOREM) + 2))[:n] for n in LENGTHS]

tokenizer = Tokenizer.from_file(tokenizer_path)
tokenizer.enable_padding(pad_id=0, pad_token="<pad>")
tokenizer.enable_truncation(max_length=512)
encs = tokenizer.encode_batch(texts)
input_ids = np.array([e.ids for e in encs], dtype=np.int64)
attention_mask = np.array([e.attention_mask for e in encs], dtype=np.int64)

session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
out = session.run(["sentence_embedding"], {"input_ids": input_ids, "attention_mask": attention_mask})[0]

nan_rows = sum(1 for row in out if np.isnan(row).any())
print(f"onnxruntime {ort.__version__}: {nan_rows}/{len(out)} rows contain NaN")

Expected behavior

nan_rows should be 0/20 regardless of onnxruntime version (as it is on 1.23.2).

Actual behavior

• onnxruntime==1.23.2 → 0/20 (clean)
• onnxruntime==1.24.1 through 1.30.0 → 20/20 (all rows NaN)

Verified with the pip onnxruntime package on Linux aarch64 (AWS Graviton, Amazon Linux 2023, Python 3.11) and independently via the Java Maven artifact (com.microsoft.onnxruntime:onnxruntime) on the same hardware — same result in both bindings.

Additional context

• Setting GraphOptimizationLevel (all 5 levels, ORT_DISABLE_ALL through ORT_ENABLE_ALL) and ExecutionMode (SEQUENTIAL/PARALLEL) in every combination does not change the outcome — the raw, unoptimized graph (ORT_DISABLE_ALL) overflows identically to the fully optimized one, which rules out a graph-fusion pass as the cause and points to the base kernel arithmetic itself (most likely in an op used heavily in this architecture — RMSNorm, RoPE/rotary embedding, or the residual Add/MatMul chain).
• Not reproducible on macOS aarch64 (Apple Silicon) with the identical onnxruntime version, model file (verified by MD5), and inputs — suggesting the regression is specific to the Linux aarch64 build's kernel implementation, not a universal numeric change.
• Happy to share the full per-layer intermediate-tensor trace (tensor names, shapes, min/max values at every layer) if useful for narrowing down the specific op.

System information

• OS: Amazon Linux 2023, Linux aarch64 (AWS Graviton)
• ONNX Runtime version: 1.24.1 through 1.30.0 (regression); 1.23.2 and earlier (working)
• Python version: 3.11 (pip package); also reproduced via Java bindings on the same hardware
• Model: onnx-community/embeddinggemma-300m-ONNX, onnx/model_fp16.onnx, revision 5090578d9565bb06545b4552f76e6bc2c93e4a66
• Execution provider: CPUExecutionProvider