[Feature] Triton INT4 / INT2 per-token-head KV cache quantization by JartX · Pull Request #40633 · vllm-project/vllm

JartX · 2026-04-22T16:23:03Z

This PR is stacked on #40631, so please merge that one first. It is the feature half of the split from #39074

The main change here is the addition of two new KV cache quantization modes, both using per-(token,head) dynamic scales:

INT4_PER_TOKEN_HEAD: packs 2 int4 values per byte. It applies a random Hadamard pre-rotation to make the inputs more Gaussian, then uses asymmetric quantization with a 4-bit zero-point steganographically stored in the low 4 bits of the float32 scale mantissa. The asymmetry is corrected on the score side with a sum(Q) * zp term.

INT2_PER_TOKEN_HEAD: packs 4 int2 values per byte. It uses a full Hadamard pre-rotation and a Lloyd-Max centroid lookup for dequantization, with 4 centroids fitted over N(0,1). The d1.5 scale factor is folded in at write time.

On the read side, these modes do not fit cleanly into the common kernel. The packed layout introduces several structural differences from the non-packed path: the Q split into 2 or 4 interleaved streams depending on PACKING_FACTOR, packed KV loads, inline unpacking and centroid lookup, split-dot score computation with optional zero-point correction, and per-stream accumulators. Because of that, this PR adds a dedicated _attn_packed kernel that handles both modes behind a single PACKING_FACTOR: tl.constexpr switch. Internally, that still means a fair amount of mode-specific logic, including 13 internal branches just to separate the INT4 and INT2 paths.

To keep the shared kernel clean, dispatch for these modes now lives behind a small factory registry under vllm/v1/attention/ops/triton_quant_kv/.

QuantKVFactory is the base protocol for a single KVQuantMode. It owns the write path (reshape_and_cache) and can optionally own the read path (unified_attention) as well. The "factory" naming comes from feedback from @bringlein: implementations are lazy-imported, only instantiated on first use, and add no overhead for modes that are never used.

The sub-byte packed modes override both read and write paths. INT8 and FP8 per-token-head override only the write path and continue using the common kernel’s existing constexpr branches on the read side. This also moves the write kernel for INT8/FP8 per-token-head out of triton_reshape_and_cache_flash.py and into the factory module, so each mode now has a single home instead of splitting write and read logic across different files.

The factory also reuses most of the same helper logic as the common kernel: 9 of the 12 helpers from triton_attention_helpers.py, covering the prologue, softmax bookkeeping, ALiBi, softcap, qq_bias, sliding window masking, and segment reduction, plus reduce_segments imported across files. The only pieces that remain bespoke are the structural parts required by the packed read path.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

mergify · 2026-04-22T16:23:44Z

Documentation preview: https://vllm--40633.org.readthedocs.build/en/40633/

JartX · 2026-04-22T16:25:34Z

@bringlein @tdoublep @kylesayrs @mgoin

gemini-code-assist

Code Review

This pull request introduces INT4 and INT2 per-token-head KV cache quantization for the V1 attention backend, employing Hadamard pre-rotation and sub-byte packing techniques. The implementation refactors quantization logic into a modular factory system and extracts shared Triton JIT helpers for attention operations. Feedback from the review emphasizes the necessity of validating that head sizes are powers of two for Hadamard-based modes to avoid assertion failures and suggests ensuring numerical consistency in the INT4 reshape kernel by using truncated scales during quantization to match the inference path.

Introduce sub-byte packed KV cache quantization for per-token-head dynamic scales, with two new modes: * ``INT4_PER_TOKEN_HEAD``: 2 × int4 per byte, random Hadamard pre-rotation, asymmetric scale + zero-point (zp steganographed in the low 4 bits of the float32 scale mantissa). * ``INT2_PER_TOKEN_HEAD``: 4 × int2 per byte, full Hadamard pre-rotation, Lloyd-Max centroid lookup (4 centroids over N(0, 1)). The attention read path for these modes is bespoke (split-Q dot across PACKING_FACTOR interleaved streams, inline centroid / nibble unpack, per-INT4 asymmetric zp correction, multi-stream accumulators) and lives in a new module ``vllm.v1.attention.ops.triton_quant_kv`` with a small factory registry: * ``QuantKVFactory`` is the protocol owning the write path (``reshape_and_cache``) and optionally the read path (``unified_attention``) for one ``KVQuantMode``. * Sub-byte packed modes override both (bespoke read kernel); INT8 / FP8 per-token-head modes override only the write path and use the common kernel via a constexpr branch. * Unused modes pay zero import / compile cost — factories are lazy-imported from ``_MODULES`` on first use. Other changes: * Move the INT8 / FP8 per-token-head cache write out of ``triton_reshape_and_cache_flash`` into the new factory module. * Dispatch in ``triton_unified_attention`` for the packed modes; the common kernel handles every other mode as before. * Promote ``cast_kv_tile`` to ``triton_attention_helpers`` so the common kernel and the packed kernel share a single implementation. * Extend ``triton_attn.py`` and ``cache.py`` to recognise the new ``int4_per_token_head`` / ``int2_per_token_head`` cache dtypes. Stacked on the 2D/3D unification refactor — the common-path helpers it extracts are reused by the packed kernel so the split-Q version shares prologue, softmax bookkeeping, segment reduction, ALiBi, softcap, qq_bias and sliding-window handling with the common path. Tests live in ``tests/quantization/test_per_token_kv_cache.py`` and ``tests/models/quantization/test_per_token_kv_cache.py``; a batch- invariance regression covering the new cache dtypes is added to ``tests/v1/determinism/test_matmul_batch_invariant.py``. Signed-off-by: JartX <sagformas@epdcenter.es>

