Clamp 1D attention mask_index to valid bounds

Merged
TI
4mo ago
+45-2

PR: Clamp 1D attention mask_index to valid bounds

Description

The CPU attention helper translates a 1D mask_index (right-side end position, and
optionally a left-side start position) into per-token mask-fill loop bounds. The
existing std::max(0, std::min(...)) clamp was correct but verbose; this PR replaces it
with std::clamp and adds a regression test that exercises out-of-range mask values.
This hardens the kernel against out-of-bounds writes when callers supply mask positions
outside [0, all_sequence_length].

Summary of Changes

Bounds clamping

FileChange
onnxruntime/contrib_ops/cpu/bert/attention_helper.hUse std::clamp(static_cast<int>(mask_index[...]), 0, all_sequence_length) for both end_position and start_position in the 1D mask_index branch of PrepareMask.

Test coverage

FileChange
onnxruntime/test/contrib_ops/attention_op_test.ccAdd AttentionMaskIndex1DClampOOB, covering an over-large end_position (999) and a negative start_position (-5). Both clamp cleanly and produce the fully-unmasked output.

Testing

  • Build and run the contrib attention tests:
    onnxruntime_test_all --gtest_filter='ContribOpAttentionTest.AttentionMaskIndex1DClampOOB'
  • The new test asserts that an out-of-range end position (above all_sequence_length)
    applies no right-side masking, and a negative start position applies no left-side
    masking — matching a fully-unmasked sequence.
  • Existing ContribOpAttentionTest.* cases continue to pass; behavior for in-range
    mask values is unchanged.

Motivation and Context

end_position/start_position drive raw index loops into the mask buffer. An
unclamped value past all_sequence_length (or a negative start) could otherwise step
outside the intended range. std::clamp guarantees both stay within
[0, all_sequence_length], preventing out-of-bounds writes with no behavior change for
valid inputs.

Checklist

  • Tests added/updated
  • Documentation updated (if applicable)
  • No breaking changes (or documented in description)
  • CI passes