[Baiju Meswani (baijumeswani)]

Extend the paged-KV XQA decode path to consume a native FP16 KV cache for the head_size=256, group-size-6 geometry used by Qwen3.8 full-attention layers.

• Add an FP16 query/output + FP16 paged-KV H256 XQA instantiation.
• Select it only for metadata-backed one-token-per-sequence decode on SM80+.
• Reuse the existing 128-token XQA page mapping for PagedAttention block sizes divisible by 128.
• Preserve paged Flash Attention for prefill, ragged/multi-token decode, metadata-free callers, unsupported devices, disabled XQA, BF16, and other geometries.
• Keep the runtime dynamic shared-memory fit check authoritative.
• Add CPU-reference numerical coverage and verify XQA dispatch for the Qwen group-size-6 shape.

On an A100-SXM4-80GB with Qwen3.8 27B, the native FP16 XQA kernel reduces B1 x 64K attention time from 0.764 ms to 0.176 ms per full-attention layer per token. End-to-end eager decode improves from 32.9 to 47.5 tok/s (+44%); stable CUDA graphs reach 52.6 tok/s. The 128-token 64K output matches the existing FP16 Flash path token-for-token.

Batch / context	FP16 Flash	FP16 XQA	Gain
B1 x 64K, eager	32.9 tok/s	47.5 tok/s	+44%
B8 x 8K, eager	143.5	171.3	+19%
B8 x 32K, eager	86.3	144.3	+67%
B8 x 64K, eager	43.2	77.7	+80%

[Copilot]

Pull request overview

Adds FP16 paged-KV XQA decode for Qwen3.8’s H256/group-6 CUDA geometry.

Changes:

• Adds and dispatches the FP16 XQA kernel on SM80+.
• Preserves Flash Attention fallbacks and shared-memory checks.
• Adds numerical, dispatch, CUDA Graph, and fallback tests.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Adds FP16 XQA eligibility and fallback routing.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Maps native FP16 cache to unquantized XQA.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h	Documents native-cache support.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu	Dispatches and queries the FP16 H256 kernel.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader_impl.cuh	Documents the native cache element type.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp16_256.cu	Instantiates the new FP16 kernel.
onnxruntime/test/contrib_ops/paged_attention_op_test.cc	Adds correctness, dispatch, graph, and fallback coverage.
docs/contrib_ops/cuda/paged_attention.md	Documents eligibility and fallback behavior.

Suppressed comments (2)

onnxruntime/test/contrib_ops/paged_attention_op_test.cc:917

• This dispatch test has the same unsupported assumption that every SM80+ device can launch this specialization. The runtime may correctly reject the loaded XQA image based on its dynamic shared-memory requirement and use Flash instead, causing this assertion to fail on valid configurations. Gate the XQA expectation on the production shared-memory query result.

  if (GetCudaArchitecture() < 800) {
    GTEST_SKIP() << "XQA requires compute capability 8.0 or later.";
  }

onnxruntime/test/contrib_ops/paged_attention_op_test.cc:966

• This fallback assertion also runs when Flash Attention was not compiled. With USE_FLASH_ATTENTION disabled, the multi-token bound correctly routes to the portable paged decoder instead, making this test fail. Guard or skip the test in non-Flash builds.

TEST(PagedAttention, Cuda_XqaFp16CacheMultiTokenBoundFallsBackToFlash) {

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

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

[Tianlei Wu (tianleiwu)]

Reviewed the FP16-cache paged XQA specialization, including backend eligibility and fallback, native-cache type dispatch, shared-memory and workspace handling, CUDA Graph behavior, build inclusion, and the new correctness/routing tests. The H256/group-6 path is narrowly gated and restores Flash Attention when metadata or runtime capability checks do not prove XQA runnable. No actionable findings remain at this head.