[generation] Support cache-cropping methods

[gante]

What does this PR do?

Adds the requirements to solve #35168

[EDITED after the discussion ending on this comment]

New dynamic cache compression methods may crop the cache differently at each layer. Doing so breaks our implicit assumption regarding the shape of the dynamic cache, where all layers have the same length. As we can see in the issue tagged above, we obtain an exception if we decide to do so.

We can enable it with small code changes in downstream libraries if we prepare the causal mask based on the maximum sequence length seen in all layers in the cache. A downstream library would then have to implement a custom attention forward pass, to left-crop the attention mask accordingly (e.g. using the length of the key)

---

The following tests were run, with no regressions compared to main:

• RUN_SLOW=1 py.test tests/utils/test_cache_utils.py -vv
• RUN_SLOW=1 py.test tests/models/llama/test_modeling_llama.py -vv

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zucchini-nlp]

As it solves the issue for #35168 with no major code changes from KVPress side, LGTM

My only concern is that it is slightly breaking, if users were somehow exploiting this right-side slicing feature before

[gante]

@ArthurZucker, this draft PR makes a small modeling change to enable an advanced cache feature, as used e.g. in Nvidia's KVPress library -- see the PR header for full details.

If you agree with these changes, I will propagate them to other models before requesting a final review :)

[ArthurZucker]

Hey! Sorry but I don't think it makes sense to do this!
It's a very niche case, and is not transparent for people who just want to see how llama works!

[gante]

@ArthurZucker makes sense. In fact, downstream libraries can implement their own custom attention implementation and crop the mask as they wish there.

We would be able to unblock this cool use case in other libs if we merge the non-modeling change, i.e. make the get_seq_length(layer=None) call return the longest sequence length for all layers, as opposed to the sequence length for layer 0 (current behavior). The two are equivalent except on this advanced use case. This change looks okay to you, correct?

EDIT: chatted on slack, moving forward with the non-model changes

[gante]

This comment applies to all models: the +1 not needed, we are discarding its corresponding data in the attn forward pass (e.g. here)

By removing it, we should expect a minimal improvement wrt memory and inference time. We also ensure that the causal mask has exactly the length of the (cached) input length, which is a requirement for the downstream use case this PR wants to enable.

[zucchini-nlp]

Btw, I am working with Mllama now and remembered it has weird cache where cross attention layers have a very large length. Can we make sure mllama generation tests dont fail after this change?

[gante]

@zucchini-nlp correct, it does change. I'm adding a test to catch more models like this. EDIT: mllama is the only model that fails this test :)

btw, RUN_SLOW=1 py.test tests/models/mllama/test_modeling_mllama.py results in 11 failures on both main and this branch 😬

[zucchini-nlp]

That is too much for the main branch, afaik mllama will not be able to run the slow tests because we are in EU and the repo is gated 😢

Currently I am removing skip from many generation tests that checked hasattr(config, "use_cache") otherwise all of them are skipped in multimodal models (PR will be submitted soon). In mlllama case it also skips all tests that need caching I guess. From code it seems like the get_seq_length() is used only when preparing cache_position if None and when preparing the attention mask. So I think the model might fail whenever there is no mask to infer length or when continuing generation from "initial prompt"

For other models, in VLM side that's the only one and maybe ImageGPT is a bit peculiar. From LLM side all I can think is the Hybrid Cache models where sliding window length is fixed and can be lower/higher than static layers