UPSTREAM PR #19183: [RFC] implementing Expected Attention in llama.cpp?

[loci-dev]

Note

Source pull request: ggml-org/llama.cpp#19183

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

[loci-review]

No meaningful performance changes were detected across 115123 analyzed functions in the following binaries: build.bin.libllama.so, build.bin.libmtmd.so, build.bin.llama-tts, build.bin.llama-cvector-generator, build.bin.libggml-base.so, build.bin.libggml-cpu.so, build.bin.libggml.so, build.bin.llama-tokenize, build.bin.llama-quantize, build.bin.llama-qwen2vl-cli, build.bin.llama-bench, build.bin.llama-gemma3-cli, build.bin.llama-gguf-split, build.bin.llama-llava-cli, build.bin.llama-minicpmv-cli.

🔎 Full breakdown: Loci Inspector.
💬 Questions? Tag @loci-dev.