AIME dataset support?

[Zhaojp-Frank]

Bug

Hi, there.
Thanks for the wonderful kvpress.
I'm reading the paper "Expected Attention", and it's claimed AIME2025 is tested with good accuracy.
However, there's no AIME dataset supported in evaluate_registry.py.

1. first just to confirm at first, does ExpectedAttn on AIME2025 supported e.g. using DeepSeek-R1-Qwen-7b ? is there a script to reproduce the Fig4.

• or simply run with ExpectedAttn pipeline and input AIME25 question?

2. since AIME25 is kind of short prompt + long decoding tokens, so KV eviction happens on-demand during decode stage?

To Reproduce

No AIME dataset listed in evaluation/evaluate_registry.py

Repository version

master

[maxjeblick]

Thanks a lot for the nice words!

We are currently working on adding compression during decoding functionality into kvpress:

• Compression during decoding functionality will probably be merged today Add decoding press functionality #139.
• AIME25 dataset evaluation is on this branch: https://github.com/NVIDIA/kvpress/tree/aledev/decoding_eval We plan to open a PR once Add decoding press functionality #139 has been merged. In particular, the dataset has been published here.

Moving the code into main took a bit longer than the paper, as we did some additional quality checks on #139.

[Zhaojp-Frank]

great! I'd like to try it~