Sink Cache Attention Scores are strange. CausalMask seems not working. about transformers HOT 5 OPEN

Supplementary:

It works with DynamicCache.

So it must be something wrong with SinkCache and relevant control code.

from transformers.

It's been a bit since I worked on this, but I think that -self.window_length + self.num_sink_tokens + key_states.shape[-2] >= 0 is not really possible.

• window_length is the max. size of the cache, e.g. 1024.
• num_sink_tokens is some (usually small) positive integer, e.g. 4
• key_states.shape[-2] is the size of the new additions to the cache.

In the code here:

We're in the "Shifting cache" phase, i.e. the cache already exists, and now we're adding enough tokens to make it overflow. However, if it already exists, then I think (I'm not 100% on this) we always add 1 new generated token, i.e. key_states.shape[-2] is 1. So I think a non-negative value can only happen if the num_sink_tokens >= window_length - 1, which is not normal behaviour.

However, if it's somehow possible to, when the cache already exists, add a bunch of tokens in one go, then I think it would be possible to mess this up. Then, the keys_to_keep should really be empty (as we're skipping way ahead and keeping no tokens), but the overflow of -self.window_length + self.num_sink_tokens + key_states.shape[-2] >= 0 into the positives is allowing some keys to stay. Then the new tokens will get appended and we'll accidentally get a cache that's too large here:

But I think that should probably cause a pretty easy-to-spot crash as the cache is now bigger than the window size, which should not be possible.

• Tom Aarsen

from transformers.

cc @gante @ArthurZucker

from transformers.

Have not worked on the sink cache so will let @gante answer here!

from transformers.

In cache_utils.py, I noticed that
keys_to_keep = self.key_cache[layer_idx][ :, :, -self.window_length + self.num_sink_tokens + key_states.shape[-2] : ]
might go wrong when -self.window_length + self.num_sink_tokens + key_states.shape[-2] >= 0
Not sure is it relevant

from transformers.