-
Notifications
You must be signed in to change notification settings - Fork 26.2k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Generate: fix SinkCache
on Llama models
#30581
Conversation
SinkCache
on Llama models
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. |
elif self._cos_cache.shape[0] < self.window_length: | ||
self._cos_cache = torch.cat([self._cos_cache, cos[0, ...]], dim=0) | ||
self._sin_cache = torch.cat([self._sin_cache, sin[0, ...]], dim=0) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just for my own understanding of how the cache is meant to work, I have two Qs:
-
Values passed in on
update
call
if we callupdate
withsin
andcos
passed in, is the cache keeping old values + new values i.e.self._cos_cache[:self._cos_cache_prev.shape[0]]
are the old values andself._cos_cache[self._cos_cache_prev.shape[0]:]
is the new values, or the passed incos
is just the new values to be appended? -
Window length
Is the assumption here that the window length is constant once the cache is created?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- The values passed in
cos
are new values to be appended. In RoPE models,sin
andcos
are a constant with shape[config.max_position_embeddings, rope_embedding_dims, config.hidden_size // config.num_attention_heads]
. However, with the compile-optimized modeling code, we only materialize the needed parts of these matrices, with shape[0] =input_ids.shape[1]
= input sequence length. SinceSinkCache
needs access to allsin
andcos
values up to shape[0] =self.window_length
when going beyond the window length, this cache was created.
Alternatively, we could pass the the model config to compute the full sin
and cos
, but that would be (IMO) an ugly interface (we would have to use the model config to instantiate a RoPE layer inside the cache, to then compute these values and discard the layer).
- Yes.
SinkCache
is a fixed-length cache -- its purpose is to be used withself.window_length
<config.max_position_embeddings
, while enabling coherent outputs beyond full sequence length =self.window_length
. In other words, coherent long outputs with a relatively short cache :) Its limitation is that it can only recall content back up to the size of the window length, it quickly forgets things.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Got it - thanks for taking the time to write this up and explain!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for fixing!
What does this PR do?
SinkCache
has been broken on Llama and Llama-based models since we released the static cache update (v4.38). Now that we are happy with the state of the static cache (#30476), we can move on to fix what we broke along the way 🤗In a nutshell, the static cache rework changed the
sin
andcos
tensors passed around, from the full set of values for all possible positions (up toconfig. max_position_embeddings
) to the values used in a forward pass alone. This is a non-negotiable change to achieve top compiled performance.However,
SinkCache
needs access to the wholesin
andcos
tensors, and they are not trivial to compute from scratch in the cache instance (it would need access to the RoPE class, creating a cyclical dependency). TheSinkCache
instance was changed to hold a cache [meta cache 🤯 ] ofsin
andcos
as it sees them, rebuilding the full tensor internally. Having the full tensor rebuilt, it can operate as expected.tests/test_cache_utils.py::CacheIntegrationTest::test_sink_cache_hard
is fixed as a result of this PR. All other sink cache tests were passing (they were not using llama 👼 )