Fix attn mask for static cache

[zucchini-nlp]

What does this PR do?

Fixes #30400.

It was found that when static cache returns key and values if "length=max_length", the zeros are not masked out. That is why the generation starts returning gibberish at larger max_new_tokens, and is more expressed in "sdpa" attention which just falls back to its internal causal mask.

All slow tests in the models are passing for me locally.

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

[ArthurZucker]

Not really sure about this fix, the target_length is probably not computed properly
What's weird about this is that we do take into account the previous zeros here:

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

Line 1096 in 7bfe577

causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)

with this anything that is bigger than the target length is mask, no matter the attention mask

[ArthurZucker]

Also Let's add tests in generation common or integration for llama at least.

[zucchini-nlp]

Ah i did not see it is supposed to be calculated here. The main error for teh code block in linked issue was in this line, where we skip all together the causal mask since it's batch_size=1 and no zeros in attention mask.

That is for sdpa attn but eager also had errors, even though not so big. Let me check with these new info and tests ofc :)

[ArthurZucker]

Thanks

[zucchini-nlp]

@ArthurZucker I found that indeed th causal mask was being calculated correctly and the gibberish output is only when we use SDPA and batch_size=1. Can be solved by removing the following so that we always rely on our own attn mask, which cannot be simply causal if static cache is used.

I am just not sure, if this line is crucial for Flash Attention 2? 🤔

if self.config._attn_implementation == "sdpa":
            # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument,
            # in order to dispatch on Flash Attention 2.
            if AttentionMaskConverter._ignore_causal_mask_sdpa(
                attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens
            ):
                return None

[gante]

@zucchini-nlp #30437 -- this PR also fixes the issue, and I think that is the way to go :D

The fix in this PR adds complexity outside the SDPA path and needs extra computations

[zucchini-nlp]

@gante I see, yeah the cropping of the cache seems a better solution