[RFC] implementing Expected Attention in llama.cpp?

[ddh0]

Recently, NVIDIA put out a paper describing a method to reduce the size of the KV cache. They call their preferred method Expected Attention.

Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution

Memory consumption of the Key-Value (KV) cache represents a major bottleneck for efficient large language model (LLM) inference. While attention-score-based KV cache pruning shows promise, it faces critical practical limitations: attention scores from future tokens are unavailable during compression, and modern implementations like Flash Attention do not materialize the full attention matrix, making past scores inaccessible. To overcome these challenges, we introduce Expected Attention, a training-free compression method that estimates KV pairs importance by predicting how future queries will attend to them. Our approach leverages the distributional properties of LLM activations to compute expected attention scores in closed form for each KV pair. These scores enable principled ranking and pruning of KV pairs with minimal impact on the residual stream, achieving effective compression without performance degradation. Importantly, our method operates seamlessly across both prefilling and decoding phases, consistently outperforming state-of-the-art baselines in both scenarios. Finally, we release KVPress, a comprehensive library to enable researchers to implement and benchmark KV cache compression methods, already including more than 20 techniques.

From my reading of the paper, the thought process is like this:

• Previous works have shown that the hidden states (cur) at each layer follow a gaussian distribution (and they re-confirmed this here too)
• The query projection tensor in each layer (wq) can be thought of as a linear transformation on cur, which maintains this same distribution for the resulting Qcur projection
• Using the callback cb(), we can observe the query projections Qcur in all layers over an interval, in order to model a distribution of the expected future queries (the paper also takes into account RoPE / positional encoding for future positions, but I could not figure out how to do that, so I skipped it for now)
• Using the expected future queries (query?) we can compute the expected importance of each entry currently in the KV cache. They assign an importance score to all KV pairs based on the magnitude of the expected contribution of that pair (assuming the expected future query).
• They sort the KV entries based on the importance scores, and discard the bottom X% of entries.

So for example if the user specified a "compression ratio" of 20%, we would discard the 20% of KV entries that are expected to have the least impact on the hidden state in the future.

---

I had hoped to be able to get this working and submit a PR for it on my own, but there are some problems that I don't know how to solve:

• How to get the actual KV entries from the llama_context? The llama_memory structure abstracts away the implementation of the different cache types, so this is not straightforward.
• How to efficiently compute the expected attention scores for all KV entries? (It would probably be best to score per-layer and combine all those scores into one unified importance ranking somehow.)
• The paper re-computes the expected attention 512 tokens of inference. It might be better/easier to only re-compute when we are done with prompt processing, and before we start text generation? I don't know.
• What do we do once we have the positions we want to discard? Would it be possible to use the existing defrag logic for this?

I am posting this thread in hopes of starting a discussion to see if this method will even be feasible to implement in llama.cpp, or if it would be too complicated. The code is certainly NOT in a state that is ready for use, or even guaranteed to work at all. I am just including my demo code as a starting point. I would be pleasantly surprised if this PR ever leaves the draft stage. :)

Make sure to read the contributing guidelines before submitting a PR

[ddh0]

additional references:

• https://github.com/NVIDIA/kvpress
• https://github.com/NVIDIA/kvpress/blob/main/kvpress/presses/expected_attention_press.py

[ngxson]

If I understand correctly, the idea can be explained in more simple terms:

• Whenever KV cache is used, we track the attention score of the KV entries
• If a given KV entry is used less than < X% of the time, then we simple remove/prune it (we define "being used" by compare attention score against a threshold)

If we view the model as a pure function f(x, kv) = (y, kv_new) where x the input token, kv the current cache state and y the next token, then what I understand is that we now need f(x, kv) = (y, kv_new, attn_score) for this to work, so the simple way is to add attn_score as an output node, much like how some models can output token and/or embeddings.

However, it might sounds more complicated than that. First off, flash attention won't expose the attention score to the outside world. And secondly, this may make it more tricky to reuse cache cross multiple sequences, same problem with KV shifting/SWA.

Also, just to mention that there were some similar works in #16817 , but might not very useful as the other PR is AI-coded and the author doesn't seem to fully understand it.