Validate per-column weight_scale/weight_zero_point shape in CPU QAttention; harden integer arithmetic in QAttention and AttentionBase (microsoft#28480)

### Description

The CPU `QAttention` kernel did not validate the shape of per-column
`weight_scale` and `weight_zero_point` inputs against the expected `3 *
hidden_size`. A model could supply a per-column
tensor smaller than expected, causing the GEMM dequantization loop to
read past the end of the buffer (offsets up to `~3 * hidden_size -
head_size`).

This PR adds the missing shape validation and, while in the area,
hardens integer arithmetic across `QAttention` and `AttentionBase`
against malformed shape attributes / dimensions.

### Changes

**`onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc`**
- Validate per-column `weight_scale` and `weight_zero_point` are 1-D
with size `3 * hidden_size`; reject otherwise.
- Use `narrow<int>` / `narrow<size_t>` when converting `int64_t` shape
dims, so out-of-range values throw rather than silently truncating.
- Use `SafeInt` for multiplications whose operands are not provably
bounded by upstream validation (`loop_len`, `input_offset`,
`qkv_offset`, the gemm allocation, and
`packed_weights_data_size` in `PrePack`).
- Refactor the gemm allocation and Q/K/V pointer arithmetic to share a
single `SafeInt`-validated `batch_size * sequence_length * hidden_size`
value.
- Drop a few redundant `static_cast<int>`s in the per-iteration index
math.
- Remove the `hidden_size_x3 % 3 == 0` and `hidden_size % num_heads_ ==
0` checks here; they are now enforced uniformly in
`AttentionBase::CheckInputs` with clearer error messages.

**`onnxruntime/contrib_ops/cpu/bert/attention_base.h`**
- Replace `static_cast<int>` with `narrow<int>` for `num_heads_`,
`rotary_embedding_`, the `parameters` struct outputs, and `GetPresent`'s
`past_sequence_length`. Without this, any
`int64_t` value outside the `int` range (e.g., a `num_heads` attribute
of `2^31`, or a `past` sequence length of `2^31`) silently truncates to
an unrelated `int` value that is then
propagated to downstream kernels and used in arithmetic, enabling
division by zero, sign flips, or out-of-bounds indexing.
- Drop the `static_cast<int>` from the `past_dims[2]` / `past_dims[4]`
shape comparisons so the equality check uses the full `int64_t` value;
previously a `past` tensor whose dim's low 32
bits happened to match `num_heads_` (or `k_hidden_size / num_heads_`)
would pass validation despite having the wrong physical shape.
- In `CheckInputs`, when `require_same_hidden_size_` is true, reject
`bias_dims[0]` not a multiple of 3 with a clear error (Q, K, V are
packed and share a hidden size).
- In `CheckInputs`, when `qkv_hidden_sizes` is not set, also reject
`q_hidden_size % num_heads_ != 0` (mirrors the existing check on the
`qkv_hidden_sizes` path).

**`onnxruntime/test/contrib_ops/quantize_attention_op_test.cc`**
- 4 regression tests for the per-column shape validation:
  - `InvalidWeightScalePerColumnShape`
  - `InvalidWeightScalePerColumnRank`
  - `InvalidWeightZeroPointPerColumnShape`
  - `InvalidWeightZeroPointPerColumnRank`
- 3 regression tests for the divisibility / narrowing checks (sharing a
`RunQAttentionExpectFailure` helper):
  - `InvalidBiasDimNotMultipleOfThree`
  - `InvalidHiddenSizeNotDivisibleByNumHeads`
- `InvalidNumHeadsOverflowsInt` (`num_heads = INT_MAX + 1` triggers
`gsl::narrowing_error`)

### Testing

All `QAttention*` / `AttentionTest*` / `MultiHeadAttention*` tests
(97/97) pass locally on CPU Release build.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 3 people

onnxruntime/contrib_ops/cpu/bert/attention_base.h

@@ -6,6 +6,7 @@
 #include <array>
 #include <vector>
 #include "core/common/common.h"
+#include "core/common/narrow.h"
 #include "core/providers/cpu/mlas_backend_kernel_selector_config_utils.h"
 #ifndef SHARED_PROVIDER
 #include "core/framework/op_kernel.h"
@@ -57,11 +58,11 @@ class AttentionBase {
   AttentionBase(const KernelInfoType& info, bool require_same_hidden_size) {
     int64_t num_heads = 0;
     ORT_ENFORCE(info.GetAttr("num_heads", &num_heads).IsOK() && num_heads > 0);
-    num_heads_ = static_cast<int>(num_heads);
+    num_heads_ = narrow<int>(num_heads);
 
     is_unidirectional_ = info.template GetAttrOrDefault<int64_t>("unidirectional", 0) == 1;
     do_rotary_ = info.template GetAttrOrDefault<int64_t>("do_rotary", 0) == 1;
-    rotary_embedding_ = static_cast<int>(info.template GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));
+    rotary_embedding_ = narrow<int>(info.template GetAttrOrDefault<int64_t>("rotary_embedding_dim", 0));
     mask_filter_value_ = info.template GetAttrOrDefault<float>("mask_filter_value", -10000.0f);
     scale_ = info.template GetAttrOrDefault<float>("scale", 0.0f);
     if (!info.template GetAttrs<int64_t>("qkv_hidden_sizes", qkv_hidden_sizes_).IsOK()) {
@@ -222,6 +223,14 @@ inline Status AttentionBase::CheckInputs(const TensorShape& input_shape,
                            "Input 'bias' dimension 0 should have same length as dimension 1 of input 'weights'");
   }
 
+  // Q, K, V are packed along bias_dims[0]. When their hidden sizes are required to be equal,
+  // bias_dims[0] == 3 * hidden_size must be a multiple of 3.
+  if (require_same_hidden_size_ && bias_dims[0] % 3 != 0) {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "Input 'bias' dimension 0 (", bias_dims[0],
+                           ") must be a multiple of 3 (Q, K, V are packed and have equal hidden sizes).");
+  }
+
   int64_t q_hidden_size = bias_dims[0] / static_cast<int64_t>(3);
   int64_t k_hidden_size = q_hidden_size;
   int64_t v_hidden_size = k_hidden_size;
@@ -241,6 +250,10 @@ inline Status AttentionBase::CheckInputs(const TensorShape& input_shape,
     q_hidden_size = qkv_hidden_sizes_[0];
     k_hidden_size = qkv_hidden_sizes_[1];
     v_hidden_size = qkv_hidden_sizes_[2];
+  } else if (q_hidden_size % num_heads_ != 0) {
+    // Match the error message produced by the qkv_hidden_sizes path above.
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "hidden_size should be divisible by num_heads:", q_hidden_size);
   }
 
   int64_t kv_sequence_length = sequence_length;
