Reject zero-dimension block_row_indices in SparseAttention - #29015
#29015Reject zero-dimension block_row_indices in SparseAttention#29015
Conversation
tianleiwu
left a comment
There was a problem hiding this comment.
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
RunSparseAttentionInvalidInputTesthelper; this one builds a rawOpTester. That is fine here, but a one-line comment noting that{0, 2}previously triggered a modulo-by-zero crash would clarify intent.
This pull request strengthens the validation logic for the
block_row_indicesinput in the SparseAttention operator and adds a corresponding unit test to ensure that zero-dimension cases are properly rejected.Validation improvements:
Status CheckInputs(insparse_attention_helper.h) to require that the first dimension ofblock_row_indicesis greater than zero, preventing invalid zero-dimension input.Test coverage:
RejectsZeroDimBlockRowIndicesinsparse_attention_op_test.ccto verify that the operator correctly rejects inputs whereblock_row_indiceshas a zero in its first dimension.