Validate SparseAttention CSR indices and key lengths element values - #29242
Merged
Conversation
CheckInputs only validated tensor shapes for block_row_indices and block_col_indices but never checked element values. This allowed attacker-controlled values to be used as array indices in ComputeAttentionProbs, leading to out-of-bounds memory access. Add ValidateCSRIndices and ValidateKeyLengths for CPU EP that check: - CSR row-pointer monotonicity and bounds - Column indices within [0, max_blocks) - Key lengths within valid range Add ValidateCSRIndicesOnDevice for CUDA EP that performs equivalent validation on-device using a single-warp kernel with strided iteration and a scratch buffer error flag. Add unit tests covering invalid CSR row pointers, non-monotonic values, out-of-range column indices, negative values, and large OOB values.
Contributor
There was a problem hiding this comment.
Pull request overview
This PR strengthens input validation for the SparseAttention contrib op by adding explicit CSR (block mask) and key-length element checks on both CPU and CUDA paths, aiming to prevent out-of-bounds access and improve robustness when inputs are malformed.
Changes:
- Added CPU-side validators for CSR indices and
key_total_sequence_lengthselement ranges, and wired them into the CPU compute path. - Added a CUDA device validation kernel + host helper, invoked prior to the CUDA kernel execution.
- Added new unit tests to ensure malformed CSR indices are rejected.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/sparse_attention_op_test.cc | Adds negative tests for CSR element-value validation failures. |
| onnxruntime/contrib_ops/cuda/sparse/sparse_attention.cc | Calls new device-side validation before CUDA kernel execution. |
| onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.h | Declares CSR validation error codes and device validation entrypoint. |
| onnxruntime/contrib_ops/cuda/sparse/sparse_attention_impl.cu | Implements CUDA validation kernel + host wrapper. |
| onnxruntime/contrib_ops/cpu/sparse/sparse_attention.cc | Invokes new CPU-side CSR + key-length validation helpers. |
| onnxruntime/contrib_ops/cpu/sparse/sparse_attention_helper.h | Refactors key-length value checks and adds new CPU validators. |
added 2 commits
June 23, 2026 16:48
- Move ValidateCSRIndices and ValidateKeyLengths from header to anonymous namespace in sparse_attention.cc (not shared across TUs). - Change Tensor pointer parameters to const Tensor& references since they are required (never null). - Extend CUDA ValidateCSRIndicesOnDevice kernel to also validate seqlens_k_total (key lengths) in the same launch. - Add CSRValidationError enum for device kernel error codes. - Use full warp (32 threads) with strided iteration in the kernel. - Scope scratch buffer allocation to validation call site.
…lper - Only validate block_col_indices entries in [0, nnz) where nnz is the last row pointer value r[max_blocks], since entries beyond that are padding and not consumed by the kernel. - Replace atomicExch with atomicCAS to preserve the first detected error code deterministically across concurrent threads/blocks. - Mark RunSparseAttentionCSRValidationTest as static to avoid external linkage.
tianleiwu
reviewed
Jun 24, 2026
tianleiwu
reviewed
Jun 24, 2026
tianleiwu
left a comment
tianleiwu
left a comment
Contributor
There was a problem hiding this comment.
One new finding on the device-side validation (the CPU path and the previously raised CUDA points look good). Also worth noting as a positive: moving the key-length value check out of the shared CheckInputs fixes a latent bug — CheckInputs used to dereference total_key_lengths->Data<int32_t>() on the host, which on the CUDA path is a device pointer (UB). That read is now correctly done on device.
- Fix existing key-length tests to use valid CSR indices (col_count=4). - Rename RunSparseAttentionInvalidInputTest to RunSparseAttentionInvalidKeyLengthsTest. - Fix CUDA kernel race: thread 0 validates row pointers sequentially, then __syncthreads ensures r[max_blocks] is safe before all threads cooperate on col-index validation. - Update kernel comment to reflect actual thread cooperation model. - Make error message layout-agnostic for kCSRValidationRowFirstNotZero. - Add tests: multi-layout invalid col/row in second layout, invalid col within NNZ range, padding acceptance beyond NNZ.
yuslepukhin
force-pushed
the
yuslepukhin/validate_sparse_attention
branch
from
June 24, 2026 18:54
1c97efb to
af1f999
Compare
- Add ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATION env var to skip device-side CSR and key-length validation when inputs are known to be well-formed (addresses @tianleiwu latency concern). - Add CUDA-specific validation tests under #ifdef USE_CUDA that use head_size=128, sparse_block_size=64, MLFloat16 (CUDA EP requirements). Tests cover: row first element, row monotonicity, col OOB, col large value, and key length OOB. Tests run only with DefaultCudaExecutionProvider.
SparseAttention requires past/present to share buffers (in-place). OpTester allocates separate buffers, causing pointer equality check to fail before reaching compute. The padding non-rejection behavior is already implicitly verified by RejectsBlockColIndicesInvalidWithinNNZ.
yuslepukhin
force-pushed
the
yuslepukhin/validate_sparse_attention
branch
from
June 24, 2026 20:28
31f4a40 to
bbde0f8
Compare
- Move ValidateCSRIndicesOnDevice call to right after CheckInputs, before the past/present shared-buffer enforcement. The validation only reads input tensors (CSR indices, key lengths) and does not need the KV cache buffers. This allows OpTester-based CUDA tests to exercise validation without IOBinding. - Fix test comment: ORT_MAKE_STATUS returns Status, does not throw. - Document in CUDA test helper why OpTester works without shared buffers (validation runs before the buffer check).
tianleiwu
approved these changes
Jun 26, 2026
This was referenced Jul 25, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces device-side and host-side validation for sparse attention input indices and key sequence lengths, improving error detection and robustness for both CPU and CUDA implementations. It also adds an environment variable to optionally disable the new device-side validation for performance reasons. Some related test code has been updated to match these changes.
Validation improvements:
ValidateCSRIndices,ValidateKeyLengths) for CSR indices and key sequence lengths in the CPU kernel, and integrated them into the compute path. This ensures invalid inputs are caught early on the CPU. [1] [2]ValidateCSRIndicesKernel) and supporting function (ValidateCSRIndicesOnDevice) to check CSR row-pointer monotonicity, column-index range, and key lengths on device, with detailed error codes and messages. The CUDA kernel is invoked in the compute path unless disabled by environment variable. [1] [2] [3]Configurability:
ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATIONto allow users to skip device-side validation for performance when inputs are known to be valid. This is parsed and used in the CUDA kernel. [1] [2] [3]Test updates:
Other:
CheckInputsin favor of the new validation routines.