Skip to content

[Feature] KV cache per-token-head INT8/FP8 quantization#38378

Merged
mgoin merged 119 commits into
vllm-project:mainfrom
JartX:feature/kvcache_per_token
Apr 2, 2026
Merged

[Feature] KV cache per-token-head INT8/FP8 quantization#38378
mgoin merged 119 commits into
vllm-project:mainfrom
JartX:feature/kvcache_per_token

Conversation

@JartX

@JartX JartX commented Mar 27, 2026 •
edited by mgoin
Loading

Copy link
Copy Markdown
Contributor

Continue of PR: #36893 with changes requested by @mgoin
At comment:
#36893 (comment)

This PR adds per-token-head kv cache quantization to the Triton attention backend

INT8_PER_TOKEN, FP8_PER_TOKEN

| Metric | FP16 | FP8 | INT8_PER_TOKEN | FP8_PER_TOKEN |
|---|---|---|---|---|
| Duration (s) | 89.17 | 91.21 | **80.95 ✅** | 104.45 |
| Request throughput (req/s) | 2.24 | 2.19 | **2.47 ✅** | 1.91 |
| Output tok/s | 1148.30 | 1122.64 | **1264.91 ✅** | 980.33 |
| Total tok/s | 2296.61 | 2245.28 | **2529.83 ✅** | 1960.66 |
| Mean TTFT (ms) | 10130 | 11453 | **10072 ✅** | 11572 |
| Median TTFT (ms) | 9689 | 11017 | **9628 ✅** | 11078 |
| P99 TTFT (ms) | 21286 | 22612 | **21119 ✅** | 24015 |
| Mean TPOT (ms) | 150.93 | 152.30 | **135.66 ✅** | 178.22 |
| Median TPOT (ms) | 152.19 | 153.55 | **136.88 ✅** | 179.77 |
| P99 TPOT (ms) | 164.52 | 165.89 | **150.01 ✅** | 193.52 |
| Mean ITL (ms) | 150.93 | 152.30 | **135.66 ✅** | 178.22 |
| Median ITL (ms) | 137.34 | 139.27 | **121.57 ✅** | 164.21 |
| P99 ITL (ms) | 468.12 | 467.41 | **461.64 ✅** | 518.82 |

JartX and others added 30 commits March 12, 2026 15:32
@JartX @EricccYang
int8 kv cache per-token per-head scale
5da77fe 
Signed-off-by: JartX <sagformas@epdcenter.es>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
@JartX
pre-commit
6f830fd 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
pre-commit
230d37e 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
isolate flows
325f7b8 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
isolate flows
7b15ebc 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Merge branch 'main' into feature/kvcache_int8_scale_per_token_continu…
fab75a6 
…e_pr_34327_thanks_ericccyang

Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
changed warning message
c5a61f2 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Fix assertion for kv_cache_torch_dtype
503f7b8 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
working on kvcache int8
eb96acf 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
pre-commit
42e914e 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
working in the test
6411697 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
pre-commit
6dcbc44 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Merge branch 'main' into feature/kvcache_int8_scale_per_token_continu…
3d293cc 
…e_pr_34327_thanks_ericccyang
@JartX
Merge branch 'main' into feature/kvcache_int8_scale_per_token_continu…
1a2bd45 
…e_pr_34327_thanks_ericccyang
@JartX
NotImplementedError at detect has_kv_transfer_group
0e30d8a 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
pre-commit
d351be8 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Merge branch 'main' into feature/kvcache_int8_scale_per_token_continu…
70ff12c 
…e_pr_34327_thanks_ericccyang
@yewentao256 @Isotr0py
Apply suggestions from code review
9f16e7a 
Co-authored-by: Isotr0py <2037008807@qq.com>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
@JartX
fix key_cache_shape with [:3]
573b6cd 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
post-dot dequantization
82f96cf 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
removed notimplementederror requested by @NickLucche
1e91fdc 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
test changes
af878f8 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
test changes
d935061 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Reserve GPU memory for INT8 per-token scale caches
781306e 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
changed doc string
446fae6 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
refactor flex
ca3f02f 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Generalize per-token KV cache scale infrastructure for extensibility
4ee2689 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
refactor to use kv_cache_uses_per_token_scales
9794165 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
torch.empty in ensure_per_token_scale_caches
a96abe5 
Signed-off-by: JartX <sagformas@epdcenter.es>
@JartX
Merge branch 'main' into feature/kvcache_int8_scale_per_token_continu…
c6c8b8e 
…e_pr_34327_thanks_ericccyang
@tjtanaa

tjtanaa commented Apr 2, 2026

Copy link
Copy Markdown
Member

@tjtanaa could you help benchmark the triton backend to make sure we aren't introducing any regressions to the non-kv quant case?

sure, let me get some figures.

tjtanaa
tjtanaa approved these changes Apr 2, 2026
View reviewed changes

@tjtanaa tjtanaa left a comment •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, amazing work.

There is no noticeable regression before and after this PR changes.

