[CUDA] Serve quantized KV caches with XQA: attention sinks, independent and per-channel scales

[tianleiwu]

Description

Extends the XQA decode kernel to serve the quantized (INT8/FP8) KV-cache configurations that
previously fell back to dequantize-the-whole-cache-then-Flash-Attention. Four changes, in
contrib_ops/cuda/bert/:

1. Attention sinks on quantized XQA. use_xqa_attention_sinks no longer requires
   !is_inputs_quantized. The sink is folded into the softmax row sum by the same code the
   non-quantized path uses, and the dequant scale is already applied to the scores before the row
   max/sum, so sink and score live in the same domain — true for both the single-block and the
   multi-block (Flash Decoding) reduction, where the sink is added to the merged row sum.

2. Independent k_scale / v_scale. XQA previously required the two scale tensors to be the
   same pointer, which no builder-emitted model satisfies. The kernel now keeps them apart:
   k_scale → qkScale (applied to Q·Kᵀ), v_scale → voScale (applied to the P·V accumulator).

3. PER_CHANNEL scales folded out of the kernel. XQA only accepts a scalar dequant scale.
   But dequantization is linear and the channel index is the contracted dim for K and the free
   dim for V:

   scores_t = sum_d q_d * (k_td * sk_d) = sum_d (q_d * sk_d) * k_td
   out_d    = sum_t p_t * (v_td * sv_d) = (sum_t p_t * v_td) * sv_d
   

   so the per-channel scales can be moved out of the kernel exactly: pre-scale Q by sk,
   post-scale the attention output by sv. p_t is unchanged because softmax only ever sees the
   (already correct) scores, and every step after the P·V accumulation is linear, so attention
   sinks, sliding window and the multi-block reduction all stay valid. XQA reads a null scale
   pointer as "scale = 1", which is how the caller signals it folded the scale itself.

   The Q side costs nothing extra: UnpackRoPEAppend already loads Q, applies RoPE and stores it to
   the very scratch buffer XQA reads, so it takes a q_fold_scale argument and applies the multiply
   to the registers it is about to store. Only the V side needs its own pass, since it consumes
   XQA's output. Both are O(num_heads * head_size) — O(1) in context length, versus the
   O(context) full-cache dequant they replace.

4. Vectorized / length-aware dequant fallback for what XQA still cannot serve (INT4,
   unsupported head_size/group size). The fallback dequant kernel now uses 32-byte vector loads,
   hoists per-channel scales out of the sequence loop, maps (batch, head) to blockIdx.y so the
   inner addressing is 32-bit, and walks only the rows inside the sequence instead of the padded
   cache capacity.

docs/contrib_ops/cuda/gqa.md is updated to match: quantized-decode description, XQA eligibility
conditions, backend-selection table, and the parity-test list.

Motivation and Context

The quantized KV-cache decode path was ~2.8x slower than an FP16 KV cache at 16k context, and
the gap grew with context length because the fallback rebuilt the entire dequantized K+V cache on
every decode step. The binding gate was attention sinks: a model carrying head_sink sets
use_smooth_softmax, which combined with a quantized cache disqualified XQA regardless of scale
granularity, head size or group size. Scale-granularity workarounds do not help — all quantized
configurations previously landed within 0.5% of each other, because none of them reached XQA.

Decode latency, batch 1, 64 generated tokens, gpt-oss-20b geometry (24 layers, head_size 64,
group 8), H200 / SM90, median of 5, idle GPU:

KV cache	1k	4k	16k	vs FP16 KV @16k
FP16 (control)	2.38 -> 2.370	2.44 -> 2.434	2.57 -> 2.573	—
INT8 per-tensor	3.22 -> 2.421	4.04 -> 2.490	7.42 -> 2.626	+2.1%
INT8 per-channel	3.21 -> 2.476	4.07 -> 2.542	7.42 -> 2.683	+4.3%
FP8 per-tensor	3.18 -> 2.390	4.08 -> 2.449	7.39 -> 2.561	−0.5%
FP8 per-channel	— -> 2.443	— -> 2.492	— -> 2.611	+1.5%

(ms/token, before -> after. The FP16 control is unchanged, so the comparison is sound.) Quantized
KV decode goes from −66% to −4% versus an FP16 cache at 16k, while keeping its −26% TTFT
(749 ms vs 1017 ms @16k). Scale granularity now costs ~0.05 ms/token (2% at 16k), so per-channel —
by far the most accurate INT8 option — no longer has to be traded away for decode speed.

