[remi-or]

This PR aims to formalize / standardize KV cache sharing in transformers in case more models add KV cache sharing patterns.
In recent years, there has been more efforts to reduce the amount of KV cache generated by a model, in the hope of reducing the memory footprint of long-context generation. One technique is to share KV cache between layers, like in YOCO (https://arxiv.org/abs/2405.05254) and CLA (https://arxiv.org/html/2405.12981v1).
Recent models, like those from the Gemma family, were implemented in transformers with KV cache sharing. Although the actual implementation is simple and correct, it would not allow for complex KV cache sharing patterns like CLA.
To make this possible, we implement the following:

• A dedicated KVCacheSharingStorage which function as an alternative to the UserDict that was previously used to share KV cache across layers. Like the UserDict, it is not an instance of a dict object, so it does not get picked up by torch's FSDP2 helpers, which would lead to a crash. But it also implements a to method, so accelerate can dispatch it across device using the same hooks as the rest of the inputs. Furthermore, the to method acts in place (documented) so once this container crosses a device boundary, it does not need to be moved again if next layer is on the same device (not the case atm). Finally, this means we can remove the manual .to calls from the modeling code, making it lighter.
• A new type of cache layer, DummyLayer, which is used for layers that use KV cache sharing and read their cache from another layer. Before, since only the last layers were shared (YOCO scheme), what was done was to simply truncate the number of cache layers, so that KV sharing layers would not have a corresponding cache layer that would store cache. This led to a discrepancy between the number of layers and the number of cache layers, and it was not comaptible with KV sharing schemes where some layers in the middle of the model share their KV cache (CLA). The DummyLayer solves both these problems.
• A new @property on the configs with KV cache sharing (and only those), named kv_sharing_roles, which the Cache object can read when constructing a new cache. There are 3 roles: consumer (the layer reads KV cache from another layer), producer, (the cache produced by this layer is going to be read by another layer), and independant (the layer does not participate in cache sharing, ie. a "normal" layer). When a Cache object is created, if the config has this property, then the Cache object create DummyLayer objects for consumer layers. The distinction between producer and independent is needed because producer layer store their KV cache in the KVCacheSharingStorage for other layers to reuse (important for sliding window layers). This property centralizes everything related to KV cache sharing: to know the exact behavior of all layers, it is the only thing needed.

Quite a long description, sorry! The idea here is to standardize KV sharing-related attributes and mechanism so that we have a simple yet adaptable interface so that when more KV-sharing models and schemes get added to trasnformers, we can support them out of the box without having to change this API. Hence feedback is very much appreciated!

[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.

[Cyrilvallez]

Some general thoughts on the design!

Agreed that we can have a native cache sharing primitive, but it's just a dict with a .to method so let's inherit from existing data structures rather than reinventing here.

For the DummyLayer, it's indeed something that I thought about for some time, to remove weird mappings from layer_types such as moe or mlp which do not have a cache. But let's make the class as minimal as possible, and not inherit from the abstract attention cache layer class!

[Cyrilvallez]

BTW, what I previously envisioned was that DummyLayers that are shared layers (so not the ones used if the layer_type is moe or mlp) could keep a ref to the index of the layer feeding them their states, so could retrieve the shared states natively without having to keep the additional dict argument in all the forward
But this can be done later as well, no need to do all at once - just what I had in mind to make it as seamless as possible long term

[Cyrilvallez]

run-slow: gemma3n, gemma4, nemotron_h

[Cyrilvallez]

run-slow: gemma3n, gemma4, nemotron_h

[github-actions]

[For maintainers] Suggested jobs to run (before merge)

run-slow: diffusion_gemma, gemma3n, gemma4, gemma4_assistant, gemma4_unified, gemma4_unified_assistant, nemotron_h

[Cyrilvallez]

run-slow: gemma3n, gemma4, nemotron_h

[github-actions]

CI (Nvidia + AMD)

Nvidia · ✅ No new failures 🎉👏 · ⚙️ Workflow · 📊 Dashboard

Failures (5 pre-existing)

    Pre-existing failures (still failing) — 5 total:

        Same error (5):

        - gemma3n:
            tests/models/gemma3n/test_modeling_gemma3n.py::Gemma3nIntegrationTest::test_generation_beyond_sliding_window
            tests/models/gemma3n/test_modeling_gemma3n.py::Gemma3nIntegrationTest::test_model_4b_bf16
            tests/models/gemma3n/test_modeling_gemma3n.py::Gemma3nIntegrationTest::test_model_4b_image
            tests/models/gemma3n/test_modeling_gemma3n.py::Gemma3nIntegrationTest::test_model_4b_multiimage

        - gemma4:
            tests/models/gemma4/test_modeling_gemma4.py::Gemma4IntegrationTest::test_model_with_image

AMD · ⚠️ CI failed to report results · ⚙️ Workflow · 📊 Dashboard

[Cyrilvallez]

Ok I pushed quite a few changes, all looks good to me now! Merging this! 🤗