QuantizedCache first token processing is counterintuitive / worse than in papers
Author: goodevening13
Pull request: no
Labels: name: Quantization

name: bug

name: Generation

### System Info

transformers==4.46.3
torch==2.5.1
(though, it does not depend on library versions)

### Who can help?

@sunmarc @zucchini-nlp



### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

**TL;DR** current quantized KV cache implementation works poorly **with the first token**, causes problems due to attention sinks. I describe this in detail in 'expected behavior'.

To reproduce the problem, simply run LLM inference (e.g. Llama 3.2 3B) with QuantoQuantizedCache (or other QuantizedCache descendants), e.g.

```
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf", torch_dtype=torch.float16, device_map="cuda:0")
inputs = tokenizer("I like rock music because", return_tensors="pt").to(model.device)

model.generate(**inputs, do_sample=False, max_new_tokens=20, cache_implementation="quantized", cache_config={"backend": "quanto", "nbits": 4})
```
Note that in this example, the quantized cache should only quantize tokens that are (buffer size) positions away. However, **it mistakenly quantized the first token right away**. This is important because the first token often acts like an [Attention Sink](https://arxiv.org/abs/2309.17453) and quantizing it substantially reduces model accuracy.

I describe the problem and demonstrate the quality drawdown in detail in the "Expected Behavior" section (below).



### Expected behavior

### What exactly is wrong

The problem occurs every time inference is run with QuantizedCache or its descendants.

The current implementation of QuantizedCache always quantizes the very first token, adding it to the q_cache instead of storing it in the full precision cache. However, according to the docstring, quantization should only occur "when the length goes beyond maximum capacity, and the original precision cache is discarded and moved into the quantized cache." This implies that the first token should only be quantized when the cache length reaches its maximum capacity.

This issue may affect model quality, as recent papers on KV quantization emphasize the importance of the first token, also known as the "attention sink" ( https://arxiv.org/abs/2309.17453 ). In fact, some papers suggest not quantizing this token at all (e.g. https://arxiv.org/abs/2410.05265 ). Failing to store the first token in full precision could therefore be problematic.

### Proposed fix

The fix involves modifying the "update" method so it keeps the first token in the FP16 buffers:

https://www.diffchecker.com/Rj5qaOCj/

Here's the full code of [fixed QuantizedCache](https://gist.github.com/goodevening13/36f1d92ccf87009e03e65f3e1bc54a32) that I used for benchmarks below.


### Effect on model quality:

To verify that this fix work, I compared the two implemantations in terms of WikiText-2 perplexity (lower is better). In the table below 'Before fix' evaluates QuantoQuantizedCache and with the current code (see versions above) and 'After fix' uses the [patch](https://gist.github.com/goodevening13/36f1d92ccf87009e03e65f3e1bc54a32) described in the previous section. These were measured with Llama-3.2-3B backbone using 2-bit and 4-bit with all default hyperparameters (see details below) .

| WikiText-2 PPL | FP16 | 4-bit | 2-bit |
|----------------|------|-------|-------|
| Before fix     |    6.9865   |  7.0052     |    12.2152   |
| After fix      |    6.9865   |   7.0005    |     **11.4623**  |


The full code for obtaining these results can be found in [gist link](https://gist.github.com/goodevening13/2aa807fb79fbe135e3162a3c54601a8e). If any additional evaluations or explanations are required, I am happy to provide those.


------

Thanks to [surkovvv](https://github.com/surkovvv) for his help with the evaluation codebase


Comments:
@goodevening13 hey!

Thanks for such a detailed analysis with evaluation, indeed keeping the first token as in Attention Sink paper can give quality boost. The current `QuantizedCache` doesn't follow 100% any existing paper and only borrowed the buffer size idea from KIVI, as mentioned in the documentation.

We are open to contributions and ideas to make it more efficient / decrease negative effects on perplexity. Please feel free to open a PR with the suggested changed, will be happy to have it in `transformers` :)

I'll give this a shot as well .

@sambhavnoobcoder great, feel free to submit a PR and tag me there 🤗 

I apologize for the time it has taken to address this issue. A pull request has been created and is available for review [here](https://github.com/huggingface/transformers/pull/35760)

Adding a data point that may help triage this family of issues (cross-ref #40099, PR #35760): `QuantizedCache` silently produces fully corrupted generations on Qwen2.5-7B-Instruct, while the identical call is lossless on Mistral-7B-Instruct-v0.3.

Setup: greedy decode, ~1.4k-token prompts, `max_new_tokens` 30-56, fp16 controls exact in the same harness, measured with a 120-case behavioral probe (exact-match action routing).

| model | fp16 | quanto int4 (defaults) |
|---|---|---|
| Qwen2.5-7B-Instruct | 98/120 | **0/120** (every generation corrupted, e.g. `' FINAL FINALциально ... 🤖'`) |
| Mistral-7B-Instruct-v0.3 | 84/120 | **84/120** (zero loss) |

What we ruled out by bisection before attributing:

- transformers version: reproduces identically on 4.49.0, 4.57.6 and 5.2.0
- torch: reproduces on 2.6.0+cu124 and 2.12.0+cu130
- backend: both quanto and hqq corrupt (different garbage)
- `residual_length`: identical corruption at 128 / 256 / 1024 / 4096. Note the prefill is quantized in one shot regardless of `residual_length`, so a residual larger than the prompt does not protect it.
- calling convention: dict `cache_config` and `QuantizedCacheConfig` both
- the PR #35760 approach (keep prefill in residual): applied locally, output is clean only until the first residual overflow re-quantizes everything, so it does not fix this case. A permanent fp16 attention-sink variant (first 4 tokens kept in original precision) does not fix it either.

So this looks like a model-family limitation of naive per-channel intN KV quantization against Qwen2.5's K/V outlier channels, rather than an implementation bug. For reference, distribution-aware 4-bit KV methods (e.g. Sinkhorn-normalized) hold 96-102/120 on the same model and probe, so the model is quantizable with normalization.

Given Qwen2.5's popularity, would you consider documenting validated model families for `QuantizedCache`, or emitting a warning? Happy to provide the minimal repro script.

(probe and full tables: https://github.com/sztlink/kv-score)