Split the 1330-line packed_per_token_head.py into three files so that no module exceeds ~680 lines, making the diff easier to review: * _packed_reshape.py — INT4 / INT2 reshape (write) kernels, Lloyd-Max quantize/dequant helpers, shared Python launcher _run_reshape_kernel. * _packed_attention.py — _attn_packed read kernel (shared by INT4 and INT2 via PACKING_FACTOR constexpr) and its _launch_packed_attn wrapper. * packed_per_token_head.py — only the _PackedFactory base class, Int4PerTokenHeadFactory / Int2PerTokenHeadFactory, and the register() calls that wire them into the lazy registry. No behavior changes; this is purely a file split. Signed-off-by: JartX <sagformas@epdcenter.es>

JartX · 2026-04-24T17:19:00Z

UNQUANT KVCACHE

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.928	±	0.0164
		strict-match	5	exact_match	↑	0.912	±	0.0180

INT8 PER TOKEN HEAD

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.916	±	0.0176
		strict-match	5	exact_match	↑	0.912	±	0.0180

FP8 PER TOKEN HEAD

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.924	±	0.0168
		strict-match	5	exact_match	↑	0.912	±	0.0180

FP8

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.932	±	0.0160
		strict-match	5	exact_match	↑	0.908	±	0.0183

INT4 PER TOKEN HEAD

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.932	±	0.0160
		strict-match	5	exact_match	↑	0.916	±	0.0176

INT2_PER_TOKEN_HEAD

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.876	±	0.0209
		strict-match	5	exact_match	↑	0.844	±	0.0230

JartX · 2026-04-24T19:39:05Z

KV Cache Quantization Benchmark — Serving Results

Configuration: 100 concurrent requests · 512 input tokens / 128 output tokens per request · 51,200 total input tokens · 12,800 total generated tokens.

Summary Table

Metric	FP16 (baseline)	FP8	FP8 per-token-head	INT8 per-token-head	INT4 per-token-head	INT2 per-token-head
Duration (s)	17.08	17.53	18.14	18.30	19.64	20.26
Request throughput (req/s)	5.85	5.70	5.51	5.47	5.09	4.94
Output throughput (tok/s)	749.36	729.99	705.67	699.62	651.83	631.90
Peak output throughput (tok/s)	1700.00	1600.00	1600.00	1600.00	1400.00	1400.00
Total throughput (tok/s)	3746.78	3649.93	3528.37	3498.08	3259.17	3159.49
Mean TTFT (ms)	5445.90	5537.39	5771.07	5842.62	6105.02	6449.02
Median TTFT (ms)	6016.20	6117.63	6376.84	6457.99	6745.12	7135.07
P99 TTFT (ms)	9419.52	9579.00	9991.29	10107.75	10555.47	11123.32
Mean TPOT (ms)	90.67	93.50	96.38	97.05	105.41	107.57
Median TPOT (ms)	86.38	89.13	91.82	92.41	100.60	102.41
P99 TPOT (ms)	120.02	123.38	127.50	128.57	138.26	142.35
Median ITL (ms)	61.25	63.80	65.30	65.63	72.84	73.29
P99 ITL (ms)	1497.38	1526.65	1588.63	1601.77	1673.71	1759.27

Observations

All variants completed 100/100 requests with zero failures. Performance ranking is consistent across throughput and latency dimensions: FP16 > FP8 > FP8/INT8 per-token-head > INT4 per-token-head > INT2 per-token-head. The gap between FP8 and INT8 per-token-head is essentially noise (~0.9% on total throughput, ~0.7% on Mean TPOT), which is the expected behavior on this hardware given equivalent compute paths at that bit width. INT4 and INT2 pay a visible throughput cost (~13% and ~16% respectively) — the tradeoff only makes sense when the memory budget justifies it.

