Static Cache: no mandatory cache_positions input

[gante]

What does this PR do?

Removes the hard requirement of cache_position from the public model classes (e.g. classes that can be loaded with AutoModel). Its contents can be derived from the cache's get_seq_length().

---

Performance and Quality checks

Slow tests run, getting the same outcome as in main:

• RUN_SLOW=1 py.test tests/models/llama/test_modeling_llama.py -vv
• RUN_SLOW=1 py.test tests/models/gemma/test_modeling_gemma.py -vv
• RUN_SLOW=1 py.test tests/test_cache_utils.py -vv [Note: two new tests were added here]

👉 Note that these tests ensure that torch.compile works, with and without attention_mask being passed.

Local benchmark (RTX3090, tiny llama) -- no changes
(main)

(this pr)

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

[gante]

cc @fxmarty :)

[ArthurZucker]

I don't think that is the way we should go. For now have not seen issue with having a new argument, we added paddin_mask at some point.
IMO we should move to having cache_positions in generate as this is more explicit, easier to maintain and less error prone.
More than that, these operation are input dependent while being vectorized they add complexity when there should not be!

[fxmarty]

No specific opinon, just that using cache_position in the modeling code when use_cache=False is not very intuitive to me:

transformers/src/transformers/models/llama/modeling_llama.py

Lines 975 to 983 in 75ed76e

	if cache_position is None:
	if isinstance(past_key_values, StaticCache):
	raise ValueError("cache_position is a required argument when using StaticCache.")
	cache_position = torch.arange(
	past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
	)
	if position_ids is None:
	position_ids = cache_position.unsqueeze(0)

Only thing I think: cache_position (if it remains) should be a strictly optional argument & the model should work without for BC.

[gante]

EDIT -- before you read why I believe removing cache_positions from the signature is an important change for our users, have a look at the current diff :)

✅ much simpler computations (@ArthurZucker addresses your concerns, although the previous version also did not introduce slowdowns)
✅ working StaticCache.get_seq_length(), which removes a few exceptions
✅ no exception is raised when an input tensor is missing, with the StaticCache (with cache_positions on main, position_ids in the previous commit)
[There are still a few if/elses for cache-type-dependent logic in the model, but #29180 will standardize the cache API so that all caches match the static cache]

Updated performance numbers (i.e. still no degradation in throughput)

---

No specific opinon, just that using cache_position in the modeling code when use_cache=False is not very intuitive to me:

@fxmarty I agree here, that it should be renamed :) e.g. data_position would be more appropriate, as it may be used without cache. Precise variable names help us in the long run.

IMO we should move to having cache_positions in generate as this is more explicit, easier to maintain and less error prone.
More than that, these operation are input dependent while being vectorized they add complexity when there should not be!

@ArthurZucker I strongly disagree, more redundant inputs result in more bugs 🐛 From the full input_ids and pad_token_id we can derive attention_mask, position_ids, and cache_positions -- it's not like we are giving more options to the users, there is only one set of valid tensors given the first two. In other words, we give users three different paths to get things wrong! As a result, we have to add safety warnings and exceptions, such as:

1. warning if there is a pad token in the input but no attention mask
2. raising an exception if cache positions are not passed
3. raise an exception if the attention mask is not of the right type
4. warning if both attention mask and pad token are missing in generate
   (plus other error-prone combinations that we are not even checking 😬)

This is arguably much more work than enforcing the correct pairing in the models themselves! On top of that, other functions like generate or pipeline have to create the correct logic anyways (e.g. create attn mask; create position ids), so why not push the logic to the models and prevent these bugs from occurring everywhere? 🤗

My goal with this PR is to iron out future sources of bugs before we roll out these changes to other models. Better than well-explained interfaces is... no need for that interface at all! Without these changes, users have to learn how to prepare a new input tensor to use static caches.

In terms of complexity, it will always exist. The difference is whether it exists outside generate (=users can get it wrong, one more argument in generate), inside generate (=we have to maintain the logic, users can get custom generation loops wrong, one more argument in forward), or inside forward (=we have to maintain the logic). In my experience, if some input can be used incorrectly, it WILL be used incorrectly. Having few bugs is one of the strong points of transformers, which we should strive to keep. [In the same direction, I think we should get rid of attention_mask and position_ids, but that would be a breaking change 😉 There are so many issues in this repo due to these variables 😭 I remember that at some point you were also trying to get rid of position_ids]

In terms of performance, we can see that it is negligible in eager forwards and none in compiled static forwards. Again, if not implemented here, it will be implemented in generate, resulting in the same net performance.

[fxmarty]

I'm not a fan of this. This would be run at each forward. Couldn't we just increment somewhere, or not use get_seq_length at all in the forward?

Also given the issues we've had: did you check it works with compile (meaning: self.model.layers[0].past_key_values.get_seq_length() is correct inbetween each forward call in generate when the model is compiled with model = torch.compile(model, mode="reduce-overhead")? I removed the requirement on get_seq_length + seen_tokens as we were getting wrong values.

