Guard CUDA LayerNorm/RMSNorm int32 offset range - #31650
Merged
Akshay Sonawane (apsonawane) merged 4 commits intoAug 5, 2026
Merged
Conversation
Reject inputs where num_rows * norm_size exceeds INT_MAX before launching kernels that use int32 row offsets. This prevents overflow in per-row pointer offset arithmetic in LayerNormalization and RMSNormalization CUDA paths.
Copilot started reviewing on behalf of
Akshay Sonawane (apsonawane)
August 4, 2026 23:04
View session
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds input-size validation to the CUDA implementations of LayerNormalization and RMSNormalization to prevent int32 overflow in CUDA kernel indexing (notably around num_rows * norm_size).
Changes:
- Added a pre-launch guard in
LayerNorm::ComputeInternalto reject inputs wherenum_rows * norm_sizeexceedsINT_MAX. - Added a similar guard in
RMSNorm::ComputeInternal. - Included
<limits>to support the new bounds checks.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/core/providers/cuda/nn/layer_norm.cc | Adds an INT_MAX-based size guard before launching the CUDA LayerNorm kernel. |
| onnxruntime/core/providers/cuda/nn/rms_norm.cc | Adds the same INT_MAX-based size guard for CUDA RMSNorm (LayerNorm-based) execution. |
…tion This commit addresses the security vulnerability identified in PR #31650 review comments from Copilot reviewer regarding integer overflow in host-side CUDA kernel indexing calculations. SECURITY FIX (Copilot Comments on LayerNorm and RMSNorm): - Enhanced validation to prevent integer overflow in HostApplyLayerNorm calculations - Original validation allowed norm_size == INT_MAX when num_rows == 1 - This caused overflow in host arithmetic: (n2 + 4 * warp_size - 1) - New validation uses two-pronged approach: * Checks that norm_size doesn't exceed INT_MAX - 256 (accounts for warp_size overhead) * Still checks that num_rows * norm_size doesn't exceed INT_MAX * Either check failing rejects the input DETAILED CHANGES: LayerNormalization (layer_norm.cc): - Added constexpr MAX_WARP_FACTOR = 256 (conservative upper bound for 4 * warp_size) - Compute MAX_NORM_SIZE = INT_MAX - MAX_WARP_FACTOR - Validate: norm_size <= MAX_NORM_SIZE AND norm_size <= INT_MAX / num_rows - Prevents: (n2 + 4 * warp_size - 1) overflow for n2 = norm_size - Prevents: num_rows * norm_size > INT_MAX overflow RMSNormalization (rms_norm.cc): - Applied identical validation as LayerNormalization - Both use the same HostApplyLayerNorm function, same vulnerability pattern - Ensures consistency across both normalization operators BENEFITS: ✅ Prevents undefined behavior from integer overflow ✅ Blocks exploitation attempts using extreme input sizes ✅ Maintains safety margin for host-side arithmetic calculations ✅ Conservative bounds ensure correctness Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Akshay Sonawane (apsonawane)
enabled auto-merge (squash)
August 5, 2026 16:50
Tianlei Wu (tianleiwu)
approved these changes
Aug 5, 2026
Akshay Sonawane (apsonawane)
deleted the
msrc/cuda-layernorm-offset-overflow-fix
branch
August 5, 2026 23:02
Tianlei Wu (tianleiwu)
pushed a commit
that referenced
this pull request
Aug 6, 2026
This pull request adds input validation checks to prevent integer overflow issues during CUDA kernel indexing in the LayerNorm and RMSNorm CUDA operators. The main goal is to ensure that the product of `num_rows` and `norm_size` does not exceed `INT_MAX`, which could lead to incorrect behavior or crashes. Input validation for CUDA kernel indexing: * Added a check in `LayerNorm::ComputeInternal` (in `layer_norm.cc`) to return an error if `num_rows * norm_size` exceeds `INT_MAX`, preventing integer overflow during CUDA kernel indexing. * Added a similar check in `RMSNorm::ComputeInternal` (in `rms_norm.cc`) to ensure the input size does not exceed CUDA kernel indexing limits. Code maintenance: * Included the `<limits>` header in both `layer_norm.cc` and `rms_norm.cc` to support the new input validation logic. [[1]](diffhunk://#diff-ebda3d3b7054f5d14c679ebe8e6520a2c84a5a558d8fdcc0798c08e7032345feR9) [[2]](diffhunk://#diff-adebea99f15800767eb7b84d40be4873e81ca4df99ffe9dc7f849eabe0a3c563R9) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> (cherry picked from commit 24f1aed)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request adds input validation checks to prevent integer overflow issues during CUDA kernel indexing in the LayerNorm and RMSNorm CUDA operators. The main goal is to ensure that the product of
num_rowsandnorm_sizedoes not exceedINT_MAX, which could lead to incorrect behavior or crashes.Input validation for CUDA kernel indexing:
LayerNorm::ComputeInternal(inlayer_norm.cc) to return an error ifnum_rows * norm_sizeexceedsINT_MAX, preventing integer overflow during CUDA kernel indexing.RMSNorm::ComputeInternal(inrms_norm.cc) to ensure the input size does not exceed CUDA kernel indexing limits.Code maintenance:
<limits>header in bothlayer_norm.ccandrms_norm.ccto support the new input validation logic. [1] [2]