Llama: fix custom 4D masks, v2
Author: poedator
Pull request: yes

this is an attempt to rebase  #29930 started initially by @gante 

Fixes the issue raised by @poedator in https://github.com/huggingface/transformers/pull/29753#issuecomment-2014530814.

Causal mask is now of shape [..., seq_len, full_len], as opposed to [..., full_len, full_len]. This means custom 4D attention masks are now the whole causal mask, so we don't need a sliced copy -- we can copy the whole thing :)

This PR also expands the support of custom 4D attention mask: we can pass both the full mask ([..., full_len, full_len]) or the partial mask ([..., seq_len, full_len]).
 - - - - - - - - 

as of 18.04.24 it is not passing the 4D mask tests because of the `_ignore_causal_mask_sdpa()` method, most recently edited in #30317 (merged today).

`tests/models/llama/test_modeling_llama.py::Mask4DTestHard::test_partial_stacked_causal_mask - ValueError: Incorrect 4D attention_mask shape: (1, 1, 12, 12); expected: (1, 1, 9, 12).`
apparently  `_ignore_causal_mask_sdpa()` expects that `attention_mask.shape[-2] == query_length`. This may only be true if the input_ids are contiguous, which is not always the case in some intended 4D mask applications.

`tests/models/llama/test_modeling_llama.py::Mask4DTestHard::test_stacked_causal_mask_static_cache - ValueError: Incorrect 4D attention_mask shape: (1, 1, 12, 16); expected: (1, 1, 12, 12).`
in this test  `_ignore_causal_mask_sdpa()` expects that `attention_mask.shape[-1] == key_value_length` which is set by `past seen tokens`. However, in the test I make this dimension equal to the static cache size, so that the mask always has same shape and the whole graph may be compiled and reused.

I hesitate to make edits to `_ignore_causal_mask_sdpa()` because there may be some greater context. 

Summoning @younesbelkada @ArthurZucker @gante  to help

Comments:
As s a solution, I added additional `expected_shapes` to `_ignore_causal_mask_sdpa()` and improved StaticCache detection code. 
Note: it is inconvenient to have StaticCache as layer.self_attn objects and other Caches as model-level object. Perhaps there may be a model-level plug to avoid referencing the layer levels.

Please review soon - I need this for my paper code. It's been broken for quite long now.

The LONG tests look OK.

all CI tests are green, SLOW tests were OK on my side yesterday

I noticed that mistral model support for 4D masks stayed broken after these fixes. So I added similar lines to `src/transformers/modeling_attn_mask_utils.py::_prepare_4d_causal_attention_mask_for_sdpa()`

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_30348). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

I added `Mask4DTestHard` tests (without static cache part) to `tests/models/mistral/test_modeling_mistral.py` to ensure that the 4d masks keep working in the models that use `_prepare_4d_causal_attention_mask_for_sdpa()`. These new tests would fail without the fixes from commit d488f35 just above.
Tested the SLOW tests for `./tests/models/mistral/` branch - all fine

Is there anything left to do before the merge is possible?
@gante @ArthurZucker 

> Let's remove unrelated changes!

sorry, but without these changes, the fixes and tests will not work. I looked for related PRs, all I found was #30476 but it is not fixing the relevant parts of the code.

I tried to follow Arthur's advice to streamline the path for the 4D masks and it seems to work. The relevant tests do pass. 
@ArthurZucker @gante , please review

I combined the 2 tests from common, which were very similar. Added tolerance - now Mixtral passes it OK.
@ArthurZucker, @gante - please see if it is good to merge now

Just waiting for the commits to be push to check the `done` parts and we can merge