Here according to my understanding, key_states and value_states are once and for all computed based on encoder hidden states when we are doing cross-attention, so we can cache them ideally, but I think these are a little bit different from the existing caches because we don't need to update the cache in every generation step, we just have to do it once in the first step, so we could either try creating another Cache class and in that cases we need to pass in two caches(one for self attention and one for cross attention), even if we can create a cache class to wrap both, we still need to modify the current get_cache logic because it definitely will need more parameters to initiate the cache, and I don't know if the use case is general enough to create this new cache class. Or we can just initiate and update the cross attention kv cache within every layer like recurrent gemma, but this will require manually reset the cache between generations because the current logic in generation seems not considering the case when there is an inherent cache possessed by the model. I am personally for the latter solution, however I think both will need some changes in the generation utils file, and not sure which way might make dynamo or cudagraphs unhappy, what do you think is best @ArthurZucker

Your understanding is correct: we compute the k/v states for the cross-attention once in the first forward pass, and then save them to cache. We would indeed want to cache the k/v states to avoid re-computing them at every step.

What do you think about the proposed design in this PR? #28931 (comment) Similar to your solution, it uses two cache's: one for the self-attention, and one for the cross-attention. Note that the cache specifics are out-of-date, given the recent changes to the cache API, but the high-level design remains valid!

Yeah, it would be great if two StaticCache would work, but the issue is in that case we can not tell whether we are in the first generation step or not because the shape of cache now is always fixed to hold the maximum generation size, and use get_seq_length(to see if we have processed any tokens yet during tracing) in branch condition will indeed cause graph breaks

my current solution is use two StaticCache and add another flag in attention layer to mark whether we are doing the first generation step, this will cause a recompile in the second step just like in llama and mistral but should work fine for the subsequent steps, but it's still nasty because it breaks the current stand-alone cache design, maybe I should indeed create another new OneShotCache class which hides the flag inside? cc @ArthurZucker
@gante @sanchit-gandhi

The current solution utilizes a new OneShotCache, and a tuple of two caches is expected for conditional generation

I wonder if there is a better way to do this, because in the encoder-decoder scenario we need a tuple of two caches here according to the current design, but this seems hardcode and easy to break

Actually we can always use OneShotStaticCache for kv cache here for simplicity, beccause using Dynamic Cache won't give us memory benefits on cross-atten kv cache