cuda pagedattention seqlens shape fix

[Akshay Sonawane (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: ",

[Tianlei Wu (tianleiwu)]

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

[github-actions]

You can commit the suggested changes from lintrunner.

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h (CheckBlockTableAndPastSeqLensValues)

The new validation rejects every negative block-table entry, including -1. The PagedAttention contract explicitly defines block_table == -1 as an unmapped/evicted block; padded and sliding-window tables legitimately contain this sentinel. Because validation scans the entire padded table rather than only dereferenced entries, valid requests fail before execution.

The repository contract documents this at docs/contrib_ops/cuda/paged_attention.md (the input table and sections 9.2/9.3): -1 means unmapped, not invalid.

Please preserve the -1 sentinel and validate only values that can actually be dereferenced (or accept -1 while continuing to reject values < -1 and >= num_blocks).

[Ti-Tai Wang (titaiwangms)]

Critical review finding

onnxruntime/contrib_ops/cuda/bert/paged_attention.cc:532

CheckBlockTableAndPastSeqLensValues runs only inside the needs_readback branch:

const bool needs_readback = !has_metadata_bounds && (needs_dense_kv || xqa_candidate);

Supplying attention_metadata makes has_metadata_bounds true and skips the validation entirely. This is specifically the CUDA-graph-compatible path. The common unquantized paged FlashAttention path can also have needs_dense_kv == false and skip it.

An out-of-range runtime block_table value then reaches CUDA consumers and is used to compute raw K/V-cache offsets without a device-side block_id >= num_blocks guard. For example, block_table=[num_blocks] with otherwise valid metadata can address beyond the cache allocation, causing GPU memory corruption or a CUDA context fault.

Please decouple index-safety validation from the decision to perform host readback. CUDA-graph paths need capture-safe device-side bounds protection at every block-table address resolver; host validation can remain diagnostic defense-in-depth outside capture.

[Copilot]

Copilot review overview

🟡 Changes recommended

Some dispatch paths bypass validation, invalid blocks can corrupt block zero, and the new metadata requirement conflicts with documentation.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 2 · 1

Open (3)

• Validate cumulative sequence lengths on every dispatch path · New
• Map invalid block IDs to the unmapped sentinel · New
• Preserve optional attention metadata for CUDA graph capture · New

Resolved since last review (3)

• CheckBlockTableAndPastSeqLensValues validates non-decreasing cumulative_sequence_length, but it…
• ComputeInternal now always performs 3 device→host copies plus a cudaStreamSynchronize() to validate…
• The new unit tests cover several invalid cases, but don’t cover the important boundary case of…

[Tianlei Wu (tianleiwu)]

The latest revision addresses the CUDA-graph synchronization concern and includes the requested decode benchmark. Two input-validation paths still need correction: sequence-value checks are skipped on normal no-readback execution paths, and a zero-width block table reaches an invalid CUDA launch.

[Ti-Tai Wang (titaiwangms)]

The all-path device sanitization and zero-width block-table rejection address the previous review. The new scratch-buffer counts still use unchecked signed int arithmetic: batch_size * max_num_blocks_per_seq and 2 * batch_size + 1, followed by grid rounding with element_count + 255. Large valid dimensions can overflow before allocation/launch and produce an undersized buffer or invalid indexing.

Please compute these counts with checked size_t/SafeInt, reject values that exceed the kernel indexing representation, and widen b * max_num_blocks_per_seq and launch-size arithmetic consistently.

Non-blocking performance concern: SanitizeSequenceLengths<<<1,1>>> serially scans every mapped block on every decode step. Please provide benchmark evidence for representative large batch/context shapes or parallelize the scan before claiming the previous batch-size limit is removed without qualification.

[Copilot]

Copilot review overview

🟡 Changes recommended

Sliding-window metadata is handled incorrectly, malformed offsets can evade validation, the scan is serialized, and documentation remains contradictory.

Get a fresh assessment by requesting another Copilot review.

Review effort: Balanced
Findings: 1 · 1 · 1

Open (3)

• Preserve past length for leading evicted blocks · New
• Avoid serializing batch metadata scan on one thread · New
• Update contradictory batch-size limit documentation · New

Resolved since last review (3)

• Map invalid block IDs to the unmapped sentinel
• Validate cumulative sequence lengths on every dispatch path
• Preserve optional attention metadata for CUDA graph capture

[Ti-Tai Wang (titaiwangms)]

Re-reviewed the amended head. Scratch counts now use checked size_t/SafeInt arithmetic with explicit kernel indexing limits, and sequence sanitization is parallelized with grid-wide scans plus coverage above the former 256-sequence limit. The prior arithmetic and serialization concerns are resolved. The required Android minimal binary-size check is still failing and must be green before merge.