Adding metal kernels for the gated delta nets. - #4020

Open

tpegolotti wants to merge 63 commits into

mainfrom

gated-delta-update

Commits on Aug 26, 2026

add basic skeleton for gated delta update forward
tpegolotti
committed
a134bd5
add skeleton for metal kernel - compile and run ok
tpegolotti
committed
46e5ea2
copied implementation over from mlx_lm as baseline
tpegolotti
committed
b6248b2
make shapes consistent with reference
tpegolotti
committed
cae3597
Add skeleton for chunkwise implementation
tpegolotti
committed
1aee8ca
base chunkwise implementation
tpegolotti
committed
ca066af
Delete local_files directory
tpegolotti
committed
e6f8798
simdgroup matrices work for C=8
tpegolotti
committed
2736da7
ran pre-commit
tpegolotti
committed
e4d6b9b
full simdgroup main loop
tpegolotti
committed
ac42ef6
fused wy into gated computation
tpegolotti
committed
ad8a59f
add elementwise macros
tpegolotti
committed
95b54cd
add gated delta benchmark
tpegolotti
committed
e84a5e5
fix name
tpegolotti
committed
58cfbbb
fix default C value
tpegolotti
committed
432c63e
update bench script
tpegolotti
committed
fb87390
Remove old chunkwise implementation
tpegolotti
committed
38ff71a
Start fallback implementation
tpegolotti
committed
59a182c
Add contiguous memory copy and padding to handle generic T to get it to work on Qwen
tpegolotti
committed
cf9a43e
Update gated delta benchmarking. Added Qwen3.5 dimensions
tpegolotti
committed
3589c21
Added first nax version
tpegolotti
committed
ced278f
improve inverse computation
tpegolotti
committed
2092f20
Fuse status update with output computation
tpegolotti
committed
4f9f765
Half matmuls in invers by fusion
tpegolotti
committed
dba88c8
Remove KP
tpegolotti
committed
905fbb7
Add log decay for stability
tpegolotti
committed
79ffa4a
Make ensure row contiguous in eval_gpu
tpegolotti
committed
5de378c
Add fallback
tpegolotti
committed
4985b4a
Reverted invert
tpegolotti
committed
39bf1cb
Change threadgroup grid
tpegolotti
committed
2ea0b97
Remove one inverse
tpegolotti
committed
d8218fa
Improve inverse
tpegolotti
committed
63dfb13
Pre PR changes: remove C as parameter, remove explicit padding, fix typing
tpegolotti
committed
3de7cef
remove log space for C=8 and force cleanup inline
tpegolotti
committed
be8bda3
update typing
tpegolotti
committed
62cab7b
added back log space
tpegolotti
committed
ad06dcf
cleanup .metal file
tpegolotti
committed
ac90714
Add clamping back
tpegolotti
committed
55ab541
Back to Horner for NAX
tpegolotti
committed
db41d8e
Add test and benchmark
tpegolotti
committed
1169017
Fix sign
tpegolotti
committed
2a026e8
Increase nax atol
tpegolotti
committed
1a493c4
Removed scaling from input
tpegolotti
committed
5bf71d0
Add mask support in fallback
tpegolotti
committed
d00103b
Fix non metal build link errors
tpegolotti
committed
8c11d66
Fixing more linker errors
tpegolotti
committed
59f2ba9
Changed lambda function to macro
tpegolotti
committed
0c236a0
adding nax files
tpegolotti
committed
6805a09
Updating kernel getters
tpegolotti
committed
db69d20
Fix format error
tpegolotti
committed
c083e9c
add jit getter
tpegolotti
committed
f39122d
adding torch check on test
tpegolotti
committed
c76d84c
adding torch check on test
tpegolotti
committed
f7ed8f3
Fixing jit linkin error
tpegolotti
committed
cdd4afd
add support for different head sizes in fallback
tpegolotti
committed
7dcbb1f
Remove debug macros
tpegolotti
committed
1fba67f
Fallback output dtype corresponds to input dtype
tpegolotti
committed
0a09840
Addressing PR comments
tpegolotti
committed
3a895f0
Add force_fused option to scaled_dot_product_attention (#4185)

hojin12312
authored and
tpegolotti
committed
262d262
Add a fused full-attention path for head_dim 256 on NAX devices (#3842)

authored and
tpegolotti
committed
9d57f23
use fallback now considers new head dimensions
tpegolotti
committed
60af0df
Run formatter
tpegolotti
committed
8b38e31

Commits on Aug 27, 2026

Addressing comments
tpegolotti
committed
c7e1a2a