feat: add varlen attention on cpu

[michaelfeil]

What does this PR do?

This PR brings varlen-flash-attention to CPU/Metal. (Its not softmax-fused / Flash), but at least its not-padded, so we don't do OOM.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the documentation guidelines.
• Did you write any new necessary tests? If applicable, did you include or update the insta snapshots?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[kozistr]

aside from one thing, everything else looks good to me!

[kozistr]

i was just wondering that it looks like causual_mask and window_mask below are always fp32 type while attention_scores could be fp16. I'm not sure if I'm right, it might fail due to a type mismatch!

[michaelfeil]

Thanks, you are right. i figured that out too.

[michaelfeil]

In this case its always fp32, but for apple metal backend it could be fp16 afaik.

[michaelfeil]

I am going to mark this PR as draft. I implemented a pretty fast attention primitive here: huggingface/candle#3250 once that is merged (which i am eagerly waiting for) we can do a simple copy of the function here (without the tests).