Skip to content

[Improvement] Formalize KV sharing#47290

Open
remi-or wants to merge 13 commits into
mainfrom
formalize-kv-sharing
Open

[Improvement] Formalize KV sharing#47290
remi-or wants to merge 13 commits into
mainfrom
formalize-kv-sharing

Conversation

@remi-or

@remi-or remi-or commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

CI

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!

@remi-or remi-or requested a review from Cyrilvallez July 13, 2026 04:24
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

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.

@github-actions

Copy link
Copy Markdown
Contributor

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

run-slow: diffusion_gemma, gemma3n, gemma4, gemma4_unified

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 29232223385:1
Result: failure | Jobs: 15 | Tests: 140,428 | Failures: 1 | Duration: 13h 23m

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants