Eric Buehler e2b6b367fa Add some fast Metal MLX SDPA kernels (#2584) ...

* Add some fast Metal MLX SDPA kernels (#32)

* Sketch the sdpa kernel

* Add full sdpa kernel,

* Add test

* Add vectorized kernel for decoding

* Update tests

* Add some docs

* Fix sdpa_vector names

* Add softcapping for vectorized sdpa

* Add softcapping for full sdpa

* Add support for head dim 32, 96, 256

* Add support for head dim 32, 96, 256

* Update docs

* Add update notice

* Clippy and format

* Conditional compilation for bf16

* Use it in quantized llama

* Some review comments

* Use set_params!

* Remove unused

* Remove feature

* Fix metal sdpa for v stride

* Remove comma

* Add the dim method to layout and shape.

---------

Co-authored-by: Laurent <laurent.mazare@gmail.com>

2024-11-05 09:28:00 +01:00

..

batch_norm.rs

Bug Fix: When converting a tensor to a variable, clone if the tensor is already a variable. (#2124)

2024-04-29 11:21:53 +02:00

group_norm.rs

Move the test-utils bits to a shared place. (#619)

2023-08-27 09:42:22 +01:00

kv_cache.rs

Add a RotatingKVCache. (#2493)

2024-09-23 13:14:32 +02:00

layer_norm.rs

Enable the new layer-norm. (#2213)

2024-05-24 16:48:21 +02:00

loss.rs

Add Binary Cross Entropy With Logit Loss to nn crate (#1157)

2023-10-23 17:12:44 +01:00

one_hot.rs

Add one-hot/cold encoding (#1489)

2024-01-01 11:18:40 +01:00

ops.rs

Improved launch config for layer-norm/rms-norm. (#2591)

2024-11-04 10:42:18 +01:00

optim.rs

Tweak the VarMap set type. (#1758)

2024-02-25 20:50:08 +01:00

rnn.rs

Add a GRU layer. (#688)

2023-08-31 08:43:10 +01:00

sdpa.rs

Add some fast Metal MLX SDPA kernels (#2584)

2024-11-05 09:28:00 +01:00