@@ -282,14 +295,14 @@ inline Status AttentionBase::CheckInputs(const TensorShape& input_shape,
                              "Inputs 'past' dimension 1 shall have same length as dimension 0 of input 0");
     }
 
-    if (static_cast<int>(past_dims[2]) != num_heads_) {
+    if (past_dims[2] != num_heads_) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                              "Inputs 'past' dimension 2 shall have length of num_heads", num_heads_);
     }
 
-    if (static_cast<int>(past_dims[4]) != k_hidden_size / num_heads_) {
+    if (past_dims[4] != k_hidden_size / num_heads_) {
       return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Inputs 'past' dimension 2 shall have length of ", k_hidden_size / num_heads_);
+                             "Inputs 'past' dimension 4 shall have length of ", k_hidden_size / num_heads_);
     }
 
     if (!past_present_share_buffer_) {
@@ -348,17 +361,17 @@ inline Status AttentionBase::CheckInputs(const TensorShape& input_shape,
 
   if (parameters != nullptr) {
     AttentionParameters* output_parameters = reinterpret_cast<AttentionParameters*>(parameters);
-    output_parameters->batch_size = static_cast<int>(batch_size);
-    output_parameters->sequence_length = static_cast<int>(sequence_length);
-    output_parameters->past_sequence_length = static_cast<int>(past_sequence_length);
-    output_parameters->kv_sequence_length = static_cast<int>(kv_sequence_length);
-    output_parameters->total_sequence_length = static_cast<int>(total_sequence_length);
-    output_parameters->max_sequence_length = static_cast<int>(max_sequence_length);
-    output_parameters->input_hidden_size = static_cast<int>(input_hidden_size);
-    output_parameters->hidden_size = static_cast<int>(q_hidden_size);
-    output_parameters->v_hidden_size = static_cast<int>(v_hidden_size);
-    output_parameters->head_size = static_cast<int>(q_hidden_size) / num_heads_;
-    output_parameters->v_head_size = static_cast<int>(v_hidden_size) / num_heads_;
+    output_parameters->batch_size = narrow<int>(batch_size);
+    output_parameters->sequence_length = narrow<int>(sequence_length);
+    output_parameters->past_sequence_length = narrow<int>(past_sequence_length);
+    output_parameters->kv_sequence_length = narrow<int>(kv_sequence_length);
+    output_parameters->total_sequence_length = narrow<int>(total_sequence_length);
+    output_parameters->max_sequence_length = narrow<int>(max_sequence_length);
+    output_parameters->input_hidden_size = narrow<int>(input_hidden_size);
+    output_parameters->hidden_size = narrow<int>(q_hidden_size);
+    output_parameters->v_hidden_size = narrow<int>(v_hidden_size);
+    output_parameters->head_size = narrow<int>(q_hidden_size) / num_heads_;
+    output_parameters->v_head_size = narrow<int>(v_hidden_size) / num_heads_;
     output_parameters->num_heads = num_heads_;
     output_parameters->is_unidirectional = is_unidirectional_;
     output_parameters->past_present_share_buffer = (past_present_share_buffer_ != 0 && past != nullptr);
@@ -398,7 +411,7 @@ inline Tensor* AttentionBase::GetPresent(TOpKernelContext* context,
                                          int head_size,
                                          int kv_sequence_length,
                                          int& past_sequence_length) const {
-  past_sequence_length = (nullptr != past) ? static_cast<int>(past->Shape().GetDims()[3]) : 0;
+  past_sequence_length = (nullptr != past) ? narrow<int>(past->Shape().GetDims()[3]) : 0;
   std::array<int64_t, 5> present_dims{2, batch_size, num_heads_,
                                       static_cast<int64_t>(kv_sequence_length) + past_sequence_length, head_size};
 

onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc

@@ -7,6 +7,7 @@
 #include "core/util/math.h"
 #include "core/util/qmath.h"
 #include "core/util/math_cpuonly.h"
+#include "core/common/narrow.h"
 #include "core/common/safeint.h"
 #include "core/platform/threadpool.h"
 #include "core/mlas/inc/mlas.h"
@@ -71,8 +72,8 @@ Status QAttention<T>::PrePack(const Tensor& weights, int input_idx, AllocatorPtr
     return Status::OK();
   }
 
-  const size_t input_hidden_size = static_cast<size_t>(weights_dims[0]);
-  const size_t hidden_size_x3 = static_cast<size_t>(weights_dims[1]);
+  const size_t input_hidden_size = narrow<size_t>(weights_dims[0]);
+  const size_t hidden_size_x3 = narrow<size_t>(weights_dims[1]);
   const size_t hidden_size = hidden_size_x3 / 3;
   const size_t head_size = hidden_size / num_heads_;
 
@@ -89,8 +90,8 @@ Status QAttention<T>::PrePack(const Tensor& weights, int input_idx, AllocatorPtr
     return Status::OK();
   }
 
-  const size_t loop_len = 3 * static_cast<size_t>(num_heads_);
-  size_t packed_weights_data_size = packed_weights_size_ * loop_len;
+  const size_t loop_len = SafeInt<size_t>(3) * num_heads_;
+  size_t packed_weights_data_size = SafeInt<size_t>(packed_weights_size_) * loop_len;
 
   packed_weights_ = IAllocator::MakeUniquePtr<void>(alloc, packed_weights_data_size, true);
   std::byte* packed_weights_data = static_cast<std::byte*>(packed_weights_.get());
@@ -171,12 +172,6 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
   T input_scale = *(input_scale_tensor->Data<T>());
 
   bool is_weight_scale_per_column = !IsScalarOr1ElementVector(weight_scale_tensor);
-  const T* weight_scale_data = weight_scale_tensor->Data<T>();
-
-  std::vector<T> dequant_scales(weight_scale_data, weight_scale_data + weight_scale_tensor->Shape().Size());
-  std::for_each(dequant_scales.begin(), dequant_scales.end(), [&input_scale](float& dequant_scale) {
-    return dequant_scale *= input_scale;
-  });
 
   uint8_t input_zero_point = 0;
   if (i_zp_tensor != nullptr) {
@@ -194,14 +189,42 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
   }
 
   const auto& shape = input->Shape();
-  const int batch_size = static_cast<int>(shape[0]);
-  const int sequence_length = static_cast<int>(shape[1]);
-  const int input_hidden_size = static_cast<int>(shape[2]);
-
-  const auto hidden_size_x3 = weights_shape.GetDims()[1];
-  const int hidden_size = static_cast<int>(hidden_size_x3) / 3;
+  const int batch_size = narrow<int>(shape[0]);
+  const int sequence_length = narrow<int>(shape[1]);
+  const int input_hidden_size = narrow<int>(shape[2]);
+
+  const int hidden_size_x3 = narrow<int>(weights_shape.GetDims()[1]);
+  // AttentionBase::CheckInputs verifies that weights_dims[1] (== bias_dims[0]) is a multiple of 3
+  // and that hidden_size = bias_dims[0] / 3 is divisible by num_heads_.
+  const int hidden_size = hidden_size_x3 / 3;
   const int head_size = hidden_size / num_heads_;
 
+  // Validate per-column 'weight_scale' / 'weight_zero_point' shapes against the expected
+  // 3 * hidden_size. Without this check, a malicious or malformed model can supply a
+  // smaller per-column tensor and cause an out-of-bounds read in the GEMM loop below
+  // (which indexes scales/zero-points using offsets up to ~3 * hidden_size - head_size).
+  if (is_weight_scale_per_column) {
+    ORT_RETURN_IF_NOT(weight_scale_tensor->Shape().NumDimensions() == 1 &&
+                          weight_scale_tensor->Shape().Size() == hidden_size_x3,
+                      "Input 'weight_scale' must be a scalar or a 1D tensor of size 3 * hidden_size (= ",
+                      hidden_size_x3, "), got shape ", weight_scale_tensor->Shape().ToString());
+  }
+
+  if (is_weight_zp_per_column) {
+    ORT_RETURN_IF_NOT(w_zp_tensor->Shape().NumDimensions() == 1 &&
+                          w_zp_tensor->Shape().Size() == hidden_size_x3,
+                      "Input 'weight_zero_point' must be a scalar or a 1D tensor of size 3 * hidden_size (= ",
+                      hidden_size_x3, "), got shape ", w_zp_tensor->Shape().ToString());
+  }
+
+  // Build the dequantization scales after shape validation so that malformed
+  // inputs are rejected before any allocation/copy work.
+  const T* weight_scale_data = weight_scale_tensor->Data<T>();
+  std::vector<T> dequant_scales(weight_scale_data, weight_scale_data + weight_scale_tensor->Shape().Size());
+  std::for_each(dequant_scales.begin(), dequant_scales.end(), [&input_scale](T& dequant_scale) {
+    return dequant_scale *= input_scale;
+  });
+
   std::vector<int64_t> output_shape(3);
   output_shape[0] = shape[0];
   output_shape[1] = shape[1];
@@ -216,16 +239,17 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
   auto* tp = context->GetOperatorThreadPool();
   // STEP.1: gemm_data(BS, 3NH) = Scale(input(BS, D) x weights(D, 3NH)) + bias(3NH)
   // D is hidden dimension of input, where input_hidden_size (D) could be larger than hidden_size (NH) when model is pruned.
-  auto gemm_data = allocator->Alloc(SafeInt<size_t>(batch_size) * sequence_length * 3 * hidden_size * element_size);
+  const auto batch_size_x_sequence_length_x_hidden_size = SafeInt<size_t>(batch_size) * sequence_length * hidden_size;
+  void* gemm_data = allocator->Alloc(batch_size_x_sequence_length_x_hidden_size * 3 * element_size);
   BufferUniquePtr gemm_buffer(gemm_data, BufferDeleter(std::move(allocator)));
 
   auto Q = reinterpret_cast<T*>(gemm_data);
-  auto K = Q + static_cast<int64_t>(batch_size) * sequence_length * hidden_size;
-  auto V = K + static_cast<int64_t>(batch_size) * sequence_length * hidden_size;
+  auto K = Q + static_cast<size_t>(batch_size_x_sequence_length_x_hidden_size);
+  auto V = K + static_cast<size_t>(batch_size_x_sequence_length_x_hidden_size);
   T* QKV[3] = {Q, K, V};
 
   {
-    const int loop_len = 3 * batch_size * num_heads_;
+    const int loop_len = SafeInt<int>(3) * batch_size * num_heads_;
     const auto* input_data = input->Data<uint8_t>();
     const auto* bias_data = bias->Data<T>();
 
@@ -243,16 +267,17 @@ Status QAttention<T>::Compute(OpKernelContext* context) const {
     scale_bias_procs.reserve(loop_len);
 
     for (int i = 0; i < loop_len; i++) {
-      const int batch_index = static_cast<int>((i / 3) / num_heads_);
-      const int head_index = static_cast<int>((i / 3) % num_heads_);
-      const int qkv_index = static_cast<int>(i % 3);
+      const int batch_index = (i / 3) / num_heads_;
+      const int head_index = (i / 3) % num_heads_;
+      const int qkv_index = i % 3;
 
-      int input_offset = batch_index * sequence_length * input_hidden_size;
+      int input_offset = SafeInt<int>(batch_index) * sequence_length * input_hidden_size;
       int weights_offset = qkv_index * hidden_size + head_index * head_size;
       int weights_scale_offset = is_weight_scale_per_column ? weights_offset : 0;
       int weights_zp_offset = is_weight_zp_per_column ? weights_offset : 0;
       float* qkv_dest = QKV[qkv_index];
-      int qkv_offset = (batch_index * num_heads_ + head_index) * (sequence_length * head_size);
+      int qkv_offset = (SafeInt<int>(batch_index) * num_heads_ + head_index) *
+                       (SafeInt<int>(sequence_length) * head_size);
 
       //                   original           transposed            iteration
       // A: input          (BxSxD)            (B.)S x D             S x D