zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-22 02:58:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity
1,440 Commits 42 Branches 73 Tags
9cbb1b014821797c3910e883feb65c634ae604ac
Commit Graph

443 Commits

Author SHA1 Message Date
Alex Barron
0fd2a1f4b0 Custom Metal Kernels from Python (#1325) 
* start

* simple kernels working

* restructure

* inverse example working

* docs + fixes

* missing file

* fix imports

* address comments

* add docs + fix test

* Review comments + refactor to a single function

* update docs

* remove hashing

* fix contig bug in test

* back to a class

* trailing whitespace

* fix tests

* match c++ and python apis

* add link + make args kw_only
 2024-08-22 13:46:29 -07:00
Awni Hannun
df3233454d 2d gather specialization (#1339) 2024-08-22 10:48:24 -07:00
Awni Hannun
d40e76809f Fix rope (#1340) 
* add test

* fix rope

* fix test
 2024-08-20 17:37:52 -07:00
Awni Hannun
bb1b76d9dc RoPE with frequencies as optional input (#1337) 
* start rope with freq input

* rope with frequencies

* nits

* fix bug

* fix bug + test

* cleanup

* optional base
 2024-08-19 18:30:50 -07:00
Angelos Katharopoulos
9d26441224 Fix contiguity check (#1336) 
Co-authored-by: Alex Barron <abarron22@apple.com>
 2024-08-19 16:05:06 -07:00
Awni Hannun
f12f24a77c fix compiling with space in paths (#1332) 2024-08-15 16:39:24 -07:00
Alex Barron
99bb7d3a58 GPU mx.sign for complex64 (#1326) 2024-08-14 07:54:53 -07: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
Awni Hannun
58d0e199e1 add bfloat conv for windograd (#1306) 
* add bfloat conv for windograd

* accumulate in fp32

* accumulate in fp32

* accumulate in bf16
 2024-08-05 15:51:13 -07:00
Awni Hannun
43ffdab172 fix rope and random (#1301) 
* fix rope and random

* comment
 2024-07-31 16:18:25 -07:00
Awni Hannun
40b6d67333 Fixes for large arrays with a few ops (#1299) 
* fixes for large arrays with a few ops

* fix bug

* fix all of copy
 2024-07-30 17:18:39 -07:00
Alex Barron
c52d1600f0 Fused Affine Quantize/Dequantize ops (#1282) 
* Add fast affine dequantize

* add full quantize kernel

* fused kernel with scale/bias computation

* fix docstring

* fix no jit error

* fix test

* test fix

* reduce fast api to only affine_quantize
 2024-07-29 15:11:38 -07:00
Jagrit Digani
7f914365fd Fix GPU sort for large arrays (#1285) 
* Fix GPU sort for large arrays
 2024-07-24 14:37:10 -07:00
Alex Barron
c34a5ae7f7 Fix bfloat16 Hadamard (#1283) 
* fix bfloat16 hadamard

* add scale

* review comments

---------

Co-authored-by: Alex Barron <abarron22@apple.com>
 2024-07-23 14:54:43 -07:00
Awni Hannun
e2aa6ec8ae some fixes (#1281) 2024-07-23 11:49:05 -07:00
Tim Gymnich
6307d166eb Fix overflow / underflow handling for expm1f (#1278) 
* Fix overflow / underflow handling for expm1f

* update tests
 2024-07-23 07:29:06 -07:00
Awni Hannun
1fba87b0df Fix leak with multi-output primitives (#1274) 
* fix leak with multi-output primitives

* hopefully an actual fix
 2024-07-23 06:34:18 -07:00
Cheng
2f83d6e4b7 Do not release buffers on exit (#1142) 2024-07-15 15:12:24 -07:00
Angelos Katharopoulos
5c1fa64fb0 Custom transforms (#1246) 2024-07-10 18:00:01 -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
Angelos Katharopoulos
03cf033f82 Fix reshape copy bug (#1253) 2024-07-07 21:37:00 -07:00
Alex Barron
2615660e62 Fix strided sort bug (#1236) 
* Use output strides in sort kernel

* fix zero strides bug
 2024-06-26 14:32:11 -07:00
Awni Hannun
56c8a33439 Get metal version from xcode (#1228) 
* get metal version from xcode

* typo

* fix
 2024-06-26 07:02:11 -07:00
Jagrit Digani
2d6cd47713 Masked gemv (#1211) 2024-06-14 09:52:26 -07:00
Awni Hannun
31e134be35 Build for macOS 15 (#1208) 
* Build for macos 15

* metal32 as well

* comment

---------

Co-authored-by: Awni Hannun <Awni Hannun>
 2024-06-13 13:31:44 -07:00
Fangjun Kuang
f20e97b092 minor fixes (#1194) 
* minor fixes

* fix build errors
 2024-06-12 22:06:49 -07:00
Alex Barron
934683088e Refactor JIT for unary/binary/ternary ops (#1206) 
* refactor unary/binary/ternary ops

* get_primitive_string util

---------
 2024-06-12 14:22:12 -07:00
Awni Hannun
de2b9e7d0a Fix kernel deps to reduce build times (#1205) 2024-06-12 11:17:39 -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
Awni Hannun
df964132fb fix scatter + test (#1202) 
* fix scatter + test

* fix test warnings

* fix metal validation
 2024-06-11 14:35: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
578842954c fix jit scan when output doesn't have primitive (#1190) 2024-06-06 07:24:58 -07:00
Awni Hannun
496315fe1d Fix scan (#1188) 
* fix scan

* improve grid size

* fix cpu cummax
 2024-06-05 14:21:58 -07:00
Awni Hannun
83b11bc58d Fix Metal API validation for empty concat (#1183) 2024-06-04 13:17:08 -07:00
Alex Barron
375a8bbdcc Add some internal GPU apis (#1177) 
* Add unary/binary/ternay/slice/concat internal GPU ops

* add pad internal op

* formatting + no_cpu fix
 2024-06-04 09:24:26 -07:00
Awni Hannun
ea9090bbc4 Add view op (#1179) 
* add view primitive

* nit

* fix view
 2024-06-04 08:05:27 -07:00
Alex Barron
4d485fca24 Add defines include (#1176) 
Co-authored-by: Alex Barron <abarron22@apple.com>
 2024-06-03 09:50:10 -07:00
Brian Keene
1865299a30 Metal shaders for memory efficient self attention on large sequences (#964) 
* Metal shaders for efficient self attention on large sequences

Updated fast attention: GEMM-ified with Steel primitives
Uses flash attention 1 for scale correction

* more compiler silencing

* Address rebase issues

* Templatize kernel instantiation, revise cpu bindings

* Safer writes to output

* Permit batch size > 1

* Numerical fixes for sdpa self attention

* Re-enable test, remove unused variable

* add benchmarking script

* Disable sdpa prior to perf tuning, and simplify tests for per-patch CI
 2024-06-03 09:16:19 -07:00
Jagrit Digani
76b6cece46 Fix multi-block sort stride management (#1169) 
* Fix multi-block sort stride management

* Add seed to tests
 2024-05-31 11:10:54 -07:00
Jagrit Digani
9f0df51f8d Fix matvec vector stride bug (#1168) 2024-05-29 12:18:28 -07:00
Awni Hannun
e7a2a3dcd1 Fix a couple bugs (#1161) 
* fix jit reduce for RMS norm

* make strides a single buffer

* better eval error message

* fix compiling with inf and bf16

* fix cpu compile with bf16
 2024-05-28 15:18:18 -07:00
Awni Hannun
7e26fd8032 Option to JIT steel gemm / conv (#1139) 2024-05-23 18:07:34 -07:00
Jagrit Digani
eab2685c67 Float mask update (#1152) 
* Float mask update

* Update CPU impl
 2024-05-23 17:20:44 -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
Rifur13
9401507336 Add groups to 2-D convolutions (#1129) 
* Added groups to 2-D convolutions. Only implemented for **some** specializations.

Also fixed 1D grouped convs with different kernel strides and added more tests.

* fix channels condition
 2024-05-22 20:01:44 -07:00
Awni Hannun
e110ca11e2 Fix offset bug for device buffers (#1151) 
* fix bug with large offsets for buffers

* add a test

* remove test as its too big for small machine
 2024-05-22 15:50:05 -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
Awni Hannun
d568c7ee36 Rename block sparse (#1149) 
* block_sparse_mm to gather_mm

* rename

* nit

* nit
 2024-05-22 07:48:34 -07:00
Angelos Katharopoulos
da83f899bb Improve qvm speed (#1140) 2024-05-20 09:20:44 -07:00
Awni Hannun
23406c9e9e Choose the right MLX bf16 for extensions (#1135) 
* default to custom bf

* choose right bf

* fix extensions

* fix circle conf
 2024-05-17 15:09:28 -07:00