Use native block tables for 128-token XQA pages

[Baiju Meswani (baijumeswani)]

Avoid expanding the PagedAttention block table when its block size already matches XQA’s 128-token page size.

For 128-token blocks, the existing block table is passed directly to XQA. Larger block sizes continue using the existing expanded page-table path.

This removes an unnecessary temporary allocation and conversion kernel from each eligible XQA invocation.

[Copilot]

Pull request overview

Optimizes CUDA PagedAttention by passing native 128-token block tables directly to XQA.

Changes:

• Skips page-table allocation and expansion for 128-token blocks.
• Adds debug telemetry distinguishing native and expanded paths.
• Adds parity tests and documentation.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Conditionally allocates expansion scratch.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Passes native tables directly to XQA.
onnxruntime/contrib_ops/cuda/bert/attention_data.h	Renames expansion scratch field.
onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.h	Adds page-table debug state.
onnxruntime/contrib_ops/cuda/bert/attention_kernel_options.cc	Prints the selected table mode.
onnxruntime/test/python/transformers/test_paged_attention.py	Tests native/expanded parity and dispatch.
docs/contrib_ops/cuda/paged_attention.md	Documents the optimized path.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

LGTM

[Baiju Meswani (baijumeswani)]

Thank you for the review.