Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

Appearance settings

Reject zero-dimension block_row_indices in SparseAttention - #29015

#29015
Merged
apsonawane merged 1 commit into
mainmicrosoft/onnxruntime:mainfrom
asonawane/gemmmicrosoft/onnxruntime:asonawane/gemmCopy head branch name to clipboard
Jun 12, 2026
Merged

Reject zero-dimension block_row_indices in SparseAttention#29015
apsonawane merged 1 commit into
mainmicrosoft/onnxruntime:mainfrom
asonawane/gemmmicrosoft/onnxruntime:asonawane/gemmCopy head branch name to clipboard

Conversation

@apsonawane

Copy link
Copy Markdown
Contributor

This pull request strengthens the validation logic for the block_row_indices input in the SparseAttention operator and adds a corresponding unit test to ensure that zero-dimension cases are properly rejected.

Validation improvements:

  • Updated the check in Status CheckInputs (in sparse_attention_helper.h) to require that the first dimension of block_row_indices is greater than zero, preventing invalid zero-dimension input.

Test coverage:

  • Added a new unit test RejectsZeroDimBlockRowIndices in sparse_attention_op_test.cc to verify that the operator correctly rejects inputs where block_row_indices has a zero in its first dimension.

@apsonawane
apsonawane enabled auto-merge (squash) June 11, 2026 19:49

@tianleiwu tianleiwu left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sound, minimal bugfix — looks good to merge.

What this fixes: Adding block_row_indices_dim[0] > 0 to the block_row_indices validation prevents num_heads % block_row_indices_dim[0] from executing when the first dimension is 0. An integer modulo-by-zero is undefined behavior and raises SIGFPE on x86, crashing the process instead of returning a clean INVALID_ARGUMENT status. The guard is correctly placed before the modulo term, so C++ short-circuit && evaluation is what actually prevents the UB. This is the only site in the codebase with this pattern, and since this CPU helper is shared with the CUDA SparseAttention op, the single fix covers both EPs.

Test: RejectsZeroDimBlockRowIndices is a genuine regression guard — the {0, 2} input would have crashed before this change. Inputs are internally consistent (head_size = 32/4 = 8; past_key/value {1,4,4,8} match kv_num_heads/head_size), so execution reaches the block_row_indices check and fails there with the expected message.

Minor nits (non-blocking):

  • The PR title "Add conditional logic to improve functionality" is generic; something like "Reject zero-dimension block_row_indices in SparseAttention" would better match the (accurate) description.
  • Other invalid-input tests in this file use the RunSparseAttentionInvalidInputTest helper; this one builds a raw OpTester. That is fine here, but a one-line comment noting that {0, 2} previously triggered a modulo-by-zero crash would clarify intent.

Comment thread onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h
@tianleiwu tianleiwu changed the title Add conditional logic to improve functionality Reject zero-dimension block_row_indices in SparseAttention Jun 12, 2026
@apsonawane
apsonawane merged commit ed7a29b into main Jun 12, 2026
86 checks passed
@apsonawane
apsonawane deleted the asonawane/gemm branch June 12, 2026 22:42
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

Morty Proxy This is a proxified and sanitized view of the page, visit original site.