cuda pagedattention seqlens shape fix

[apsonawane]

This pull request introduces improved validation for input tensors in the Paged Attention CUDA operator, ensuring that sequence lengths and block table values are properly checked before computation. It adds a new comprehensive validation function, integrates it into the main compute path, and provides unit tests to verify correctness. Additionally, a minor bug in shape checking is fixed.

Input validation improvements:

• Added a new function CheckBlockTableAndPastSeqLensValues in paged_attention_helper.h to validate that cumulative_sequence_length, past_seqlens, and block_table values are within expected ranges, preventing invalid memory accesses and logical errors.
• Integrated this validation into the main compute function in paged_attention.cc, copying relevant device data to CPU and running the checks before computation proceeds.
• Fixed a bug in CheckSequenceLengthTensors where the logical condition for validating the shape of seqlens was incorrect (&& replaced with ||).

Testing:

• Added a new test file paged_attention_helper_test.cc with unit tests covering valid and invalid cases for the new validation logic, ensuring robustness against edge cases.

[Copilot]

Pull request overview

This PR strengthens runtime validation for the CUDA PagedAttention contrib op by adding host-side checks for cumulative_sequence_length, past_seqlens, and block_table, fixes a shape-check bug in CheckSequenceLengthTensors, and adds unit tests for the helper validation.

Changes:

• Fix CheckSequenceLengthTensors shape validation logic (&& → ||) for seqlens.
• Add CheckBlockTableAndPastSeqLensValues and invoke it from PagedAttention<T>::ComputeInternal after copying device inputs to host.
• Add CUDA-kernel unit tests for sequence-length and block-table/past-seqlens validation.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h	Fixes seqlens shape guard and adds new value-level validation helper for paged attention inputs.
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Integrates the new validation helper into the CUDA PagedAttention compute path via D→H copies + sync.
onnxruntime/test/contrib_ops/cuda_kernels/paged_attention_helper_test.cc	Adds unit tests covering the new helper validation (plus the seqlens shape fix).

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h:220

• The past_seqlens range check currently rejects past_length == max_num_blocks_per_seq * block_size even when q_len == 0. The paged attention implementation explicitly supports batches with zero new tokens, so a “no-write” sequence with a full cache can be a valid edge case. Consider allowing past_length == max_cache_sequence_length when q_len == 0, while still rejecting it when q_len > 0.

    const int32_t q_len = q_end - q_start;
    const int32_t past_length = past_seqlens[b];
    if (past_length < 0 || static_cast<int64_t>(past_length) >= max_cache_sequence_length) {
      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
                             "past_seqlens values must be in [0, max_num_blocks_per_seq * block_size). Invalid value: ",

[tianleiwu]

I agree we shall not introduce memory copy for performance reason.

[github-actions]

You can commit the suggested changes from lintrunner.