Validate MatMulNBits 8-bit g_idx bounds on CUDA - #31643
Open
apsonawane wants to merge 3 commits into
Open
Conversation
Add missing rid bounds assert/clamp in Dequantize8BitsKernelReOrder and validate non-CPU g_idx tensors on host before launch. Includes CUDA regression tests for out-of-range and negative g_idx values.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens the CUDA MatMulNBits (8-bit blockwise) path by validating the g_idx/group_index input to prevent out-of-bounds indexing into per-block scales/zero_points, and adds negative tests to ensure invalid indices fail with a clear error.
Changes:
- Added host-side
group_indexrange validation for CUDA tensors inMatMulNBits<T>::ComputeInternal()(pre-kernel launch). - Added device-side assertion + clamping in the reorder dequantization CUDA kernel as defense-in-depth.
- Added CUDA tests that expect failures for negative and out-of-range
g_idxvalues.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/matmul_8bits_test.cc | Adds CUDA tests asserting invalid g_idx causes operator failure. |
| onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc | Adds CUDA-side group_index validation before launching dequant/GEMM. |
| onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_8bits.cu | Adds kernel-level assert/clamp around rid (group index) usage. |
apsonawane
enabled auto-merge (squash)
August 5, 2026 17:41
Add CUDA graph capture detection to the group_index validation helper so we fail fast with a clear ORT status instead of performing a host round-trip and stream synchronization during capture. Also switch the device-to-host copy to an explicit cudaMemcpyDeviceToHost for clarity. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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 improves the robustness and reliability of the CUDA implementation for blockwise 8-bit quantized matrix multiplication in ONNX Runtime. The main focus is on validating the group index (
g_idx) input to prevent out-of-bounds memory access, and on adding tests to verify error handling for invalid group indices.Validation and Error Handling Improvements:
ValidateGroupIndexRangeForCudainmatmul_nbits.ccto check that all values in thegroup_indextensor are within the valid range[0, k_blocks), returning an error if any value is out of bounds. This is called before launching the CUDA kernel. [1] [2]Dequantize8BitsKernelReOrder, added an assertion and clamping to ensure thatrid(the group index) is within the valid range, further protecting against invalid memory access.Testing Enhancements:
matmul_8bits_test.ccto verify that the operator fails as expected wheng_idxcontains out-of-range or negative values, ensuring the new validation logic is exercised.Miscellaneous:
<vector>header to support host-side validation logic.These changes collectively improve the safety and reliability of quantized matrix multiplication on CUDA by proactively catching invalid input and providing clear error messages.