Validate SparseAttention CSR indices and key lengths element values

Merged
YU
4mo ago
+533-33

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:

  • Added host-side validation functions (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]
  • Implemented a CUDA kernel (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:

  • Introduced the environment variable ORT_DISABLE_SPARSE_ATTENTION_INPUT_VALIDATION to 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:

  • Refactored test helper function names and test input shapes to match the new validation logic and error messages. [1] [2] [3]

Other:

  • Removed redundant key length value checks from CheckInputs in favor of the new validation routines.