onnxruntime
a243f8cb - Validate PackedAttention token offsets (#33330)

Commit
4 days ago
Validate PackedAttention token offsets (#33330) This pull request introduces validation for the `token_offset` input in the CUDA implementation of the PackedAttention operator, specifically for the unfused backend. The main goal is to ensure that invalid `token_offset` values are detected early and reported with a clear error message, improving robustness and debuggability. The change includes both the implementation of the validation logic and new tests to verify its correctness. **Validation improvements for PackedAttention:** * Added a CUDA kernel and a host function (`ValidatePackedAttentionTokenOffset`) to check that all `token_offset` values are within the valid range `[0, B * S)` for the unfused backend, and integrated this validation into `PackedAttention<T>::ComputeInternal`. [[1]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR31-R84) [[2]](diffhunk://#diff-450dced704ee2fc03ee3611224808eb8a00671f7587e48599ce17c21d5e7c9e8R16-R21) [[3]](diffhunk://#diff-fd7731431d1e80901e297b2ce7f1018fc7305faf58145cb36733bd19b0aeb9d2R270-R279) * Updated relevant CUDA kernels to handle out-of-range `token_offset` values gracefully by skipping computation or zeroing out invalid output regions. [[1]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR115-R117) [[2]](diffhunk://#diff-4e4c1980e908e7c273f52041c3dc123cb7a8c4ee20779fb28b4975911556c9dbR371-R374) **Testing enhancements:** * Added a new unit test (`InvalidTokenOffset_Unfused_CUDA`) to verify that invalid `token_offset` values for the unfused backend are correctly detected and result in a descriptive error message.
Author
Parents
Loading