onnxruntime
066b90ef - Validate MatMulNBits 8-bit g_idx bounds on CUDA (#31643)

Commit
9 days ago
Validate MatMulNBits 8-bit g_idx bounds on CUDA (#31643) 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:** * Added a new function `ValidateGroupIndexRangeForCuda` in `matmul_nbits.cc` to check that all values in the `group_index` tensor 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]](diffhunk://#diff-06d2e44c42a93a3c7e6b6f717115421091ee66639d0d324515042bd30d308871R37-R64) [[2]](diffhunk://#diff-06d2e44c42a93a3c7e6b6f717115421091ee66639d0d324515042bd30d308871R634-R638) * In the CUDA kernel `Dequantize8BitsKernelReOrder`, added an assertion and clamping to ensure that `rid` (the group index) is within the valid range, further protecting against invalid memory access. **Testing Enhancements:** * Added two new tests in `matmul_8bits_test.cc` to verify that the operator fails as expected when `g_idx` contains out-of-range or negative values, ensuring the new validation logic is exercised. **Miscellaneous:** * Included the `<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.
Author
Parents
Loading