Clamp 1D attention mask_index to valid bounds
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
| File | Change |
|---|---|
onnxruntime/contrib_ops/cpu/bert/attention_helper.h | Use 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
| File | Change |
|---|---|
onnxruntime/test/contrib_ops/attention_op_test.cc | Add 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