Validate per-column weight_scale/weight_zero_point shape in CPU QAttention; harden integer arithmetic in QAttention and AttentionBase

[edgchen1]

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.

[Copilot]

Pull request overview

This PR hardens CPU QAttention/AttentionBase against malformed model inputs by adding missing shape validation for per-column weight_scale/weight_zero_point and by tightening integer conversions/overflow behavior to prevent OOB reads and unsafe arithmetic.

Changes:

• Add validation that per-column weight_scale and weight_zero_point are 1-D tensors of size 3 * hidden_size in CPU QAttention.
• Replace truncating static_cast<int> conversions with narrow<int> and introduce additional SafeInt-checked arithmetic in QAttention and AttentionBase.
• Add regression tests covering invalid per-column scale/zero-point rank/size for CPU QAttention.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/contrib_ops/cpu/quantization/attention_quant.cc	Adds per-column scale/zp shape validation and hardens integer arithmetic/offset calculations in CPU QAttention.
onnxruntime/contrib_ops/cpu/bert/attention_base.h	Hardens conversions/shape comparisons to avoid truncation-based validation bypasses.
onnxruntime/test/contrib_ops/quantize_attention_op_test.cc	Adds CPU-only regression tests for invalid per-column scale/zero-point shapes/ranks.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[yuslepukhin]

Test coverage gaps:

No test for the new hidden_size % num_heads_ check (commit 098e4ea). A test with e.g. hidden_size=5, num_heads=2 that expects failure with "must be divisible by num_heads" should be added.

No test for the hidden_size_x3 % 3 == 0 check. A test with weights_dims[1] not divisible by 3 (e.g., 13) would cover this.

No test for narrow<> throwing on overflow (e.g., num_heads attribute = INT_MAX + 1).

Existing happy-path tests (17 QAttention* tests) cover normal operation and are reported as passing.

[edgchen1]

Test coverage gaps:

No test for the new hidden_size % num_heads_ check (commit 098e4ea). A test with e.g. hidden_size=5, num_heads=2 that expects failure with "must be divisible by num_heads" should be added.

No test for the hidden_size_x3 % 3 == 0 check. A test with weights_dims[1] not divisible by 3 (e.g., 13) would cover this.

No test for narrow<> throwing on overflow (e.g., num_heads attribute = INT_MAX + 1).

Existing happy-path tests (17 QAttention* tests) cover normal operation and are reported as passing.

Added some additional tests.

[yuslepukhin]

LGTM