Batched flex attention to reduce max memory usage

[cathalobrien]

🚀 The feature, motivation and pitch

I'm working with flex attention. For longer sequence lengths e.g. 40320, I get OOMs at the line "scores = (query @ key.transpose(-2, -1)).to(dtype=working_precision)" where it tries to allocate a 40320^2 matrix of floats (48GB) all at once. Would it be possible to batch this score calculation to reduce the maximum memory usage, like what the memory-efficient attention algorithm does. Thanks

Alternatives

No response

Additional context

No response

### Tasks

[cathalobrien]

never mind, seems like the problem was I had to compile the function.

meta-pytorch/attention-gym#15