Skip to content

Add tensor size validation for MatMulBnb4 to prevent OOB read via K/N attribute mismatch#27995

Open
vraspar wants to merge 3 commits intomainfrom
vraspar/fix-matmulbnb4-oob-read
Open

Add tensor size validation for MatMulBnb4 to prevent OOB read via K/N attribute mismatch#27995
vraspar wants to merge 3 commits intomainfrom
vraspar/fix-matmulbnb4-oob-read

Conversation

@vraspar
Copy link
Copy Markdown
Contributor

@vraspar vraspar commented Apr 6, 2026

Description

Validates that b_quant and absmax tensor sizes are consistent with K/N/block_size attributes before dequantization in the MatMulBnb4 operator.

Fixes https://portal.microsofticm.com/imp/v5/incidents/details/31000000559964/summary

Changes

  • Constructor: Validate K > 0, N > 0, block_size > 0
  • Compute(): Validate b_quant size >= (KN+1)/2 and absmax size >= ceil(KN/block_size)
  • Tests: Two regression tests for undersized b_quant and absmax tensors

Motivation and Context

MSRC case 109215: A crafted model can set K/N attributes larger than actual tensor sizes, causing OOB reads from b_quant_data and absmax_data during dequantization.

Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

This PR hardens the MatMulBnb4 contrib operator against crafted models by validating that the b_quant and absmax tensors are large enough for the K/N/block_size attributes before dequantization, preventing potential out-of-bounds reads.

Changes:

  • Add basic attribute validation for K, N, and block_size in the CPU kernel constructor.
  • Add runtime size checks in Compute() for b_quant and absmax before dequantization.
  • Add two regression tests that expect failure on undersized b_quant / absmax inputs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 5 comments.

File Description
onnxruntime/contrib_ops/cpu/quantization/matmul_bnb4.cc Adds size/attribute validation intended to prevent OOB reads during BNB4 dequantization on CPU.
onnxruntime/test/contrib_ops/matmul_bnb4_test.cc Adds negative/regression tests that expect failures when b_quant/absmax are undersized vs K/N/block_size.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

… attribute mismatch

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@vraspar vraspar force-pushed the vraspar/fix-matmulbnb4-oob-read branch from 1caf670 to 3a8d28f Compare April 7, 2026 00:23
…ts to CPU EP

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

@vraspar vraspar requested review from jambayk and yuslepukhin April 8, 2026 21:31
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.

2 participants