Commit f2c39fe
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?v=4&size=40){kind=avatar}
[CUDA] Add INT4 paged KV cache with per-channel scales (#32515)
### 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>1 parent 5894ba8 commit f2c39fe
23 files changed
Lines changed: 1698 additions & 151 deletions
File tree
- cmake
- docs
- contrib_ops/cuda
- onnxruntime
- contrib_ops
- cpu/bert
- cuda
- bert
- xqa
- core/graph/contrib_ops
- test/python/transformers
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
133 | 133 | | |
134 | 134 | | |
135 | 135 | | |
136 | | - | |
| 136 | + | |
137 | 137 | | |
138 | 138 | | |
139 | 139 | | |
| |||
1513 | 1513 | | |
1514 | 1514 | | |
1515 | 1515 | | |
1516 | | - | |
| 1516 | + | |
1517 | 1517 | | |
1518 | 1518 | | |
1519 | 1519 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
4804 | 4804 | | |
4805 | 4805 | | |
4806 | 4806 | | |
4807 | | - | |
| 4807 | + | |
4808 | 4808 | | |
4809 | | - | |
| 4809 | + | |
4810 | 4810 | | |
4811 | 4811 | | |
4812 | 4812 | | |
| |||
4839 | 4839 | | |
4840 | 4840 | | |
4841 | 4841 | | |
4842 | | - | |
| 4842 | + | |
4843 | 4843 | | |
4844 | | - | |
| 4844 | + | |
4845 | 4845 | | |
4846 | 4846 | | |
4847 | 4847 | | |
4848 | 4848 | | |
4849 | 4849 | | |
4850 | 4850 | | |
4851 | 4851 | | |
4852 | | - | |
| 4852 | + | |
4853 | 4853 | | |
4854 | 4854 | | |
4855 | 4855 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1104 | 1104 | | |
1105 | 1105 | | |
1106 | 1106 | | |
1107 | | - | |
| 1107 | + | |
1108 | 1108 | | |
1109 | 1109 | | |
1110 | 1110 | | |
| |||
1122 | 1122 | | |
1123 | 1123 | | |
1124 | 1124 | | |
1125 | | - | |
| 1125 | + | |
1126 | 1126 | | |
1127 | 1127 | | |
1128 | 1128 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
238 | 238 | | |
239 | 239 | | |
240 | 240 | | |
241 | | - | |
| 241 | + | |
242 | 242 | | |
243 | 243 | | |
244 | 244 | | |
| |||
466 | 466 | | |
467 | 467 | | |
468 | 468 | | |
469 | | - | |
| 469 | + | |
470 | 470 | | |
471 | 471 | | |
472 | 472 | | |
| |||
0 commit comments