[CUDA] PagedAttention: quantized KV cache, XQA decode, MLA, QK-Norm and head sink

[tianleiwu]

Description

PagedAttention is ORT's continuous-batching attention operator, but on main it only supports
FP16/BF16 caches with RoPE and softcap, has no paged decode kernel, and forces a device→host
synchronization on every node on every step (which makes it uncapturable by CUDA graphs). This PR
brings it to feature parity with GroupQueryAttention for the popular LLM families and adds the
paging and latent-cache primitives that serving frameworks need, additively — every model valid
under the shipped com.microsoft::PagedAttention opset-1 schema keeps working unchanged.

The design rationale, the compatibility invariant, and the alternatives that were considered and
rejected are written up in the new design document
docs/contrib_ops/cuda/paged_attention.md; the section
numbers referenced below point into it.

Summary of Changes

Schema (bert_defs.cc, docs/ContribOperators.md)

All additions are trailing optional inputs, new attributes whose defaults reproduce current
behavior, or widened type constraints (§4).

New input	Idx	Purpose
slot_mapping	10	Explicit per-token cache slot, so the scheduler owns placement instead of the kernel re-deriving it (§5)
head_sink	11	Attention sink / smooth softmax, matching GQA (§6)
q_norm_weight / k_norm_weight	12, 13	Fused QK-RMSNorm (Qwen3, gpt-oss) (§7)
k_scale / v_scale	14, 15	Per-tensor or per-channel dequantization scales for a quantized cache (§8)
attention_metadata	16	Optional CPU input carrying replay-wide upper bounds [max_query_len, max_kv_len], which removes the per-node per-step D→H sync (§4.7)

New attribute	Default	Purpose
qk_norm_epsilon	1e-6	Epsilon for the fused QK-Norm
k_quant_type / v_quant_type	NONE	NONE | PER_TENSOR | PER_CHANNEL
k_cache_dtype / v_cache_dtype	""	Logical cache element type, named after the ONNX type it denotes
kv_cache_layout	SEPARATE	SEPARATE | LATENT (absorbed MLA: one cache, no value/value_cache)
v_head_size	0	Narrower V head, LATENT only (DeepSeek-V3 uses 576/512)
rotary_offset	0	Applies RoPE to [rotary_offset, rotary_offset + rotary_dim) so MLA can rotate only the positional suffix

key_cache / value_cache move from T to a new T_CACHE constraint (float16, bfloat16,
int8, float8e4m3fn), and value_cache / value_cache_out become optional so a LATENT node can
omit them. Shape inference now takes the cache element type from inputs 3/4 rather than from query,
which was wrong for a quantized cache.

CUDA kernels (paged_attention_impl.cu, paged_attention.cc/.h, paged_attention_helper.h)

• Paged decode kernel (LaunchPagedDecodeAttention) — split-KV, block-table-aware decode with
  native head-sink, softcap, sliding-window and on-the-fly cache dequantization.
• XQA paged decode (onnxruntime/contrib_ops/cuda/bert/xqa/) — TensorRT-LLM's XQA kernels
  extended to the paged block layout: 8 new translation units
  xqa_paged_{fp16,bf16}_{int8,fp8}_{64,128}.cu plus a shared paged loader. Selected for
  quantized-cache decode.
• Quantized paged cache — ReshapeAndCache quantizes on write; all read paths dequantize with
  k_scale/v_scale under PER_TENSOR or PER_CHANNEL granularity.
• ApplyHeadSink — exact post-hoc LSE rescale (1/(1+exp(s_h − lse))) applied after the
  quantized/unquantized branch, so no backend can silently drop the sink (§6).
• QkNormRotaryTNH — fuses QK-RMSNorm, RoPE (with rotary_offset) and the packed-QKV unpack
  into one pass.
• Absorbed MLA (PagedLatentAttentionKernel / LatentAttention) — single latent cache, V read
  as the leading v_head_size channels of the same row that supplies K (§12).
