onnxruntime
b1f76d58 - [CUDA] Add FP16-cache paged XQA decode for head size 256 (#32263)

Commit
28 days ago
[CUDA] Add FP16-cache paged XQA decode for head size 256 (#32263) 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% |
Author
Parents
Loading