[gante]

@fxmarty yes, this version works! e.g. try running RUN_SLOW=1 py.test tests/test_cache_utils.py::CacheIntegrationTest::test_static_cache_greedy_decoding_pad_left_0_eager on the latest commit, which confirms that all dynamic, eager static cache, and compiled static cache get the same sensible outputs.

I'm not a fan of this. This would be run at each forward. Couldn't we just increment somewhere?

I am very much aligned with your comment! Indeed that was my original plan by adding seen_tokens to the original cache implementation. However, as you mentioned, we haven't found a way to make it work at compile time. Perhaps we shall add a TODO to revisit this in the future, after the issue you opened on the PyTorch side has been addressed?

[ArthurZucker]

• if part of the previous cache happens to be zero tensor (very rare in full precision, maybe not in half precision, we had a similar issue with llava when relying on tensors being zeros)
• cost is ok I guess
• it's totally implicit 😢 and that's less aligned with our overall philosophy.

[ArthurZucker]

Thanks for your great explanation.
TLDR; I think part of this choice comes down to being implicit vs explicit.
Your approach works, but it's implicit. Both have ➕ and ➖ for maintainability.
I am leaning a lot more towards the explicit for a few reasons:

1. If you generate the cache_positions in the model with torch.arange, the generations you will get with cuda graphs will be wrong precisely because the cache positions are not given as input. I tried this and it does not work.
2. ➕ Instead of relying on hidden seen_tokens and 3 different get_length, get_usable_lenght, get_max_length, you just need cache_positions in generate.
3. ➕ Stateless is less error prone, cache_positions is also pretty known from GPT-Fast.

4. From the full input_ids and pad_token_id we can derive attention_mask, position_ids, and cache_positions
   Is not true: if you want to do paged attention / packed training with a custom 4D attention mask or anything a bit outside the classic forward (any custom Cache class), then you cannot only rely on input_ids and pad token.

5. The generate function should keep track of where it is, for simplicity but also because otherwise your not really able to manipulate anything outside the forward of the model. So at the end of the day you have a loop in generate, which is the only public place where we use the past_key_values, but you don't let it handle the cache_positions which is a bit strange!

Let's keep the cache positions, gradually remove calls to get_seq_length, and explicitly pass the arguments to the model instead of relying on the model.

[ArthurZucker]

• if part of the previous cache happens to be zero tensor (very rare in full precision, maybe not in half precision, we had a similar issue with llava when relying on tensors being zeros)
• cost is ok I guess
• it's totally implicit 😢 and that's less aligned with our overall philosophy.

[ArthurZucker]

if the cache positions are using internally anyways, then that means we still have to create them at some point.

[gante]

@ArthurZucker as we talked on Slack, the latest set of commits makes cache_position follow the same logic as set for attention_mask or position_ids: enable both explicit [full control of forward] and implicit [lazy usage of forward] usage 🤗

NOTE: ignore the changes in prepare_inputs_for_generation for now, from what I'm seeing while enabling torch.compile we will need more touches there!

In summary, this PR:

1. Enables implicit cache_position, compatible with torch.compile;
2. Fixes StaticCache.get_seq_length() (needed for implicit cache_position in prepare_inputs_for_generation);
3. Deprecates the seen_tokens attribute in the other caches, as it is redundant with cache_position;
4. Slicing in prepare_inputs_for_generation with StaticCache now shares the original path, as opposed to having custom logic (a bunch of TODOs were added here, only feasible after future PRs);
5. Adds docstrings for cache_position in forward;
6. Adds several tests related to cache_position, to ensure it is working as expected.

[gante]

This restructure prioritizes cache_position, falling back to .get_seq_length() in its absence. This replaces seen_tokens, which is now deprecated.

Note that past_length [all seen tokens] and cache_length [tokens in the cache] are both needed, otherwise SinkCache won't work.

[gante]

This is already on main (see here), not sure why this shows up 👀

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[ArthurZucker]

@gante could you rebase / fix meges and I ll review!

[ArthurZucker]

Late review, not sure what is the status.
Overall the goal is to not have to pass cache positions, but the logic is very complicated, while asking users to pass the cache positions seems a lot simpler no?

[ArthurZucker]

nice 😉

[ArthurZucker]

that's a lot of work 😓
Does not seem like this is needed? Two cases:

1. No cache positions -> not using generate or not using cache positions -> use the DynamicCache, thus the previous code works for the past length
2. cache positions -> use them

[ArthurZucker]

IMO we should go towards everybody should pass the cache positions and we should not use past_seen_tokens = static_cache.get_seq_length().

[ArthurZucker]

does that also fix the ONNX export we had?

[ArthurZucker]

we broke AWQ a few times with this, let's check generation_config.cache_implementation ?

[gante]

Closing this PR and other cache PRs, as we want to move in the opposite direction (static cache behaving like the other caches)