onnxruntime

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

Validate per-column weight_scale/weight_zero_point shape in CPU QAttention; harden integer arithmetic in QAttention and AttentionBase (#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>