Raw per-run outputs preserved below for reference.

Raw Results

FP16 (baseline)

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  17.08
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              5.85
Output token throughput (tok/s):         749.36
Peak output token throughput (tok/s):    1700.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3746.78
Mean TTFT (ms):                          5445.90
Median TTFT (ms):                        6016.20
P99 TTFT (ms):                           9419.52
Mean TPOT (ms):                          90.67
Median TPOT (ms):                        86.38
P99 TPOT (ms):                           120.02
Mean ITL (ms):                           90.67
Median ITL (ms):                         61.25
P99 ITL (ms):                            1497.38

FP8

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  17.53
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              5.70
Output token throughput (tok/s):         729.99
Peak output token throughput (tok/s):    1600.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3649.93
Mean TTFT (ms):                          5537.39
Median TTFT (ms):                        6117.63
P99 TTFT (ms):                           9579.00
Mean TPOT (ms):                          93.50
Median TPOT (ms):                        89.13
P99 TPOT (ms):                           123.38
Mean ITL (ms):                           93.50
Median ITL (ms):                         63.80
P99 ITL (ms):                            1526.65

FP8 per-token-head

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  18.14
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              5.51
Output token throughput (tok/s):         705.67
Peak output token throughput (tok/s):    1600.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3528.37
Mean TTFT (ms):                          5771.07
Median TTFT (ms):                        6376.84
P99 TTFT (ms):                           9991.29
Mean TPOT (ms):                          96.38
Median TPOT (ms):                        91.82
P99 TPOT (ms):                           127.50
Mean ITL (ms):                           96.38
Median ITL (ms):                         65.30
P99 ITL (ms):                            1588.63

INT8 per-token-head

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  18.30
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              5.47
Output token throughput (tok/s):         699.62
Peak output token throughput (tok/s):    1600.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3498.08
Mean TTFT (ms):                          5842.62
Median TTFT (ms):                        6457.99
P99 TTFT (ms):                           10107.75
Mean TPOT (ms):                          97.05
Median TPOT (ms):                        92.41
P99 TPOT (ms):                           128.57
Mean ITL (ms):                           97.05
Median ITL (ms):                         65.63
P99 ITL (ms):                            1601.77

INT4 per-token-head

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  19.64
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              5.09
Output token throughput (tok/s):         651.83
Peak output token throughput (tok/s):    1400.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3259.17
Mean TTFT (ms):                          6105.02
Median TTFT (ms):                        6745.12
P99 TTFT (ms):                           10555.47
Mean TPOT (ms):                          105.41
Median TPOT (ms):                        100.60
P99 TPOT (ms):                           138.26
Mean ITL (ms):                           105.41
Median ITL (ms):                         72.84
P99 ITL (ms):                            1673.71

INT2 per-token-head

Successful requests:                     100
Failed requests:                         0
Benchmark duration (s):                  20.26
Total input tokens:                      51200
Total generated tokens:                  12800
Request throughput (req/s):              4.94
Output token throughput (tok/s):         631.90
Peak output token throughput (tok/s):    1400.00
Peak concurrent requests:                100.00
Total token throughput (tok/s):          3159.49
Mean TTFT (ms):                          6449.02
Median TTFT (ms):                        7135.07
P99 TTFT (ms):                           11123.32
Mean TPOT (ms):                          107.57
Median TPOT (ms):                        102.41
P99 TPOT (ms):                           142.35
Mean ITL (ms):                           107.57
Median ITL (ms):                         73.29
P99 ITL (ms):                            1759.27

@tdoublep @bringlein

JartX · 2026-04-24T21:36:50Z

Before starting the review @tdoublep @bringlein @mgoin @kylesayrs , please note the following pull request (PR), which includes the decode and prefill kernels for per token head. This is based on the previous PR, but I've been working on them for a few weeks.

If it's alright with you, since only the per token head branches are being addressed, I'd like to continue in the next PR, even though it has one more file and the dispatcher flow to the kernels.

#40835

I know these are additional flows to consider compared to int2/int4, but I'd say the gains are significant enough to justify the effort. Thank you so much for your time.

If you agree, I'll continue with that PR; otherwise, we'll continue with the stable and tested changes from the current one and add the kernels in another PR.

mergify · 2026-05-04T09:36:48Z

Hi @JartX, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

