feat: flex-attention backend

[cathalobrien]

This PR adds a flex attention backend to the attention layers.

Flex attention has been available in pytorch since ~2.6 but it wasn't stable for our use case. Now it seems more stable (I tested aifs-single v2 with torch 2.11 and it worked fine).

The specific use case here was that a user running on CPU installed the triton cpu build and wanted to use flex attention rather than SDPA.

I have also enabled support for the flash attention v4 backend of flex attention by checking if flash attention v4 is installed and passing the right kernel options to flex attention so it can use it.

Performance is slightly slower then flash attention v2 (0.34 it/s vs 0.38 it/s), but much faster then sdpa.

I added a correctness test which compares against sdpa, and a small benchmark test which compares all 3 attention kernels standalone.

anemoi-core/models/tests/integration/attention/test_benchmark_attention.py::test_attention_backend_benchmark[forward-global] Attention benchmark: window_size=None, dtype=torch.float16, shape=(1, 16, 40320, 32)
  flex: 36.25 ms
  sdpa: 27.32 ms
  flash: 21.16 ms
PASSED
anemoi-core/models/tests/integration/attention/test_benchmark_attention.py::test_attention_backend_benchmark[forward-sliding_window] Attention benchmark: window_size=1120, dtype=torch.float16, shape=(1, 16, 40320, 32)
  flex: 4.88 ms
  sdpa: 112.95 ms
  flash: 1.95 ms
PASSED
anemoi-core/models/tests/integration/attention/test_benchmark_attention.py::test_attention_backend_benchmark[fowrward & backward-global] Attention benchmark: window_size=None, dtype=torch.float16, shape=(1, 16, 40320, 32)
  flex: 28.18 ms
  sdpa: 26.16 ms
  flash: 21.11 ms
PASSED
anemoi-core/models/tests/integration/attention/test_benchmark_attention.py::test_attention_backend_benchmark[fowrward & backward-sliding_window] Attention benchmark: window_size=1120, dtype=torch.float16, shape=(1, 16, 40320, 32)
  flex: 5.04 ms
  sdpa: 114.19 ms
  flash: 1.96 ms
PASSED

[HCookie]

Awesome work, looks good to me

[HCookie]

Suggested change

	# One reason to use flex attention with the flash attewntion v4 backend, ratehr then using flash attention v4 directly, is
	# One reason to use flex attention with the flash attention v4 backend, rather then using flash attention v4 directly, is

[HCookie]

Suggested change

	# if flash attention is not available then the trion backend will be used for flex attention
	# if flash attention is not available then the triton backend will be used for flex attention