• CUDA-graph safety — backend dispatch, grid sizing and workspace extents now come from static
  shapes and the block_table.shape[1] * block_size capacity bound; per-step quantities are read on
  device. The unconditional cudaStreamSynchronize is gone from the capturable path (§4.7).
• int8→fp16 conversion fast path (xqa/utils.cuh, cvtS8x4ToF16x4) — replaces a scalar
  I2F loop with a prmt + sub.f16x2 sequence (5 full-rate instructions per 4 elements, bit
  identical). Shared with the non-paged GQA loader.
• Kernel registration is now <T, T_CACHE>-typed; FP8 combinations are behind
  USE_FP8_KV_CACHE && !DISABLE_FLOAT8_TYPES.

GQA bug fix (flash_api.{h,cc}, group_query_attention_impl.cu)

mha_fwd had constexpr void* head_sink = nullptr; hardcoded inside it, and
FlashAttentionAndQuantizeKV — the only GQA prompt path taken when the KV cache is quantized —
called it. So for gpt-oss with an INT8/FP8 KV cache, the attention sinks were silently dropped for
the entire prompt on every layer while decode stayed correct. mha_fwd now takes head_sink and
GQA forwards it.

Op-level prefill error drops 0.074829 → 0.000122. On gpt-oss-20b (int4 body, INT8 per-channel
KV), MMLU-Pro-800 goes 0.6175 (494/800) → 0.7200 (576/800). Existing CI missed this because
atol["int8_fp16"] = 1e-1 in test_gqa.py is ~800× wider than the post-fix error.

Tooling and docs

• symbolic_shape_infer.py: correct output width for packed-QKV and LATENT nodes, and cache
  outputs typed from the cache inputs.
• New docs/contrib_ops/cuda/paged_attention.md design document; regenerated ContribOperators.md
  and OperatorKernels.md.

Testing

test_paged_attention_cuda.py grows from a smoke test to ~2k lines / 198 cases, with new suites for
features (slot_mapping, head sink, QK-Norm), quantized cache (int8/fp8 × per-tensor/per-channel),
the paged decode kernel, the XQA decode path, attention_metadata, and MLA — each against a PyTorch
reference.

python onnxruntime/test/python/transformers/test_paged_attention_cuda.py   # 198 passed
python onnxruntime/test/python/transformers/test_gqa.py -k xqa             # 714 passed

The GQA suite is included because the int8 conversion fast path is shared with the non-paged loader.

Backward compatibility. A node with none of the new inputs/attributes takes exactly the code path
it does today: T_CACHE == T, value_cache present, kv_cache_layout == SEPARATE, all quantization
NONE. The compatibility invariant is stated normatively in §4.2.

Experimental Results

Measured on gpt-oss-20b, H200. E2E numbers are driven through onnxruntime-genai; the CUDA-graph and
engine-side plumbing they depend on is not part of this PR — they are included to show what the
operator-side changes enable, not as a claim about this diff alone.

Paged decode kernel (isolated, nh=64 / kvh=8 / hs=64 / block=256)

XQA on/off at b=8, ctx=4096, per decode call:

cache	before	after
int8 PER_TENSOR	2315 µs	122 µs
int8 PER_CHANNEL	2339 µs	128 µs
fp8 PER_TENSOR	1633 µs	57 µs
fp8 PER_CHANNEL	1566 µs	57 µs

Before XQA the quantized paths were ~2.5× slower than fp16 — the generic kernel was the bottleneck,
not the KV bytes.

The cvtS8x4ToF16x4 conversion path then closes the residual int8-vs-fp8 gap (nsys median, SASS goes
from 3928 to 3592 instructions with 192 → 0 I2F):

ctx	batch	int8 before	int8 after	gain	fp8
1024	32	19.91 µs	13.60 µs	−31.7%	12.64 µs
4096	8	24.32 µs	17.41 µs	−28.4%	16.48 µs
4096	32	66.62 µs	42.66 µs	−36.0%	41.98 µs
16384	8	78.11 µs	52.58 µs	−32.7%	49.18 µs
16384	32	244.71 µs	162.27 µs	−33.7%	176.70 µs