Configuration Output Throughput Diff % Mean TTFT Diff % Mean TPOT Diff % Duration Diff %
Ori noquant no prefix 7434.96 tok/s 0% 3366.76 ms 0% 30.90 ms 0% 35.26 s 0%
Orifp8quant No cache 7309.90 tok/s -1.68% 3411.33 ms +1.32% 31.44 ms +1.75% 35.86 s +1.70%
Afternoquant no cache 7416.04 tok/s -0.25% 3504.87 ms +4.10% 30.85 ms -0.16% 35.35 s +0.26%
Afterfp8quant no cache 7292.52 tok/s -1.92% 3428.95 ms +1.85% 31.51 ms +1.97% 35.95 s +1.96%

@mgoin mgoin merged commit 2ce3d0c into vllm-project:main Apr 2, 2026
79 of 80 checks passed
@github-project-automation github-project-automation Bot moved this from Todo to Done in AMD Apr 2, 2026
@JartX JartX mentioned this pull request Apr 2, 2026
Closed
orozery added a commit that referenced this pull request Apr 2, 2026
@orozery
Revert "[Feature] KV cache per-token-head INT8/FP8 quantization (#38378…
f36b324 
…)"

This reverts commit 2ce3d0c.
@ZhanqiuHu ZhanqiuHu mentioned this pull request Apr 2, 2026
Merged
1 task
@MatthewBonanni MatthewBonanni mentioned this pull request Apr 2, 2026
Merged
5 tasks
@wizzense wizzense mentioned this pull request Apr 5, 2026
Closed
7 tasks
HenryTangDev pushed a commit to HenryTangMain/vllm that referenced this pull request Apr 6, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py @HenryTangDev
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
63796dc 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
@EricccYang EricccYang mentioned this pull request Apr 7, 2026
Open
1 task
@ashok-arora

ashok-arora commented Apr 7, 2026

Copy link
Copy Markdown

@JartX INT8 seems to be lagging behind FP16, is that expected even for longer context length and higher batch sizes?

@JartX

JartX commented Apr 7, 2026

Copy link
Copy Markdown
Contributor Author

@ashok-arora
Hi! Could you give me more details, please? Tell me which graphics cards you're using and what you mean by "behind." Performance? Accuracy?

Thanks a lot!

puririshi98 pushed a commit to puririshi98/vllm that referenced this pull request Apr 7, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py @puririshi98
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
61f18b8 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
Signed-off-by: Rishi Puri <riship@nvidia.com>
mtparet pushed a commit to blackfuel-ai/vllm that referenced this pull request Apr 9, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py @mtparet
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
1d3ec42 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
This was referenced Apr 11, 2026
Closed
Closed
Closed
Draft
mystous pushed a commit to mystous/vllm_hybrid that referenced this pull request May 10, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
764bafe 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
ec5953c 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
my-other-github-account pushed a commit to my-other-github-account/vllm that referenced this pull request May 15, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
132542c 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
jhu960213 pushed a commit to jhu960213/vllm that referenced this pull request May 20, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
d740bf0 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
mvanhorn pushed a commit to mvanhorn/vllm that referenced this pull request Jun 4, 2026
@JartX @EricccYang
@yewentao256 @Isotr0py @mvanhorn
[Feature] KV cache per-token-head INT8/FP8 quantization (vllm-project…
b394a88 
…#38378)

Signed-off-by: JartX <sagformas@epdcenter.es>
Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: yangyang4991 <yangyang4991@gmail.com>
Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Co-authored-by: Isotr0py <2037008807@qq.com>
Signed-off-by: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@MatthewBonanni MatthewBonanni MatthewBonanni left review comments

@claude claude[bot] claude[bot] left review comments

@mgoin mgoin mgoin approved these changes

@NickLucche NickLucche NickLucche approved these changes

@tjtanaa tjtanaa tjtanaa approved these changes

@AndreasKaratzas AndreasKaratzas AndreasKaratzas approved these changes

@tdoublep tdoublep Awaiting requested review from tdoublep tdoublep is a code owner

@robertgshaw2-redhat robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat is a code owner

@yewentao256 yewentao256 Awaiting requested review from yewentao256 yewentao256 is a code owner

@pavanimajety pavanimajety Awaiting requested review from pavanimajety pavanimajety is a code owner

@DarkLight1337 DarkLight1337 Awaiting requested review from DarkLight1337 DarkLight1337 is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 is a code owner

@njhill njhill Awaiting requested review from njhill njhill is a code owner

@heheda12345 heheda12345 Awaiting requested review from heheda12345 heheda12345 is a code owner

@WoosukKwon WoosukKwon Awaiting requested review from WoosukKwon WoosukKwon is a code owner

@alexm-redhat alexm-redhat Awaiting requested review from alexm-redhat alexm-redhat is a code owner

@ApostaC ApostaC Awaiting requested review from ApostaC

@orozery orozery Awaiting requested review from orozery

@zhuohan123 zhuohan123 Awaiting requested review from zhuohan123 zhuohan123 is a code owner

@youkaichao youkaichao Awaiting requested review from youkaichao youkaichao is a code owner

@LucasWilkinson LucasWilkinson Awaiting requested review from LucasWilkinson LucasWilkinson is a code owner

@tlrmchlsmth tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth is a code owner

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

documentation Improvements or additions to documentation quantization ready ONLY add when PR is ready to merge/full CI is needed rocm Related to AMD ROCm v1

Projects

AMD
Status: Done

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

9 participants

@JartX @tjtanaa @ashok-arora @mgoin @LucasWilkinson @NickLucche @MatthewBonanni @AndreasKaratzas @yewentao256