MatmulNBits prepacking scales fix - #27412
Conversation
…es processing the B input
There was a problem hiding this comment.
Pull request overview
This PR fixes a critical bug in the MatMulNBits operator's PrePack method for the MLFloat16 specialization. The bug occurs when prepacking the B input tensor (weights) and scales need to be converted from MLFloat16 to float32. The code was incorrectly using the B tensor's size instead of the scales tensor's size for buffer allocation and conversion, which could lead to buffer overruns or underruns depending on the relative sizes of B and scales tensors.
Changes:
- Fixed incorrect size calculation when prepacking scales for MLFloat16 MatMulNBits operator
- Changed from using B tensor size to scales tensor size for scales conversion buffer allocation
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 3 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
9d40b7d
6f5da27
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Description
Fix incorrect scales element count while pre-packing scales while we processing the B input in the Prepack() method of MatmulNBits operator
Motivation and Context
Fix potential crash due to incorrect element count