Skip to content

Clamp 1D attention mask_index to valid bounds - #29449

Merged
tianleiwu merged 1 commit into
mainfrom
tlwu/20260630/icm_fix_att_helper
Jul 7, 2026
Merged

Clamp 1D attention mask_index to valid bounds#29449
tianleiwu merged 1 commit into
mainfrom
tlwu/20260630/icm_fix_att_helper

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

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

Use std::clamp for end_position and start_position in the 1D mask_index
branch of attention_helper.h so out-of-range mask values cannot drive
out-of-bounds writes into the mask buffer. Add AttentionMaskIndex1DClampOOB
test covering an over-large end_position and a negative start_position.
@tianleiwu
tianleiwu merged commit f2559a5 into main Jul 7, 2026
85 of 86 checks passed
@tianleiwu
tianleiwu deleted the tlwu/20260630/icm_fix_att_helper branch July 7, 2026 17:46
tianleiwu added a commit that referenced this pull request Jul 7, 2026
# 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

- [x] Tests added/updated
- [ ] Documentation updated (if applicable)
- [x] No breaking changes (or documented in description)
- [ ] CI passes
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants