Interface with new Metal primitives

[maleadt]

Uses JuliaGPU/Metal.jl#854

[maleadt]

I did wonder about moving this into Metal.jl though. The MPS/MPSGraph interfaces used by the extension aren't exactly stable, while the NNlib ones presumably are? Although looking at e.g. #718 it seems like the interfaces are changing here too, so not sure what's best.

[CarloLucibello]

You can consider the current interface of NNlib stable. The only deprecation planned is for dot_product_attention, which will be replaced by a new method that doesn't returns the weight along the output (so that we can support flash attention).

I think having these methods here nicely parallels the CUDA/cuDNN/NNlib interaction, so I think it is very reasonable.

Buildkite tests are failing...

[maleadt]

Yes, Metal.jl hasn't been released with these changes yet. But I wanted to put up the PR already.

[maleadt]

You can consider the current interface of NNlib stable.

Then I think we should consider moving the extensions into the GPU back-ends, no? The code they call is far less stable than the NNlib interfaces (e.g. the low-level cuDNN or MPS APIs).

[CarloLucibello]

I think the current division, which also this PR follows, is fine. The dispatch structure of NNlib's methods can be understood by just looking at NNlib's repo.

[CarloLucibello]

tag a Metal.jl release?