[Baiju Meswani (baijumeswani)]

Extend the quantized paged-KV XQA decode path to head_size=256 for FP16 and BF16 activations with INT8 or FP8 KV caches.

This enables XQA for models such as Qwen3.8 whose full-attention layers use 24 query heads, 4 KV heads, and head size 256.

• Add FP16/BF16 paged XQA instantiations for head size 256 with INT8 and FP8 KV caches.
• Add H256 loader declarations, launch dispatch, and dynamic shared-memory queries.
• Extend PagedAttention's XQA eligibility check to head size 256.
• Add an operator test for the Qwen group-size-6 INT8 geometry that verifies numerical output and confirms SdpaKernel=XQA dispatch.

The existing runtime shared-memory check remains authoritative. If an H256 kernel does not fit the device's opt-in shared-memory limit, PagedAttention uses the existing fallback path.

[Copilot]

Pull request overview

Extends CUDA paged XQA decoding to support head size 256.

Changes:

• Adds FP16/BF16 kernels for INT8 and FP8 caches.
• Extends dispatch, shared-memory checks, and eligibility.
• Adds an INT8 group-size-6 operator test.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Enables H256 XQA eligibility.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h	Documents H256 support.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu	Adds H256 dispatch and memory queries.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_256.cu	Instantiates FP16/INT8 H256.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_256.cu	Instantiates FP16/FP8 H256.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_256.cu	Instantiates BF16/INT8 H256.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_256.cu	Instantiates BF16/FP8 H256.
onnxruntime/test/contrib_ops/paged_attention_op_test.cc	Tests FP16/INT8 group-size-6 dispatch.

---

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