Validate GatherBlockQuantized CUDA indices bounds

onnxruntime/contrib_ops/cuda/quantization/gather_block_quantized.cu

@@ -53,6 +53,14 @@ __global__ void GatherBlockQuantizedKernel(
   int64_t idx_after = out_idx % after_gather_dim;
   int64_t idx = (out_idx % (after_gather_dim * ind_dim)) / after_gather_dim;
   int64_t idx_at_g = indices[idx];
+  // Keep invalid dynamic indices from participating in device address arithmetic.
+  if (idx_at_g < -gather_axis_dim || idx_at_g >= gather_axis_dim) {
+    output[out_idx] = static_cast<T2>(0);
+    return;
+  }
+  if (idx_at_g < 0) {
+    idx_at_g += gather_axis_dim;
+  }
   int64_t in_idx = idx_before * gather_axis_dim * after_gather_dim + idx_at_g * after_gather_dim + idx_after;
 
   int64_t block_id = in_idx / block_size;

onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc

@@ -473,7 +473,7 @@ TEST(GatherBlockQuantizedOpTest, ShapeMismatch) {
 #endif
 
 template <typename T1, typename T2, typename Tind>
-void Test_InvalidIndices_WithZeroPoints() {
+void Test_InvalidIndices_WithZeroPoints(bool expect_safe_cuda_output = false) {
   std::vector<int> data = {-8, -7, -6, -5,
                            -4, -3, -2, -1,
                            0, 1, 2, 3,
@@ -495,8 +495,43 @@ void Test_InvalidIndices_WithZeroPoints() {
   constexpr int64_t quantize_axis = 2;
   constexpr int64_t block_size = 16;
   constexpr int64_t bits = 4;
+  if (expect_safe_cuda_output) {
+    output.assign(output.size(), 0.0f);
+  }
+  RunUnpackedData<T1, T2, Tind>(data, data_shape, indices, indices_shape, scales, scales_shape, zero_points,
+                                gather_axis, quantize_axis, block_size, bits, output, output_shape,
+                                expect_safe_cuda_output, true);
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_NegativeInvalidIndices_WithZeroPoints(bool expect_safe_cuda_output = false) {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {-3};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 1};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               3.f, 4.f, 5.f, 6.f,
+                               -6.f, -4.f, -2.f, 0.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  constexpr int64_t gather_axis = 0;
+  constexpr int64_t quantize_axis = 2;
+  constexpr int64_t block_size = 16;
+  constexpr int64_t bits = 4;
+  if (expect_safe_cuda_output) {
+    output.assign(output.size(), 0.0f);
+  }
   RunUnpackedData<T1, T2, Tind>(data, data_shape, indices, indices_shape, scales, scales_shape, zero_points,
-                                gather_axis, quantize_axis, block_size, bits, output, output_shape, false, true);
+                                gather_axis, quantize_axis, block_size, bits, output, output_shape,
+                                expect_safe_cuda_output, true);
 }
 
 #ifndef USE_CUDA
@@ -507,6 +542,28 @@ TEST(GatherBlockQuantizedOpTest, InvalidIndices) {
 }
 #endif
 
+#ifdef USE_CUDA
+TEST(GatherBlockQuantizedOpTest, InvalidIndicesSafelyHandled_Cuda) {
+  if (!HasCudaEnvironment(0)) {
+    GTEST_SKIP() << "CUDA not available";
+  }
+
+  Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>(true);
+  Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int64_t>(true);
+  Test_InvalidIndices_WithZeroPoints<uint8_t, float, int32_t>(true);
+}
+
+TEST(GatherBlockQuantizedOpTest, NegativeInvalidIndicesSafelyHandled_Cuda) {
+  if (!HasCudaEnvironment(0)) {
+    GTEST_SKIP() << "CUDA not available";
+  }
+
+  Test_NegativeInvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>(true);
+  Test_NegativeInvalidIndices_WithZeroPoints<UInt4x2, float, int64_t>(true);
+  Test_NegativeInvalidIndices_WithZeroPoints<uint8_t, float, int32_t>(true);
+}
+#endif
+
 template <typename T1, typename T2, typename Tind>
 void Test_GatherAxis0_WithZeroPoints(int bits = 4) {
   std::vector<int> data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8,