The INT8 reshape kernel computed q = clamp(x / scale, -128, 127) and relied on the implicit fp32 → int8 cast inside tl.store. That cast truncates toward zero, biasing every K and V slot by up to 0.5 LSB in the direction of zero. The bias is signed and accumulates over context length, contaminating the KV cache progressively and showing up as repetition / N-gram loops in long generations (e.g. think mode on Qwen3.5 with kv_cache_dtype=int8_per_token_head). The kernel docstring already describes the intended behaviour (q = clamp(round(x / scale), ...)) — only the round call was missing. Add an IS_INT_QUANT constexpr that selects explicit round-half-away- from-zero (+0.5 / -0.5 then to int32, same pattern the INT4 packed quantizer uses) when the destination is int8. FP8 leaves the implicit cast intact: fp8's native RNE on the cvt instruction is the correct rounding for sub-integer values, and forcing integer rounding there would destroy the fp8 precision below 1.0. INT4 path (separate file) already rounds correctly; INT2 uses Lloyd-Max centroids and does not round. The read-side HIP prefill kernel (paged_prefill_attn_rdna3_v2_int8.cu) and the Triton decode kernel were both audited as part of this investigation — softmax guards, block-boundary lookup and mask handling are all correct. The bug was strictly on the write side. Signed-off-by: JartX <sagformas@epdcenter.es>

mergify · 2026-05-12T22:43:57Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @JartX.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Cache RHT D₁ signs per (d, round_idx, device, dtype) instead of just (d, round_idx, device). This prevents implicit fp32 promotion when multiplying fp16 tensors by fp32 signs (53× faster for fp16). Signs are ±1.0, exactly representable in any IEEE float format. Signed-off-by: JartX <sagformas@epdcenter.es>

JartX · 2026-05-25T10:35:54Z

@tdoublep closed in favor of: #40835

JartX requested review from DarkLight1337, LucasWilkinson, MatthewBonanni, heheda12345, mgoin, pavanimajety, robertgshaw2-redhat, tdoublep, yewentao256 and ywang96 as code owners April 22, 2026 16:23

claude Bot reviewed Apr 22, 2026

View reviewed changes

mergify Bot added documentation Improvements or additions to documentation v1 labels Apr 22, 2026

JartX force-pushed the feature/int4-int2-per-token-head-v2 branch from ab04bea to ba7a480 Compare April 22, 2026 16:24

gemini-code-assist Bot reviewed Apr 22, 2026

View reviewed changes

Comment thread vllm/v1/attention/ops/triton_quant_kv/packed_per_token_head.py Resolved

Comment thread vllm/v1/attention/ops/triton_quant_kv/packed_per_token_head.py Resolved

JartX force-pushed the feature/int4-int2-per-token-head-v2 branch from ba7a480 to 0757de1 Compare April 23, 2026 15:30

JartX force-pushed the feature/int4-int2-per-token-head-v2 branch from 0757de1 to d485ea1 Compare April 24, 2026 15:22

yewentao256 added the ready ONLY add when PR is ready to merge/full CI is needed label Apr 24, 2026

JartX mentioned this pull request Apr 24, 2026

[Feature] Triton INT4 per-token-head KV cache quantization #40835

Merged

JartX added 3 commits April 25, 2026 14:19

Merge branch 'main' into feature/int4-int2-per-token-head-v2

d67bcf4

Merge branch 'main' into feature/int4-int2-per-token-head-v2

8e129ec

Merge branch 'main' into feature/int4-int2-per-token-head-v2

48708c8

JartX requested a review from zyongye as a code owner May 12, 2026 22:42

mergify Bot added the needs-rebase label May 12, 2026

JartX closed this May 25, 2026

Navigation Menu

Repository navigation

[Feature] Triton INT4 / INT2 per-token-head KV cache quantization - #40633

[Feature] Triton INT4 / INT2 per-token-head KV cache quantization#40633
JartX wants to merge 7 commits into
vllm-project:mainfrom
JartX:feature/int4-int2-per-token-head-v2

JartX commented Apr 22, 2026

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Apr 22, 2026

Uh oh!

JartX commented Apr 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

JartX commented Apr 24, 2026

Uh oh!

JartX commented Apr 24, 2026 •

edited

Loading

Uh oh!

JartX commented Apr 24, 2026 •

edited

Loading

Uh oh!

mergify Bot commented May 4, 2026

Uh oh!

mergify Bot commented May 12, 2026

Uh oh!

JartX commented May 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

JartX commented Apr 22, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Apr 22, 2026

Uh oh!

JartX commented Apr 22, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

JartX commented Apr 24, 2026

Uh oh!

JartX commented Apr 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

KV Cache Quantization Benchmark — Serving Results

Summary Table

Observations

Raw Results

FP16 (baseline)

FP8

FP8 per-token-head

INT8 per-token-head

INT4 per-token-head

INT2 per-token-head

Uh oh!

JartX commented Apr 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

mergify Bot commented May 4, 2026

Uh oh!

mergify Bot commented May 12, 2026

Uh oh!

JartX commented May 25, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

JartX commented Apr 24, 2026 •

edited

Loading

JartX commented Apr 24, 2026 •

edited

Loading