fix(security): validate INT4 W4A8 MoE inputs to prevent OOB read and …#48580
Open
jperezdealgaba wants to merge 2 commits into
Open
fix(security): validate INT4 W4A8 MoE inputs to prevent OOB read and …#48580jperezdealgaba wants to merge 2 commits into
jperezdealgaba wants to merge 2 commits into
Conversation
…div-by-zero The CPU fused-MoE INT4_W4A8 path derived its intermediate size N and group count from the w1_scale tensor shape (attacker-controlled checkpoint geometry), skipped the weight-shape checks that every other quant branch performs, and never validated the scale/zero sizes. A crafted int4 MoE checkpoint that inflates the scale shape while shipping small packed weights makes the tinygemm read out of bounds of packed_w1/packed_w2 (CWE-125). num_groups = 0 additionally causes a divide-by-zero (CWE-369). Validations added: - w1_scale/w2_scale/w1_zero/w2_zero presence and dimensionality checks before deriving N from scale shape - num_groups > 0 guard to prevent divide-by-zero - Cross-validate scale N-block count against packed weight dimensions to prevent inflated scale shapes from causing OOB reads - Full scale/zero element count validation (matching INT8, FP8, MXFP4) - Packed weight size check against stride computed from scale-derived N Signed-off-by: Juan Perez de Algaba <jperezde@redhat.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: jperezde <jperezde@redhat.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
…ulations Use __int128 wide arithmetic for required_w1, required_w2, and buffer_size_nbytes multiplication chains to prevent integer overflow from bypassing the size validation checks. An attacker-controlled checkpoint geometry with large dimensions could wrap around int64_t, making the subsequent TORCH_CHECK pass despite actual buffer sizes being insufficient. Signed-off-by: Juan Perez de Algaba <jperezde@redhat.com> Co-authored-by: Cursor <cursoragent@cursor.com> Signed-off-by: jperezde <jperezde@redhat.com>
| int64_t buffer_size_nbytes = | ||
| M * topk * N * 2 + M * topk * K * 2 + | ||
| num_threads * BLOCK_M * K * | ||
| __int128 buffer_size_wide = |
Contributor
There was a problem hiding this comment.
⚪ Severity: LOW
The buffer size calculation introduces an integer overflow risk. While buffer_size_wide is checked against the maximum limit of int64_t, subsequent additions for INT4_W4A8 and other quantization methods are performed directly on buffer_size_nbytes without overflow validation, potentially leading to heap-based buffer out-of-bounds writes.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Move all branch-specific buffer size additions (for INT8_W8A8, FP8_W8A16/MXFP4, and INT4_W4A8 on lines 1025-1034) into the __int128 buffer_size_wide accumulation before the overflow check on line 1021. This ensures the complete allocation size is validated against int64_t limits before the downcast.
Specifically:
- Remove the three
ifblocks that add tobuffer_size_nbytes(lines 1025-1034). - Instead, add those same quantities to
buffer_size_wide(usingstatic_cast<__int128>(...)casts for each multiplicand) before theTORCH_CHECK. - Also protect the sub-expressions inside
std::max()from int64_t overflow by computing them in__int128, e.g.:
std::max(static_cast<__int128>(M) * K, static_cast<__int128>(M) * topk * N) - After the single
TORCH_CHECK, castbuffer_size_widetoint64_tto getbuffer_size_nbytes.
This consolidates all size arithmetic under the existing __int128 overflow guard.
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
1 participant
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.
The CPU fused-MoE INT4_W4A8 path derived its intermediate size N and
group count from the w1_scale tensor shape (attacker-controlled checkpoint
geometry), skipped the weight-shape checks that every other quant branch
performs, and never validated the scale/zero sizes.
A crafted int4 MoE checkpoint that inflates the scale shape while shipping
small packed weights makes the tinygemm read out of bounds of
packed_w1/packed_w2 (CWE-125). num_groups = 0 additionally causes a
divide-by-zero (CWE-369).