onnxruntime
91769bb2 - Harden contrib CPU int narrowing for attention attrs (#31648)

Commit
3 days ago
Harden contrib CPU int narrowing for attention attrs (#31648) This pull request strengthens input validation for attention-related operators and adds corresponding unit tests to ensure that large attribute values do not cause integer overflows. The changes affect the `LinearAttention`, `LongformerAttentionBase`, and GPT subgraph implementations, improving reliability and error reporting. **Input validation improvements:** * [`onnxruntime/contrib_ops/cpu/bert/linear_attention.cc`](diffhunk://#diff-5c4a68eb89a2f94cd428e245a6d7f15a6713138b94ccf04d8d8de61edfce4834L44-R55): Added checks to ensure `q_num_heads` and `kv_num_heads` attributes are within the range `[1, INT_MAX]`, preventing integer overflows during initialization. * [`onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h`](diffhunk://#diff-1da5f76768826ff4730d0c18c326a4581e3257a889df80237ec0d36b8eccc486L28-R39): Added validation to ensure `num_heads` and `window` attributes are within `[1, INT_MAX]`. * [`onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc`](diffhunk://#diff-1c3c36abe873809b84c5b9b0c4620cdd9acbac51c55cee11b2925440fd62f8b3R176-R182): Added checks to ensure certain tensor dimensions do not exceed `INT_MAX`, preventing unsafe casts. **Testing enhancements:** * [`onnxruntime/test/contrib_ops/linear_attention_op_test.cc`](diffhunk://#diff-3dd3dc17608e4a4603256f802204c16de8f89b61878359375421b659545295a5R1385-R1436): Added new tests to verify that the `LinearAttention` operator correctly rejects out-of-range `q_num_heads` and `kv_num_heads` values, ensuring the new validation logic is exercised. **Code maintenance:** * Included `<limits>` header where necessary to support the new validation checks. [[1]](diffhunk://#diff-5c4a68eb89a2f94cd428e245a6d7f15a6713138b94ccf04d8d8de61edfce4834R12) [[2]](diffhunk://#diff-1da5f76768826ff4730d0c18c326a4581e3257a889df80237ec0d36b8eccc486R7) [[3]](diffhunk://#diff-1c3c36abe873809b84c5b9b0c4620cdd9acbac51c55cee11b2925440fd62f8b3R10) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading