Can KV reordering be avoided with `flex_attention`?

**Is your feature request related to a problem? Please describe.**

In the `flex_attention` integration, we take `token_positions` into account by reordering `key, value` for every call. This adds time and especially memory requirements.

It also does not work with `sliding_window_size` given, even though this is easily supported in `flex_attention` if `token_positions` is not given.

It would be much simpler to use a `mask_mod` of the form
`q_idx + input_pos >= token_positions[kv_idx]`

Here, `input_pos, token_positions` are captured tensors, and `token_positions.ndim == 1`. We cannot use an overall fix `token_positions`, because it changes with every call. On the other hand, `input_pos` could be folded into `token_positions`.

**Describe the solution you'd like**

We need a solution where `token_positions` becomes an input to the kernel or its `BlockMask`. This is not doable as in the simple examples, where captured tensors must be scalar (so `ndim==0`). But it may work with what they call BlockMask conversion, where the BlockMask can be changed without kernel re-compilation.

Start by tracing the flex_attention integration and how token_positions currently triggers key/value reordering. Read the mask_mod examples and the BlockMask conversion path, then determine whether token_positions can be passed as a kernel or BlockMask input without recompilation. Done means avoiding per-call KV reordering while preserving token-position behavior and supporting sliding_window_size.

Written by the indexing model from the issue text.