The int8/fp8 gap goes from up to +59% down to ≤ 7.6% (int8 is faster at the largest config), so the
two cache formats can now be chosen on accuracy grounds.

End-to-end decode throughput (mxfp4 body, INT8 KV, prompt 128 / new 256)

batch	baseline	+ attention_metadata (no sync)	+ CUDA graphs	total
1	242.7 tok/s	263.8	305.7	+26.0%
8	796.1	811.5	870.5	+9.3%
32	2635.7	2682.0	2879.2	+9.2%

Including XQA, batch-1 int8 decode goes 199.2 → 305.7 tok/s (+53.5%). After this work attention
is no longer the bottleneck at b=1 — the MoE GEMMs and 49 MatMulNBits nodes dominate the step, with
XQA at 24 × 8.6 µs.

PagedAttention vs GroupQueryAttention, matched models

Two models built from the identical recipe (int4 body, INT8 per-channel KV, identical
num_heads/kv_num_heads/scale/window/rotary, byte-identical weight file), differing only in the
attention operator. Greedy generation on this stack is bit-reproducible (0/198 discordance across
replicates), so there is no sampling noise to subtract.

	GQA	PagedAttention	read as
MMLU-Pro-800	0.7200 (576/800)	0.7163 (573/800)	+0.4 pp, 3 questions
GPQA-diamond	0.6061 (120/198)	0.6212 (123/198)	−1.5 pp, 3 questions

The two benchmarks disagree in direction and both deltas are 3 questions: equivalent within noise.
(An apparent +8 pp advantage for paged in earlier runs turned out to be the GQA sink bug fixed above,
seen from the other side.)

config	GQA tok/s	paged tok/s	delta
b=1, p=128, n=256	374.4	376.3	+0.5%
b=2, p=4096, n=256	684.3	686.6	+0.3%
b=8, p=128, n=256	1792.7	1683.8	−6.1%
b=32, p=128, n=256	5285.9	4109.5	−22.3%

Peak device memory at matched KV capacity agrees to within 24 MiB (0.16%) from 16k to 128k
max_length — paged costs nothing extra, and its advantage is structural (a shared pool sized to
aggregate demand rather than batch × max_length).

The b=32 gap was profiled with nsys --cuda-graph-trace=node: the captured graph body is at parity
with GQA's eager model pass (6.220 ms vs ~6.2 ms) and the entire regression is a 2.384 ms
search/sampling tail, which the onnxruntime-genai Engine runs once per request rather than once
per batch. It is not attributable to this operator, and a partial engine-side fix already recovers
b=32 to 4517 tok/s.

Follow-ups (not in this PR)

• attention_bias and output_qk (§10, §11) — schema slots reserved, kernels deferred.
• Sub-byte (int4 / float4e2m1) packed caches — attribute vocabulary reserved and rejected at
  validation until a backend exists (§21.4).
• .Alias(3, 1).Alias(4, 2) on the kernel def, so a non-aliasing allocation plan fails at partition
  time instead of run time (§4.4).
• Re-tightening atol["int8_fp16"] in test_gqa.py now that the sink bug is fixed.

[Copilot]

Pull request overview

This PR extends the CUDA contrib com.microsoft::PagedAttention operator to cover additional modern serving features (absorbed MLA / latent KV layout, quantized KV cache, QK RMSNorm, head sinks, explicit slot mapping, and replay-stable host metadata for CUDA Graph friendliness) and adds a paged-KV XQA decode path for quantized caches.

Changes:

• Expand the PagedAttention op schema + type/shape inference to support quantized caches (T_CACHE + scales), optional/conditional cache/value IO for kv_cache_layout=LATENT, and new optional inputs (slot mapping, head sink, QK-norm, metadata).
• Update the CUDA kernel implementation to handle quantized caches, new prologues/epilogues, backend selection (Flash/MEA/paged-decode/XQA/latent), and add paged-XQA translation units/loaders.
• Update documentation and Python symbolic shape inference for the new operator behavior.