For configurations XQA still cannot serve, change 4 alone gives 7.44 -> 4.61 ms/token @16k (1.61x);
its residual overhead over the FP16 non-XQA Flash path is now 0.44 ms/token (was 3.40), i.e. roughly
75% of theoretical HBM bandwidth for the traffic it must move.

Testing

• onnxruntime/test/python/transformers/test_gqa.py: 865 passed, 384 of them new. Three new
  classes, each asserting both numerics and that the op actually reaches XQA
  (SdpaKernel == "XQA"), so a silent fallback fails the test rather than passing quietly:

  • TestXQAQuantizedHeadSinkParity (96) — quantized decode with an attention sink.
  • TestXQASeparateKVScaleParity (144) — distinct k_scale / v_scale tensors.
  • TestXQAPerChannelScaleParity (144) — per-channel K/V scales folded into Q and the output.

  Combined coverage: INT8 + FP8 x fp16 + bf16 x head_size 64/128/256 x group 4/8 x global /
  sliding-window / multi-block (past 512) x sink on/off x rotary on/off x prepacked vs
  runtime-converted sink.

• onnxruntime_provider_test --gtest_filter='*GroupQueryAttention*': 63 passed.

• The q_fold_scale fusion in change 3 was verified to be a numerical no-op: op-level max-abs and
  RMS error against the fp32 reference are bit-identical to the unfused two-pass version across
  head_size 64/128 x past 4/512.

• End to end, a gpt-oss-20b build with an INT8 per-channel KV cache produces a token-identical
  128-token greedy continuation to the same model with an FP16 KV cache.

• lintrunner clean.

[Copilot]

Pull request overview

This PR extends the contrib CUDA GroupQueryAttention (GQA) “XQA” decode path to efficiently serve quantized KV caches (INT8/FP8) that previously fell back to full-cache dequantization + Flash-Attention, adding support for attention sinks, independent K/V scales, and exact per-channel scale folding.

Changes:

• Enable attention sinks (smooth softmax via head_sink) for quantized XQA decode and thread sink/scales correctly through both single- and multi-block reductions.
• Support independent k_scale/v_scale, and serve PER_CHANNEL quantization by folding K scale into Q and V scale into the output (exact linear transform).
• Optimize the non-XQA quantized decode fallback with a vectorized, length-aware dequantization path and expand Python parity/dispatch coverage + docs.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_gqa.py	Adds new parity+dispatch tests for quantized XQA with sinks, separate K/V scales, and per-channel folding.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h	Updates XQA loader API to accept separate K/V scale pointers and document null-scale semantics.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu	Threads separate K/V scale pointers through FP16 XQA dispatch.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8.cu	Updates INT8 XQA dispatch signature to include sinks + separate K/V scales.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8_impl.cuh	Threads sinks + separate K/V scales into INT8 group-size specializations.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh	Enables sinks on quantized paths and threads separate K/V scales throughout.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_fp8_impl.cuh	Threads sinks + separate K/V scales into FP8 group-size specializations.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_64.cu	Updates explicit instantiation signature to match loader changes (H=64).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_128.cu	Updates explicit instantiation signature to match loader changes (H=128).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_256.cu	Updates explicit instantiation signature to match loader changes (H=256).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu	Threads separate K/V scales + sinks through BF16 XQA dispatch (incl. INT8 path).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8.cu	Updates BF16 INT8 XQA dispatch signature to include sinks + separate K/V scales.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8_impl.cuh	Threads sinks + separate K/V scales into BF16 INT8 group-size specializations.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh	Enables sinks on quantized BF16 paths and threads separate K/V scales throughout.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_fp8_impl.cuh	Threads sinks + separate K/V scales into BF16 FP8 group-size specializations.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_64.cu	Updates explicit instantiation signature to match loader changes (BF16 H=64).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_128.cu	Updates explicit instantiation signature to match loader changes (BF16 H=128).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_256.cu	Updates explicit instantiation signature to match loader changes (BF16 H=256).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh	Updates generated launch plumbing to pass separate K/V scale pointers.
onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh	Implements separate K/V scale handling and null-scale semantics in the kernel.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Updates XQA eligibility: allow sinks on quantized paths; relax scale restrictions; accept PER_CHANNEL for XQA via folding.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh	Adds q_fold_scale support to fold per-channel K scale into Q during preprocessing.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh	Adds vectorized/length-aware dequant fast path and output scaling kernel for per-channel V folding.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h	Reserves Q scratch buffer for per-channel K folding on XQA.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu	Wires per-channel folding into ExtremeDecoding (XQA) and length-aware dequant into fallback path.
docs/contrib_ops/cuda/gqa.md	Updates quantized-decode/XQA eligibility documentation and test inventory.