13 Commits

Author	SHA1	Message	Date
Awni Hannun	111f1e71af	Faster contiguous gather for indices in the first axis (#2552) ... * faster contiguous gather for indices in the first axis * work per thread > 1 * angelos suggestion for scales / biases	2025-08-28 21:26:30 -07:00
Awni Hannun	827003d568	fix METAL quantization in JIT (#2553)	2025-08-28 18:26:25 -07:00
Angelos Katharopoulos	4a9b29a875	MoE backward improvements (#2335)	2025-07-07 17:59:53 -07:00
Angelos Katharopoulos	99eefd2ec0	Gather mm new kernel and small refactoring (#2040)	2025-04-14 16:37:36 -07:00
Awni Hannun	de5f38fd48	Custom logsumexp (#2028) ... * initial custom logsumexp * more tests * comments + fix	2025-03-31 07:36:55 -07:00
Awni Hannun	b7c9f1d38f	scatter axis + gather axis primitives (#1813) ... * scatter axis + gather axis primitives * add transforms * comment	2025-01-31 20:48:08 -08:00
Awni Hannun	30bbea2f08	Add gemv masked to JIT plus some fixes (#1310) ... * add gemv masked to JIT plus some fixes * some cleanup * add utils * fix * fix 2 * more cleaning * fix * remove unused mps matmul support * one more nit * revert	2024-08-07 13:38:07 -07:00
Alex Barron	a3c287354f	Fast Hadamard Transform (#1249) ... * Working hadamard for powers of 2 * working for m*2^k * add scale and check contiguity * add size check * clean up * fix test * add grads + vmap * gpu only * skip on linux * test typo * add cpu impl * remove gpu only tests * fix linux build + add is_equivalent	2024-07-09 20:39:01 -07:00
Alex Barron	dd7d8e5e29	Add Quantized Ops to the JIT (#1204) ... * JIT for quantized ops * remove unused imports * address comments * fix imports * second attempt to fix imports --------- Co-authored-by: Alex Barron <abarron22@apple.com>	2024-06-12 09:47:12 -07:00
Alex Barron	27d70c7d9d	Feature complete Metal FFT (#1102) ... * feature complete metal fft * fix contiguity bug * jit fft * simplify rader/bluestein constant computation * remove kernel/utils.h dep * remove bf16.h dep * format --------- Co-authored-by: Alex Barron <abarron22@apple.com>	2024-06-06 12:57:25 -07:00
Awni Hannun	7e26fd8032	Option to JIT steel gemm / conv (#1139)	2024-05-23 18:07:34 -07:00
Awni Hannun	0189ab6ab6	More jitting (#1132) ... * docs + circle min size build * jit scan, arange, softmax * add sort * jit reductions * remove print * fix deps * clean includes / nits	2024-05-23 16:23:44 -07:00
Awni Hannun	226748b3e7	JIT compile option for binary minimization (#1091) ... * try cpp 20 for compile * unary, binary, ternary in jit * nits * fix gather/scatter * fix rebase * reorg compile * add ternary to compile * jit copy * jit compile flag * fix build * use linked function for ternary * some nits * docs + circle min size build * docs + circle min size build * fix extension * fix no cpu build * improve includes	2024-05-22 12:57:13 -07:00