Validate MatMulNBits 8-bit g_idx bounds on CUDA

onnxruntime/contrib_ops/cuda/quantization/dequantize_blockwise_8bits.cu

@@ -115,10 +115,10 @@ __global__ void Dequantize8BitsKernelReOrder(
   const int32_t* g_idx = reorder_idx + kb_idx * block_size + ((threadIdx.x * element_per_thread) & (block_size - 1));
 
   for (int i = 0; i < element_per_thread; i++) {
-    // Typical value of g_idx is in the range of [0, groups_per_K) for reordering groups.
-    // No range check here so it might have out-of-bound access if the reorder_idx is not valid.
     int32_t rid = g_idx[i];
-    ptrdiff_t scale_zp_offset = n_idx * groups_per_K + rid;
+    CUDA_KERNEL_ASSERT(rid >= 0 && rid < groups_per_K);
+    rid = max(0, min(rid, groups_per_K - 1));  // Clamp for release safety
+    ptrdiff_t scale_zp_offset = static_cast<ptrdiff_t>(n_idx) * groups_per_K + rid;
     T scale = *(scale_data + scale_zp_offset);
 
     uint8_t zp = 128;  // Default zero point

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.cc

@@ -5,6 +5,7 @@
 
 #include <cstdint>
 #include <optional>
+#include <vector>
 
 #include "core/common/status.h"
 #include "core/common/float16.h"
@@ -33,6 +34,75 @@ namespace contrib {
 namespace cuda {
 using namespace onnxruntime::cuda;
 
+namespace {
+
+Status ValidateGroupIndexRange(const Tensor* group_index, int64_t k_blocks, cudaStream_t stream) {
+  if (group_index == nullptr) {
+    return Status::OK();
+  }
+
+  const size_t g_idx_size = static_cast<size_t>(group_index->Shape().Size());
+  const int32_t* g_idx_data = group_index->Data<int32_t>();
+  std::vector<int32_t> g_idx_host;
+  if (group_index->Location().device.Type() != OrtDevice::CPU) {
+    g_idx_host.resize(g_idx_size);
+    CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(g_idx_host.data(), g_idx_data, g_idx_size * sizeof(int32_t),
+                                         cudaMemcpyDeviceToHost, stream));
+    CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream));
+    g_idx_data = g_idx_host.data();
+  }
+
+  for (size_t i = 0; i < g_idx_size; ++i) {
+    if (g_idx_data[i] < 0 || g_idx_data[i] >= k_blocks) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "group_index value at index ", i, " is ", g_idx_data[i],
+                             ", which is out of valid range [0, ", k_blocks, ")");
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace
+
+template <typename T>
+Status MatMulNBits<T>::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
+                               /*out*/ bool& is_packed,
+                               /*out*/ PrePackedWeights* /*prepacked_weights*/) {
+  is_packed = false;
+  constexpr int kInputIndexGroupIndex = 4;
+  if (input_idx == kInputIndexGroupIndex) {
+    const int64_t k_blocks = (K_ + block_size_ - 1) / block_size_;
+    ORT_RETURN_IF_ERROR(ValidateGroupIndexRange(&tensor, k_blocks, cudaStreamLegacy));
+    is_group_index_validated_ = true;
+    return Status::OK();
+  }
+
+#if USE_FPA_INTB_GEMM
+  if constexpr (std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>) {
+    if (has_fpA_intB_gemm_) {
+      cudaStream_t stream = cudaStreamLegacy;
+      if (input_idx == MatMulNBits_Input_B) {
+        ORT_RETURN_IF_ERROR(PrePack_B(tensor, alloc, stream, is_packed));
+        is_prepacked_weight_ = is_packed;
+      } else if (input_idx == MatMulNBits_Input_Scale) {
+        ORT_RETURN_IF_ERROR(PrePack_Scale(tensor, alloc, stream));
+        is_prepacked_scale_ = true;
+        is_packed = true;
+      } else if (input_idx == MatMulNBits_Input_ZeroPoint && has_zero_points_) {
+        ORT_RETURN_IF_ERROR(PrePack_ZeroPoint(tensor, alloc, stream));
+        is_prepacked_zero_point_ = true;
+        is_packed = true;
+      }
+    }
+  }
+#else
+  ORT_UNUSED_PARAMETER(alloc);
+#endif
+
+  return Status::OK();
+}
+
 #if USE_FPA_INTB_GEMM
 using onnxruntime::llm::kernels::weight_only::GemmPluginProfilerManager;
 using onnxruntime::llm::kernels::weight_only::WeightOnlyGroupwiseQuantGemmPluginProfiler;
@@ -322,34 +392,6 @@ void MatMulNBits<T>::RunGemmProfile(bool hasWeightOnlyCudaKernel, int min_m, int
   gemmProfiler_->profileTactics(weightOnlyGemmRunner_, gemmId_.dtype, dims, gemmId_, hasWeightOnlyCudaKernel);
 }
 
-template <typename T>
-Status MatMulNBits<T>::PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
-                               /*out*/ bool& is_packed,
-                               /*out*/ PrePackedWeights* /*prepacked_weights*/) {
-  is_packed = false;
-  if constexpr (std::is_same_v<T, MLFloat16> || std::is_same_v<T, BFloat16>) {
-    if (has_fpA_intB_gemm_) {
-      cudaStream_t stream = cudaStreamLegacy;  // Use default stream for prepacking.
-      if (input_idx == MatMulNBits_Input_B) {
-        ORT_RETURN_IF_ERROR(PrePack_B(tensor, alloc, stream, is_packed));
-        is_prepacked_weight_ = is_packed;
-      } else if (input_idx == MatMulNBits_Input_Scale) {
-        ORT_RETURN_IF_ERROR(PrePack_Scale(tensor, alloc, stream));
-        is_prepacked_scale_ = true;
-        is_packed = true;
-      } else if (input_idx == MatMulNBits_Input_ZeroPoint) {
-        if (has_zero_points_) {
-          ORT_RETURN_IF_ERROR(PrePack_ZeroPoint(tensor, alloc, stream));
-          is_prepacked_zero_point_ = true;
-          is_packed = true;
-        }
-      }
-    }
-  }
-
-  return Status::OK();
-}
-
 template <typename T>
 Status MatMulNBits<T>::PrePack_B([[maybe_unused]] const Tensor& tensor,
                                  [[maybe_unused]] AllocatorPtr alloc,
@@ -600,6 +642,12 @@ Status MatMulNBits<T>::ComputeInternal(OpKernelContext* ctx) const {
   if (Y->Shape().Size() == 0)
     return Status::OK();
 
+  if (reorder_idx != nullptr) {
+    ORT_RETURN_IF_NOT(is_group_index_validated_,
+                      "MatMulNBits group_index must be a constant initializer so CUDA can validate it once ",
+                      "before graph capture and avoid a device-to-host synchronization on every inference.");
+  }
+
   cudaStream_t stream = this->Stream(ctx);
 
   typedef typename onnxruntime::cuda::OrtToCudaType<T>::type CudaT;

onnxruntime/contrib_ops/cuda/quantization/matmul_nbits.h

@@ -239,10 +239,10 @@ class MatMulNBits final : public CudaKernel {
   }
 
   Status ComputeInternal(OpKernelContext* context) const override;
-#if USE_FPA_INTB_GEMM
   Status PrePack(const Tensor& tensor, int input_idx, AllocatorPtr alloc,
                  bool& is_packed, PrePackedWeights* prepacked_weights) override;
 
+#if USE_FPA_INTB_GEMM
 #ifndef BUILD_CUDA_EP_AS_PLUGIN
   // Level 2 (Phase-A memory roadmap, issue microsoft/onnxruntime#29775): instance-level workspace
   // estimate, callable after CreateKernels(). Uses the same constructed runner state that
@@ -288,6 +288,7 @@ class MatMulNBits final : public CudaKernel {
   bool column_wise_quant_blk_{true};
 
   bool has_g_idx_{false};
+  bool is_group_index_validated_{false};
   bool has_bias_{false};
   bool has_zero_points_{false};
   bool is_zero_points_scale_same_type_{false};

onnxruntime/test/contrib_ops/matmul_8bits_test.cc

@@ -556,6 +556,86 @@ TEST(MatMulNBits, BFloat16_Int8_Gemm_Cuda) {
   TestMatMul8BitsTyped<BFloat16, 32, 1024, 1024, 128>(abs_error, rel_error);
 }
 
+TEST(MatMulNBits, InvalidGIdx_OutOfRange_Cuda) {
+  if (!HasCudaEnvironment(0)) {
+    GTEST_SKIP() << "CUDA not available";
+  }
+
+  constexpr int64_t M = 2, N = 4, K = 64, block_size = 32;
+  constexpr int64_t k_blocks = (K + block_size - 1) / block_size;
+  constexpr int64_t blob_size = block_size * QBits / 8;
+
+  OpTester test("MatMulNBits", 1, kMSDomain);
+  test.AddAttribute<int64_t>("K", K);
+  test.AddAttribute<int64_t>("N", N);
+  test.AddAttribute<int64_t>("block_size", block_size);
+  test.AddAttribute<int64_t>("bits", QBits);
+  test.AddAttribute<int64_t>("accuracy_level", int64_t{0});
+
+  std::vector<float> a_data(M * K, 1.0f);
+  test.AddInput<float>("A", {M, K}, a_data, false);
+
+  std::vector<uint8_t> b_data(N * k_blocks * blob_size, 0);
+  test.AddInput<uint8_t>("B", {N, k_blocks, blob_size}, b_data, true);
+
+  std::vector<float> scales(N * k_blocks, 1.0f);
+  test.AddInput<float>("scales", {N, k_blocks}, scales, true);
+
+  test.AddOptionalInputEdge<uint8_t>();
+
+  std::vector<int32_t> g_idx(K, 99999);
+  test.AddInput<int32_t>("g_idx", {K}, g_idx, true);
+
+  test.AddOptionalInputEdge<float>();
+
+  std::vector<float> y_data(M * N, 0.0f);
+  test.AddOutput<float>("Y", {M, N}, y_data);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.emplace_back(DefaultCudaExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure, "group_index value", {}, nullptr, &execution_providers);
+}
+
+TEST(MatMulNBits, InvalidGIdx_Negative_Cuda) {
+  if (!HasCudaEnvironment(0)) {
+    GTEST_SKIP() << "CUDA not available";
+  }
+
+  constexpr int64_t M = 2, N = 4, K = 64, block_size = 32;
+  constexpr int64_t k_blocks = (K + block_size - 1) / block_size;
+  constexpr int64_t blob_size = block_size * QBits / 8;
+
+  OpTester test("MatMulNBits", 1, kMSDomain);
+  test.AddAttribute<int64_t>("K", K);
+  test.AddAttribute<int64_t>("N", N);
+  test.AddAttribute<int64_t>("block_size", block_size);
+  test.AddAttribute<int64_t>("bits", QBits);
+  test.AddAttribute<int64_t>("accuracy_level", int64_t{0});
+
+  std::vector<float> a_data(M * K, 1.0f);
+  test.AddInput<float>("A", {M, K}, a_data, false);
+
+  std::vector<uint8_t> b_data(N * k_blocks * blob_size, 0);
+  test.AddInput<uint8_t>("B", {N, k_blocks, blob_size}, b_data, true);
+
+  std::vector<float> scales(N * k_blocks, 1.0f);
+  test.AddInput<float>("scales", {N, k_blocks}, scales, true);
+
+  test.AddOptionalInputEdge<uint8_t>();
+
+  std::vector<int32_t> g_idx(K, -1);
+  test.AddInput<int32_t>("g_idx", {K}, g_idx, true);
+
+  test.AddOptionalInputEdge<float>();
+
+  std::vector<float> y_data(M * N, 0.0f);
+  test.AddOutput<float>("Y", {M, N}, y_data);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.emplace_back(DefaultCudaExecutionProvider());
+  test.Run(OpTester::ExpectResult::kExpectFailure, "group_index value", {}, nullptr, &execution_providers);
+}
+
 // Chunked dequant+GEMM path tests for INT8.
 // Force the chunked path with a small chunk size to exercise per-chunk pointer
 // arithmetic (blob, scales, zero_points offsets) and strided cuBLAS output