Guard CUDA LayerNorm/RMSNorm int32 offset range - #31650
Open
apsonawane wants to merge 1 commit into
Open
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.
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. |
Comment on lines
+80
to
+82
| ORT_RETURN_IF(params.num_rows > 0 && | ||
| params.norm_size > std::numeric_limits<int>::max() / params.num_rows, | ||
| "RMSNormalization input is too large for CUDA kernel indexing: num_rows * norm_size exceeds INT_MAX"); |
Comment on lines
+92
to
+95
| ORT_RETURN_IF(params.num_rows > 0 && | ||
| params.norm_size > std::numeric_limits<int>::max() / params.num_rows, | ||
| "LayerNormalization input is too large for CUDA kernel indexing: num_rows * norm_size exceeds INT_MAX"); | ||
|
|
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]