[CUDA] Add INT4 paged KV cache with per-channel scales - #32515
Merged
Merged
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)
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
XQA routing is not proven by the tests, the norm/RoPE test disables INT4, and launcher documentation describes the wrong scale contract.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds build-gated packed INT4 KV-cache support to CUDA PagedAttention, including portable and XQA decode paths.
Changes:
- Adds INT4 packing, validation, cache I/O, and kernel registration.
- Adds FP16 H256 XQA decode/speculative-decode specializations.
- Adds documentation and CUDA-focused tests.
File summaries
| File | Description |
|---|---|
test_paged_attention_int4.py |
Adds INT4 correctness and regression tests. |
bert_defs.cc |
Extends schema for packed uint8 caches. |
cuda_contrib_kernels.cc |
Registers INT4 CUDA kernels. |
xqa_paged_spec_dec_fp16_int4_256.cu |
Instantiates speculative INT4 XQA. |
xqa_paged_loader.h |
Documents XQA launcher contracts. |
xqa_paged_loader.cu |
Dispatches INT4 XQA kernels. |
xqa_paged_fp16_int4_256.cu |
Instantiates decode INT4 XQA. |
xqa_loader.h |
Adds the INT4 quantization enum. |
mhaUtils.cuh |
Dequantizes packed grains into shared memory. |
mha.h |
Defines packed INT4 cache heads. |
mha_impl.cuh |
Adds INT4 synchronization points. |
int4_cache.cuh |
Implements packed-grain dequantization. |
paged_attention.cc |
Adds validation, routing, and workspace selection. |
paged_attention_impl.cu |
Implements INT4 cache writes and reads. |
paged_attention_helper.h |
Validates INT4 shapes and contracts. |
ContribOperators.md |
Updates generated operator documentation. |
paged_attention.md |
Documents INT4 behavior and support. |
Review details
- Files reviewed: 17/17 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I found one correctness concern in the INT4 XQA path, plus several test and documentation gaps. I left recommendations inline.
Prove XQA routing in tests, cover the INT4 cache writer, and correct the launcher scale contract. - Assert PagedAttention dispatch telemetry so an XQA test cannot pass through the portable fallback; gate XQA tests on SM80+. - Enable INT4 in the norm/RoPE ordering test and compare packed cache bytes. - Add an XQA/portable parity test for large PER_CHANNEL K scales over zero and nonzero codes, at the FP16 fold bound. - Skip paged decode when the query-token count exceeds the device grid-Y limit, which a metadata-bounded INT4 speculative batch can reach, and regress it. - Document the static FP32 PER_CHANNEL scale folding in both INT4 XQA launchers, the speculative INT4 path, the FP16 fold bound, and the default-ON build flag.
This was referenced Sep 9, 2026
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
One correctness concern remains after the latest updates.
Baiju Meswani (baijumeswani)
previously approved these changes
Sep 9, 2026
…32520) ### Description Stacked on #32515 — review that one first. The base branch is `tlwu/20260909/int4_kv_per_channel`, so the diff here is the single commit on top; it retargets to `main` once #32515 merges. Paged XQA folds a `PER_CHANNEL` K scale into the query and stores the product in `T`, while the portable kernel keeps the same product in an FP32 shared-memory tile. A large scale saturates FP16 there, and a zero cache code then turns that infinity into `NaN`, so the two backends can disagree on otherwise valid input: ``` q = 100, k_scale = 1000 -> XQA fold: inf (NaN once multiplied by a zero code) portable: 100000.0 ``` This addresses the correctness thread on #32515. INT4 is the most exposed because its scale spans `max|K| / 7` instead of `max|K| / 127`, roughly 18x larger for identical data, and INT4 XQA is only eligible with `PER_CHANNEL` scales — but the fold itself is shared with INT8 and FP8, so this is not INT4-specific. ### Approach `kCacheScale` is already applied as a scalar into `qkScale`, once per CTA and outside the K/V loop. So the fix needs no change in the inner loop: - Divide the fold by `max|k_scale|` (`PagedMaxAbsScaleKernel`, a single-block reduction launched on device so the step stays capturable). - Pass that maximum to XQA as its scalar K scale; the kernel multiplies it back into `qkScale`. The correction is exact, and the folded query is bounded by `max|q|` at any scale magnitude. INT4 dequantizes into FP16 shared memory, so `cacheElemSize == 2` and `isKVCacheQuantized` is false, which is why the scalar scale was previously ignored on that path. That predicate is now separate from "the cache elements are narrower than `T`". Both INT4 translation units define `XQA_PAGED_INT4`, so decode and speculative decode are covered. ### Cost No inner-loop work is added and shared memory is unchanged, so XQA eligibility is unaffected. The only new work is one small max reduction (1024 floats at H256) per node per step, amortized under CUDA-graph replay. ### Validation status **Not built — there is no CUDA toolkit on the machine this was written on. This needs a GPU build and a perf comparison against #32515 before it leaves draft.** Host-side numerical checks against an FP32 reference: | scale regime | current fold | this change | |---|---|---| | `q=100, k_scale=1000`, zero code | `NaN` | `18750.0`, exact | | `max|s|` 1e2 / 1e3 | 5.3e-4 / 2.3e-4 rel err | 1.7e-5 / 3.6e-5 rel err | | `max|s|` 1e4 / 1e5 | non-finite | 2.3e-4 / 2.0e-4 rel err | Accuracy improves at ordinary magnitudes too, since normalizing makes better use of the FP16 mantissa. `test_int4_xqa_large_per_channel_k_scale_matches_portable` was moved into the regime that previously produced `NaN`: it drives `q * k_scale` to 1e6 with every K code zero, which is the `inf * 0` path. The codes are zeroed deliberately — with nonzero codes a scale that large makes softmax one-hot, and the argmax is then fp16-sensitive and flaky. ### Reviewer notes - This changes INT8 and FP8 `PER_CHANNEL` XQA as well, since they share the fold. Those paths previously passed a null scalar scale; they now receive the normalizer. Please exercise them alongside INT4. - If the preference is to keep the blast radius on the new feature only, the normalizer can be passed for `uint8` caches alone, leaving INT8/FP8 exactly as today.
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I found two correctness concerns in the new scale normalization, along with several test, performance, and build-metadata concerns. Recommendations are included inline.
Route per-channel K decode through FP32 portable kernels to avoid normalized FP16 query underflow and scalar scale overflow. Remove normalization scratch and reduction, handle tiny cache-write scales without reciprocal overflow, and gate INT4 build info on CUDA. Add dispatch-verified INT4/INT8/FP8 dynamic-range, scale-value, and CUDA graph regressions. Document the correctness-first fallback and its loss of XQA acceleration.
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
INT4 XQA is unreachable under the current dispatch gates, and the new per-channel XQA test is guaranteed to fail.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
- Files reviewed: 22/22 changed files
- Comments generated: 2
- Review effort level: Balanced
The large per-channel K scale test still required XQA, but per-channel K decode now bypasses XQA, so the assertion could no longer hold. Invert it to pin the routing that replaced it and rename the test accordingly, reusing the tri-state expect_xqa convention already used by the speculative helper.
XQA takes one scalar K scale, so a PER_CHANNEL scale is folded into the query. Storing that product in fp16 saturates on a large scale and a zero cache code then turns the infinity into a NaN, so the fold had been disabled and INT4 XQA with it, costing 1.44-6.29x on decode and demoting INT8/FP8 per-channel decode that was already XQA-eligible. Normalize the fold by the power of two just above max|k_scale| and hand that normalizer to XQA as its scalar K scale. A power of two keeps both the division and the reapplication exact, and every normalized scale lands in (0, 1], so the fold cannot overflow for any finite table. fp16 still bounds the channel range the fold can hold to 24 binades below the largest scale. Calibrated tables sit far inside it, so this defaults on and ORT_ENABLE_XQA_PER_CHANNEL_KV=0 routes wider tables to the portable FP32 kernel.
An attention scale above one with a channel scale at FLT_MAX would overflow attention_scale * normalizer in fp32 and turn every logit into a NaN. That table spans one binade, so it stays on XQA and the exponent bound is what keeps it finite; assert it there rather than only through the per-channel opt-out.
Tianlei Wu (tianleiwu)
enabled auto-merge (squash)
September 10, 2026 04:47
This was referenced Sep 12, 2026
Open
Open
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
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
Adds an INT4 paged KV cache to the CUDA
PagedAttentionkernel.The cache stores signed 4-bit values packed two per byte in a
uint8tensor, 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 existingk_cache_dtype/v_cache_dtypeattributes, whose'int4'value was previously reserved and rejected during validation.Quantization reuses the existing static
PER_TENSOR/PER_CHANNELscales 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_CHANNELscales reach them by folding the channel scale into the query: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 abovemax|k_scale|, and that normalizer is handed to XQA as its scalar K scale, which multiplies it back intoqkScaleonce per CTA outside the K/V loop. A power of two rather thanmax|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| / 7rather thanmax|K| / 127, and INT4 XQA is only eligible withPER_CHANNELscales.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=0routesPER_CHANNELK decode and metadata-bounded speculative decode to the portable FP32 kernel for tables that do exceed it.PER_TENSORK 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_CHANNELscales 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:
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_CHANNELpaged KV cache is the baseline; only the KV cache format differs between arms, and the model weights are byte-identical.PER_CHANNELThe 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.75 pp — six times the candidate's delta — and rewrites a third of the generations. So INT4
PER_CHANNELis indistinguishable from INT8 on this task, and the resolution of the measurement is roughly ±1.5 pp, not ±0.2 pp.On GPQA-diamond (198 questions) the same INT4
PER_CHANNELmodel scored 153/198 against 150/198 for INT8.The MMLU-Pro run above was measured on this head, so it covers the normalized fold as shipped. An earlier run of the same suite against the unnormalized fold scored 661 for the same arm; the two differ by three questions, inside the noise the concurrency control demonstrates. The GPQA numbers predate the fold and describe the cache format rather than the kernel.
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 onDECODER_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_portableandtest_per_channel_k_keeps_int8_xqacompare XQA against the portable kernel through the opt-out, and the wide-range and non-finite scale tests pin portable behaviour withORT_ENABLE_XQA_PER_CHANNEL_KV=0. 31/31 pass on an H200, includingtest_xqa_large_attention_scale_and_k_scale, which drives an attention scale above one against aFLT_MAXchannel scale on the XQA path.test_paged_attention.pygains a finite-output assertion for overridden scale maxima, an XQA assertion for largePER_CHANNELK scales, and an opt-out test.Build configurations:
onnxruntime_USE_INT4_KV_CACHE=ONand=OFFeach build clean with zero warnings in the changed files, and-DUSE_INT4_KV_CACHE=1was confirmed present/absent on the actual compile lines.