Reviewed changes

Copilot reviewed 27 out of 29 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
onnxruntime/python/tools/symbolic_shape_infer.py	Adds special-case output/cache inference for packed QKV and LATENT layout.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Updates PagedAttention schema + shape inference for new inputs/attrs, quantized cache types, and LATENT output rules.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers additional typed CUDA kernels for PagedAttention (activation + cache dtype pairs).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h	Declares paged-KV XQA decode launcher and shared-memory query helper.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu	Implements dispatcher for paged-XQA decode by head size / cache quant type / bf16.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader_impl.cuh	Shared TU body for generating paged-XQA kernels across group sizes and dtypes.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_impl_gen.cuh	Instantiation template mirroring XQA contiguous kernels but binding paged-KV entry points.
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_64.cu	Instantiates fp16+int8 paged-XQA kernel (head=64).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_int8_128.cu	Instantiates fp16+int8 paged-XQA kernel (head=128).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_64.cu	Instantiates fp16+fp8 paged-XQA kernel (head=64, gated).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_fp16_fp8_128.cu	Instantiates fp16+fp8 paged-XQA kernel (head=128, gated).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_64.cu	Instantiates bf16+int8 paged-XQA kernel (head=64).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_int8_128.cu	Instantiates bf16+int8 paged-XQA kernel (head=128).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_64.cu	Instantiates bf16+fp8 paged-XQA kernel (head=64, gated).
onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_bf16_fp8_128.cu	Instantiates bf16+fp8 paged-XQA kernel (head=128, gated).
onnxruntime/contrib_ops/cuda/bert/xqa/utils.cuh	Adds int8→half conversion fastpath; adjusts constants for paged kernels.
onnxruntime/contrib_ops/cuda/bert/xqa/mhaUtils.cuh	Extends paged KV cache metadata (extraSeqLen) to unify seq-len handling.
onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh	Fixes constexpr-divisor hazards; plumbs extraSeqLen into paged KV cache list.
onnxruntime/contrib_ops/cuda/bert/paged_attention.h	Extends kernel class state for new attrs/options + XQA shared-mem caching.
onnxruntime/contrib_ops/cuda/bert/paged_attention.cc	Implements new backend selection, metadata bounds, quantized cache wiring, and XQA/latent paths.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.h	Adds helpers for paged decode and latent shared-memory sizing; templates on cache dtype.
onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu	Adds quantized cache read/write, fused QK-norm+RoPE prologue, paged decode, latent backend, XQA paged decode, and updated gathers.
onnxruntime/contrib_ops/cuda/bert/paged_attention_helper.h	Validates new inputs/attrs (LATENT, slot mapping, head sink, QK-norm, cache quantization/data types, metadata).
onnxruntime/contrib_ops/cuda/bert/attention_data.h	Extends PagedAttentionData for cache dtype, scales, decode workspaces, and XQA scratch.
onnxruntime/contrib_ops/cpu/bert/attention_parameters.h	Extends PagedAttentionParameters with sink/QK-norm/quant/LATENT/offset fields.
onnxruntime/contrib_ops/cpu/bert/attention_common.h	Adds KV cache logical dtype enum + parsing helpers.
docs/ContribOperators.md	Updates PagedAttention docs to reflect new schema surface area.

[Copilot]

Pull request overview

Copilot reviewed 31 out of 33 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

onnxruntime/python/tools/symbolic_shape_infer.py:2572

• In _infer_PagedAttention, the packed-QKV divisor uses (kv_num_heads or 0). If the node is malformed and kv_num_heads is missing/None, this makes the divisor num_heads, so shape inference can incorrectly derive a head_size and emit an invalid output width instead of falling back to generic propagation (contradicting the comment about missing attributes). Guard the packed-QKV path so kv_num_heads must be present/positive before computing the divisor.