[CUDA] Add Hadamard rotation and per-token KV scales to PagedAttention - #32521
Draft
Tianlei Wu (tianleiwu) wants to merge 2 commits into
Draft
Tianlei Wu (tianleiwu) wants to merge 2 commits into
Tianlei Wu (tianleiwu) wants to merge 2 commits into
Conversation
PagedAttention can now store its KV cache as signed 4-bit values packed two per byte into a uint8 tensor, halving the cache footprint relative to INT8. The cache head dimension becomes (head_size + 1) / 2, and the logical element type is named by the existing k_cache_dtype / v_cache_dtype attributes, whose 'int4' value was previously reserved and rejected. Quantization reuses the existing static PER_TENSOR / PER_CHANNEL scales on inputs 14/15, so no new quantization granularity or operator input is introduced. Decode also gains an XQA path. The loader dequantizes a packed grain into FP16 shared memory, and PER_CHANNEL needs no per-element scaling inside the kernel: the channel scale is folded into Q beforehand and applied to the output afterwards, which is exact because dequantization is linear and the channel dimension is contracted for K and free for V. The kernel therefore runs at unit scale, reusing the fold that INT8 PER_CHANNEL already relies on. Everything sits behind the existing onnxruntime_USE_INT4_KV_CACHE option, which defaults to OFF; feature-off builds compile unchanged and reject int4 caches during input validation. Measured on Qwen3.8-27B (H200, SM90a, 2048 generated tokens, milliseconds per target forward), XQA versus the portable paged-decode fallback: prompt 512, batch 1, 7 drafts: 37.45 -> 26.04 (1.44x) prompt 8192, batch 4, 7 drafts: 213.35 -> 59.09 (3.61x) prompt 32768, batch 4, 7 drafts: 546.25 -> 124.15 (4.40x) prompt 32768, batch 4, 0 drafts: 243.98 -> 38.80 (6.29x)
Builds on the INT4 paged KV cache by adding the two pieces the rotated INT4 configuration needs. PER_TOKEN quantization gives each cached token its own scale, computed from that token's own amax during the cache write and stored in new in-place scale caches (inputs 17/18, outputs 3/4) that reuse the slot_mapping and block_table of the KV caches. Unlike the static scales this granularity adapts to outliers, which is what makes 4-bit viable for a rotated cache. qk_rotation / v_rotation apply a per-head orthonormal Walsh-Hadamard transform: Q and K are rotated after QK-Norm and rotary, V is rotated before caching and the transform is inverted on the attention output. Because the transform is orthonormal it leaves attention mathematically unchanged while spreading each head's energy across channels, which shrinks the per-token quantization range. The rotation is incompatible with PER_CHANNEL on the rotated side, since a static per-channel scale is not preserved by mixing channels, so validation rejects that combination and the rotated path uses PER_TOKEN. The XQA INT4 loader now takes an optional per-token FP16 scale array; a null pointer keeps the PER_CHANNEL behaviour introduced in the previous commit, where the scale is folded into Q and the output instead.
Tianlei Wu (tianleiwu)
marked this pull request as draft
September 9, 2026 23:55
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Sep 10, 2026
### Description Adds an INT4 paged KV cache to the CUDA `PagedAttention` kernel. The cache stores signed 4-bit values packed two per byte in a `uint8` tensor, halving the cache footprint relative to INT8. The cache head dimension becomes `(head_size + 1) / 2`, and the logical element type is named by the existing `k_cache_dtype` / `v_cache_dtype` attributes, whose `'int4'` value was previously reserved and rejected during validation. Quantization reuses the existing static `PER_TENSOR` / `PER_CHANNEL` scales on inputs 14/15, so this introduces **no new quantization granularity and no new operator input**. Reads and writes go through the portable paged decode / gather paths, which fold scales in FP32. #### XQA decode FP16 INT4 XQA decode and speculative-decode specializations are included, and `PER_CHANNEL` scales reach them by folding the channel scale into the query: ``` score = Q · (K_int · s_c) = (Q · s_c) · K_int out_c = s_c · Σ_t p_t · V_int[t, c] ``` The fold is exact in real arithmetic, but the product is held in FP16, so a large channel scale saturates it and a zero cache code then turns the infinity into a `NaN`. The fold is therefore divided by the power of two just above `max|k_scale|`, and that normalizer is handed to XQA as its scalar K scale, which multiplies it back into `qkScale` once per CTA outside the K/V loop. A power of two rather than `max|k_scale|` itself matters: the division and the reapplication are both exact, so the fold adds no rounding of its own, and every normalized scale lands in `(0, 1]`, so `|Q * s| <= |Q|` and the store cannot overflow for any finite table. The reduction is one block on the compute stream, so the path stays CUDA-graph capturable and re-reads the table on every replay. INT4 is the most exposed case because its scale spans `max|K| / 7` rather than `max|K| / 127`, and INT4 XQA is only eligible with `PER_CHANNEL` scales. **Limit and opt-out.** FP16 spans about 40 binades, and an overflow-free normalizer must be at least `max|k_scale|`, so channels more than 24 binades below the largest flush to zero in the folded query. Calibrated tables sit far inside that budget — across the 128 per-(head, side) tables of a Qwen3.8-27B INT4 export the widest spans 4.9 binades — so this defaults on. `ORT_ENABLE_XQA_PER_CHANNEL_KV=0` routes `PER_CHANNEL` K decode and metadata-bounded speculative decode to the portable FP32 kernel for tables that do exceed it. `PER_TENSOR` K is unaffected either way. #### Build flag Everything is behind `onnxruntime_USE_INT4_KV_CACHE`. Builds with the option off compile unchanged and reject int4 caches during input validation. ### Motivation and Context A 4-bit paged KV cache halves long-context cache memory versus INT8. The paged cache allocation on Qwen3.8-27B at 1024 blocks measures 8,192 MiB with an INT8 cache; because `PER_CHANNEL` scales are static initializers rather than a scale cache, INT4 stores exactly half of that and adds nothing back. #### Decode performance Measured on Qwen3.8-27B (H200, SM90a, 2048 generated tokens, milliseconds per target forward), INT4 XQA versus the portable fallback: | prompt | batch | drafts | portable | XQA | speedup | |---:|---:|---:|---:|---:|---:| | 512 | 1 | 7 | 37.45 | 26.04 | 1.44x | | 8192 | 4 | 7 | 213.35 | 59.09 | 3.61x | | 32768 | 4 | 7 | 546.25 | 124.15 | 4.40x | | 32768 | 4 | 0 | 243.98 | 38.80 | 6.29x | Target-forward counts are identical between the two arms for the 8192 and both 32768 rows, so those are like-for-like. An unchanged control arm measured across the same two sessions drifted at most 4.8%. #### Accuracy MMLU-Pro, 800 questions, Qwen3.8-27B INT4 weights with DFlash2 speculative decoding (N=7), greedy with natural EOS, one H200 held exclusively for the run. The INT8 `PER_CHANNEL` paged KV cache is the baseline; only the KV cache format differs between arms, and the model weights are byte-identical. | arm | correct / 800 | accuracy | paired vs baseline | |---|---:|---:|---| | INT8 KV (baseline) | 659 | 82.375% | — | | baseline, identical rerun | 659 | 82.375% | 0 of 800 generations changed | | baseline at concurrency 4 (neutral control) | 663 | 82.875% | +0.50 pp, 284 generations changed | | **INT4 KV, `PER_CHANNEL`** | **661** | **82.625%** | **+0.25 pp, 27 wins / 25 losses, exact McNemar p = 0.89** | The two control rows are what make this readable. The identical rerun is token-exact, so the harness and engine are deterministic and any difference between arms is attributable to the arm. The neutral control changes nothing but request concurrency, yet it moves accuracy by +0.50 pp and rewrites a third of the generations — a larger swing than INT4 produces. So INT4 `PER_CHANNEL` is indistinguishable from INT8 on this task, and the resolution of the measurement is roughly ±2 pp, not ±0.2 pp. On GPQA-diamond (198 questions) the same INT4 `PER_CHANNEL` model scored 153/198 against 150/198 for INT8. These runs used the XQA decode path with the per-channel fold, which is the configuration this PR ships. This is the first of two stacked changes; #32521 adds per-token scales and Hadamard rotation on top. ### Testing `onnxruntime/test/python/transformers/test_paged_attention_int4.py` (new) covers packing/padding layout, prefill, chunked prefill, decode, split-KV, speculative decode, CUDA-graph replay, bfloat16 activations, scale extremes including subnormal and near-overflow scale tables, dispatch assertions that per-channel INT4 decode and speculative decode land on `DECODER_ATTENTION`, XQA fallback for ineligible static scales, invalid-contract rejection, and an int8/fp8 regression check that the existing cache types are unaffected. Dispatch is asserted, not inferred: the INT4 decode, speculative-decode and CUDA-graph tests require `SdpaKernel=XQA`, `test_int4_per_channel_xqa_matches_portable` and `test_per_channel_k_keeps_int8_xqa` compare XQA against the portable kernel through the opt-out, and the wide-range and non-finite scale tests pin portable behaviour with `ORT_ENABLE_XQA_PER_CHANNEL_KV=0`. **31/31 pass on an H200**, including `test_xqa_large_attention_scale_and_k_scale`, which drives an attention scale above one against a `FLT_MAX` channel scale on the XQA path. `test_paged_attention.py` gains a finite-output assertion for overridden scale maxima, an XQA assertion for large `PER_CHANNEL` K scales, and an opt-out test. Build configurations: `onnxruntime_USE_INT4_KV_CACHE=ON` and `=OFF` each build clean with zero warnings in the changed files, and `-DUSE_INT4_KV_CACHE=1` was confirmed present/absent on the actual compile lines. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Stacked on #32515 — please review that one first; this PR targets its branch, not
main.Adds the two pieces the rotated INT4 configuration needs on top of the per-channel INT4 paged KV cache.
PER_TOKEN quantization. Each cached token gets its own scale, computed from that token's own amax during the cache write and stored in new in-place scale caches (inputs 17/18, outputs 3/4) that reuse the
slot_mappingandblock_tableof the KV caches. Unlike a static scale this granularity adapts to outliers, which is what makes 4-bit viable for a rotated cache.Hadamard rotation.
qk_rotation/v_rotationapply a per-head orthonormal Walsh-Hadamard transform: Q and K are rotated after QK-Norm and rotary, V is rotated before caching, and the transform is inverted on the attention output. Because the transform is orthonormal it leaves attention mathematically unchanged ((QH)(KH)^T = QK^Tandsoftmax(S)(VH) = (softmax(S)V)H) while spreading each head's energy across channels, which shrinks the per-token quantization range.The rotation is incompatible with
PER_CHANNELon the rotated side, because a static per-channel scale is not preserved by mixing channels, so validation rejects that combination and the rotated path usesPER_TOKEN.The XQA INT4 loader now takes an optional per-token FP16 scale array. A null pointer keeps the
PER_CHANNELbehaviour from #32515, where the scale is folded into Q and the output and the kernel runs at unit scale, so that path is unchanged bit-for-bit.Motivation and Context
INT4 KV halves the cache against INT8 and is what makes a 256K context fit on 24 GB alongside a speculative drafter. Per-token scales plus the rotation are the accuracy insurance for that: on a 27B model the rotation lifts INT4 K SNR from 11.67 dB to 19.66 dB in simulation, and the two 4-bit variants land within noise of the INT8 baseline on GPQA-diamond and MMLU-Pro.
Notes for reviewers
onnxruntime_USE_INT4_KV_CACHECMake option (default OFF). With the option off,int4_xqa_eligibleis aconstexpr falseand the dispatch is unchanged.allow_per_tokenargument when parsing the quantization-type attribute, so aPER_TOKENnode reports the existing "not supported yet" path instead of claiming the attribute value is invalid. The WebGPU kernel still rejects any quantized cache.docs/ContribOperators.mdis generated; its diff is the schema change rendered bytools/python/gen_contrib_doc.py.onnxruntime_USE_INT4_KV_CACHE. The WebGPU EP was not built locally.