Validate MatMulNBits 8-bit g_idx bounds on CUDA

[Akshay Sonawane (apsonawane)]

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] [2]
• 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.

[Copilot]

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_index range validation for CUDA tensors in MatMulNBits<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_idx values.

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.

[Tianlei Wu (tianleiwu)]

Requesting changes because the current validation path makes g_idx models incompatible with CUDA graph capture and adds a blocking device-to-host validation round trip to the inference hot path. Details inline.

[Tianlei Wu (tianleiwu)]

Thanks for moving g_idx validation out of the inference hot path; the rebuilt CUDA provider tests run both new invalid-range cases successfully. Two initialization-path regressions remain: unchecked attribute arithmetic can fault before normal input validation, and valid initializer-based models fail whenever prepacking is disabled. Details inline.

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc:682

The validation result is cached for the lifetime of the kernel, but g_idx is a normal graph input and may change between InferenceSession::Run calls. A valid first invocation sets is_group_index_validated_; every later invocation then skips validation even if it supplies different, out-of-range values.

PrePack only covers immutable initializer inputs, so the cache is safe there but not for a runtime tensor. On a later invalid run, the CUDA clamp silently substitutes another group and returns incorrect output instead of the intended INVALID_ARGUMENT.

Please validate every non-prepacked runtime g_idx; cache the result only when the tensor is an immutable initializer.

[Tianlei Wu (tianleiwu)]

One CUDA test-isolation issue remains. The implementation now addresses the earlier bounds, initialization, and graph-capture concerns, but the new regression tests do not prove that the CUDA kernel owns the node. Details inline.