Skip to content

Add KVzapPress and ThresholdPress #171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
SimJeg merged 70 commits into from
Jan 14, 2026
Merged

Add KVzapPress and ThresholdPress #171

SimJeg merged 70 commits into from
Jan 14, 2026

Conversation

@SimJeg
Copy link
Collaborator

@SimJeg SimJeg commented Jan 12, 2026

  • Two new presses, KVzapPress and ThresholdPress
  • New root directory kvzap to reproduce KVzap training and evaluation
  • Script to re-run the kvpress leaderboard
  • Small fixes to LongBench and AIME25 evaluation

SimJeg added 30 commits December 2, 2025 09:27
@SimJeg
Add KVzapPress and ThresholdPress
72326e0 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Update pipeline
3b28410 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add KVzap to evaluation
d26db4b 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Relax constraint
b649181 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix tokenization for llama3
66ffc84 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix tokenization again
fd156f4 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
e541e8d 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Merge remote-tracking branch 'origin/main' into simon/kvzap
be20457
@SimJeg
Use post_init_from_model
3a056f7 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
d79bb64 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix ThresholdPress
2129982 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Push models to hub
b38e025 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Dont reload KVzap model each call
7d55330 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Update names
9384698 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add non greedy evaluation
60f2253 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add dir_id
d62bc9e 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add no_press
0edcef9 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix triviaqa
7865b1d 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix metric
6cf7cf4 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Update eval
4ef0280 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Refactor
d7f74ed 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix init
2146fa4 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix kvzip
aca5a5e 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add LB script
c471d5c 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
7361174 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add SDPX
f216e1c 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add tests
3e26ffb 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add threshold to evaluate
b21009f 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
c7cc80b 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Merge branch 'main' into simon/kvzap
07c5aca
SimJeg added 12 commits December 16, 2025 16:58
@SimJeg
Fix again?
1660f04 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Yet again
2535c25 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Merge branch 'main' into simon/kvzap
2bf64b2
@SimJeg
Support multi-GPU
c0c2fb8 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Revert kvzip_press.py to main (moved changes to simon/kvzip branch)
e73f4e5 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Handle tokenization edge case
197d89a 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add FP8 support
2769225 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Remove kvzip_plus normalization
d967ef0 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
b8fb408 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix style
7139da2 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Fix FP8
d6beeca 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Reorder parameters
1801694 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@copy-pr-bot
Copy link

copy-pr-bot bot commented Jan 12, 2026

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@SimJeg SimJeg requested a review from maxjeblick January 12, 2026 09:41
@SimJeg
Copy link
Collaborator Author

SimJeg commented Jan 12, 2026

/ok to test 1801694

maxjeblick
maxjeblick reviewed Jan 12, 2026
View reviewed changes
Copy link
Collaborator

@maxjeblick maxjeblick left a comment

Choose a reason for hiding this comment

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

Thanks a lot for the PR!
I've added some smaller comments.

SimJeg added 4 commits January 12, 2026 11:10
@SimJeg
Address PR feedback
8cd916d 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Remove LB plot
b4e780f 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Remove LB plot
fad04c0 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Add Ridge
dc8f10a 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Copy link
Collaborator Author

SimJeg commented Jan 12, 2026

/ok to test dc8f10a

@SimJeg
Merge branch 'main' into simon/kvzap
c0ac69f 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Copy link
Collaborator Author

SimJeg commented Jan 12, 2026

/ok to test c0ac69f

maxjeblick
maxjeblick approved these changes Jan 12, 2026
View reviewed changes
Copy link
Collaborator

@maxjeblick maxjeblick left a comment

Choose a reason for hiding this comment

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

LGTM!

@SimJeg
Add arxiv url
e863021 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg
Update pipeline and README
d10e102 
Signed-off-by: SimJeg <sjegou@nvidia.com>
@SimJeg SimJeg merged commit 50c2ae5 into main Jan 14, 2026
2 checks passed
@SimJeg SimJeg deleted the simon/kvzap branch January 14, 2026 09:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@maxjeblick maxjeblick maxjeblick approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@SimJeg @maxjeblick