Improve GPU performance

[metascroy]

🚀 The feature, motivation and pitch

Improve GPU performance on CoreML (can you use MLX?).

Alternatives

No response

Additional context

No response

RFC (Optional)

No response