Skip to content

Fix out-of-bounds read in EmbedLayerNormalization shape inference for scalar beta - #32610

Merged
kunal-vaishnavi merged 1 commit into
microsoft:mainfrom
wangw-1991:Fix_ShapeInference_OOB
Sep 17, 2026
Merged

kunal-vaishnavi merged 1 commit into
microsoft:mainfrom
wangw-1991:Fix_ShapeInference_OOB

Conversation

@wangw-1991

Copy link
Copy Markdown
Contributor

Description

EmbedLayerNormalizationShapeInference validated the beta input's rank using the wrong shape: beta_dims was aliased to gamma_shape.dim() instead of beta_shape.dim(). As a result, a rank-1 gamma would let the rank check pass
even when beta is a 0-dim scalar, and execution reached beta_shape.dim(0).dim_value() — an unchecked protobuf repeated-field access on an empty dim list, causing an out-of-bounds read during session load, before provider/kernel assignment.

Fix

  • Alias beta_dims to beta_shape.dim() so the rank check guards the correct input; a scalar beta is now rejected via fail_shape_inference instead of reaching the OOB access. This matches the adjacent gamma validation.
  • Correct the gamma error message ("2 dimension" → "1 dimension") to reflect the actual check

Copilot AI balanced review requested due to automatic review settings September 15, 2026 11:13
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

A scalar-beta regression test is needed to protect the memory-safety fix.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes an out-of-bounds shape-inference read caused by validating beta against gamma dimensions.

Changes:

  • Validate beta using its own shape.
  • Correct the expected gamma rank in the error message.
File summaries
File Description
onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc Corrects gamma/beta shape validation.
Review details
  • Files reviewed: 1/1 changed files
  • Comments generated: 1
  • Review effort level: Balanced

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment thread onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc
@wangw-1991

Copy link
Copy Markdown
Contributor Author

kunal-vaishnavi Can you help review this or help find a suitable reviewer? Thanks.

@kunal-vaishnavi

Copy link
Copy Markdown
Contributor

Thanks for your contribution!

@wangw-1991

Copy link
Copy Markdown
Contributor Author

Thanks for your contribution!

Thanks for the review. Can you help merge this PR if no other issues.

@kunal-vaishnavi
kunal-vaishnavi merged commit 7054657 into microsoft:main Sep 17, 2026
90 checks passed
@wangw-1991
Wei Wang (wangw-1991) deleted the Fix_ShapeInference_OOB branch September 18, 2026 01:10
Tianlei Wu (tianleiwu) added a commit that referenced this pull request Sep 29, 2026
This cherry-picks the following commits for the release:

* #32610
* #32633
* #32641
* #32611
* #32609
* #32607
* #32679
* #32681
* #32666

CI fixes:

* #32491 (installs `uuid-dev` for the Linux minimal-build CoreML job;
that check also fails on the `rel-1.28.3` base)
* Raise the Android minimal baseline binary-size threshold from
1,440,768 to 1,443,840 bytes (release branch only). The checks added by
these fixes grow the minimal build from 1,440,538 to 1,443,098 bytes
(+2,560).

All picks except #32607 applied cleanly with `git cherry-pick -x`, in
the order they merged into main (#32666 was added afterward).

#32607 required a manual backport because `bert_defs.cc` has diverged on
main since 1.28 branched:
- `GroupQueryAttention` / `SparseAttention`: backported as-is
(`total_sequence_length_index` of 6 / 7, plus the single-element check).
The GQA call site in 1.28 has no `sliding_window_cache`, so it only gets
the new argument.
- `CausalConvWithState`: the 1.28 schema predates `channels_last`,
`state_window`, and `dilation`. Only the `ndim` in [1, 3] check and the
`weight`/`input` rank == `ndim + 2` checks are backported. These match
what the 1.28 CPU kernel already enforces.
`ChannelsLastInputRankBelowThreeIsRejected` is dropped.
- `GatedDeltaNet`: this op does not exist on 1.28.3, so its hunk and
test are dropped.

Validation (Windows x64, CPU EP, RelWithDebInfo):
- All 14 regression tests added by these PRs pass.
- `onnxruntime_test_all`: 1848 tests (including #32666's
`PluginCopiesForwardStreams`), 0 failures.
- `onnxruntime_provider_test`: 5370 tests, 0 failures.

---------

Co-authored-by: Wei Wang <wei4.wang@intel.com>
Co-authored-by: shiyi <shiyi.zou@intel.com>
Co-authored-by: Bin Miao <bin.miao@intel.com>
Co-authored-by: Bryan B <bryan.bernhart@intel.com>
Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com>
Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Xiaofei Han <xiaofeihan@microsoft.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants