[CUDA] Support array mask in SDPA

[zcbenz]

Note that cuDNN does not support boolean masks, so we have to convert boolean masks to additive masks with where(mask, full_like(mask, 0), full_like(mask, -inf)), which has some performance penalty. (PyTorch does the same thing too.)

What cuDNN does support is setting padding mask directly: we pass the sequence lengths and cuDNN will apply padding masks automatically, and it works together with the set_causal_mask flag. I don't know how much performance gain this approach brings, but I think it worths a try as a future work.

[awni]

Looks great!

As a general guideline we should work towards avoiding in dispatching differently based on the back-end because it breaks the ability to exporting from machine to another (e.g. export on cuda would not work on Metal). The fast primitives are one place where we break this guideline a lot (e.g. cpu vs gpu). Just elaborating for future reference as we may want to push the mask -> float into the primitive.