[Question] Masked SDPA attention kernel?

[EricLBuehler]

Hello MLX team!

Thank you for your excellent work here. We have been focusing on accelerating Candle performance on Metal and recently integrated some fast SDPA kernels from MLX!

I noticed that the scaled_dot_product_attention.metal file does not contain kernels for causal, non-decode step kernels. This seems to prohibit their usage in the prompt step for models such as Llama, and the fallback code for the SDPA op seems to indicate so.

Are there any plans to add masking support to the full SDPA kernel or current methods that have been taken to accelerate the prompt step?

[awni]

Are there any plans to add masking support to the full SDPA kernel or current methods that have been taken to accelerate the prompt step?

This is a work in progress. @jagrit06 has been working on that. We're hoping to ship something soon, but I can't give you an exact timeline yet.

[awni]

Closing as a dup of #129