Repository navigation

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

Merged

tpegolotti merged 64 commits into

mainfrom

gated-delta-update

Sep 15, 2026

Merged

Adding metal kernels for the gated delta nets.#4020
tpegolotti merged 64 commits into
mainfrom
gated-delta-update

Commits on Sep 14, 2026
- add basic skeleton for gated delta update forward
  tpegolotti
  committed
  c168d5c
- add skeleton for metal kernel - compile and run ok
  tpegolotti
  committed
  1f00f3c
- copied implementation over from mlx_lm as baseline
  tpegolotti
  committed
  5cbd03f
- make shapes consistent with reference
  tpegolotti
  committed
  7a81134
- Add skeleton for chunkwise implementation
  tpegolotti
  committed
  381dbd5
- base chunkwise implementation
  tpegolotti
  committed
  8afda35
- Delete local_files directory
  tpegolotti
  committed
  89ac94c
- simdgroup matrices work for C=8
  tpegolotti
  committed
  7266c2e
- ran pre-commit
  tpegolotti
  committed
  76e3c4b
- full simdgroup main loop
  tpegolotti
  committed
  8507bf1
- fused wy into gated computation
  tpegolotti
  committed
  568338a
- add elementwise macros
  tpegolotti
  committed
  77f62ff
- add gated delta benchmark
  tpegolotti
  committed
  5efb380
- fix name
  tpegolotti
  committed
  5aac26e
- fix default C value
  tpegolotti
  committed
  ef6a3e8
- update bench script
  tpegolotti
  committed
  a41fdec
- Remove old chunkwise implementation
  tpegolotti
  committed
  be961cd
- Start fallback implementation
  tpegolotti
  committed
  0251091
- Add contiguous memory copy and padding to handle generic T to get it to work on Qwen
  tpegolotti
  committed
  fc9f3a7
- Update gated delta benchmarking. Added Qwen3.5 dimensions
  tpegolotti
  committed
  22e285f
- Added first nax version
  tpegolotti
  committed
  53c3a50
- improve inverse computation
  tpegolotti
  committed
  65dc671
- Fuse status update with output computation
  tpegolotti
  committed
  7c1edbe
- Half matmuls in invers by fusion
  tpegolotti
  committed
  1b54e78
- Remove KP
  tpegolotti
  committed
  b688a77
- Add log decay for stability
  tpegolotti
  committed
  1a05e31
- Make ensure row contiguous in eval_gpu
  tpegolotti
  committed
  bc82ef6
- Add fallback
  tpegolotti
  committed
  2b8f7f3
- Reverted invert
  tpegolotti
  committed
  fa5eef0
- Change threadgroup grid
  tpegolotti
  committed
  84ad09e
- Remove one inverse
  tpegolotti
  committed
  ee185f5
- Improve inverse
  tpegolotti
  committed
  f465b66
- Pre PR changes: remove C as parameter, remove explicit padding, fix typing
  tpegolotti
  committed
  3461841
- remove log space for C=8 and force cleanup inline
  tpegolotti
  committed
  c34f133
- update typing
  tpegolotti
  committed
  69f7df7
- added back log space
  tpegolotti
  committed
  be74c6e
- cleanup .metal file
  tpegolotti
  committed
  1354488
- Add clamping back
  tpegolotti
  committed
  0c9db8b
- Back to Horner for NAX
  tpegolotti
  committed
  11fe0d8
- Add test and benchmark
  tpegolotti
  committed
  7eece10
- Fix sign
  tpegolotti
  committed
  c477ff2
- Increase nax atol
  tpegolotti
  committed
  8635a64
- Removed scaling from input
  tpegolotti
  committed
  a81605c
- Add mask support in fallback
  tpegolotti
  committed
  aa8362b
- Fix non metal build link errors
  tpegolotti
  committed
  4989ca3
- Fixing more linker errors
  tpegolotti
  committed
  b08132a
- Changed lambda function to macro
  tpegolotti
  committed
  87482d0
- adding nax files
  tpegolotti
  committed
  9553ccc
- Updating kernel getters
  tpegolotti
  committed
  9dd0fe0
- Fix format error
  tpegolotti
  committed
  b736b1a
- add jit getter
  tpegolotti
  committed
  fb245b6
- adding torch check on test
  tpegolotti
  committed
  348d12f
- adding torch check on test
  tpegolotti
  committed
  c610b6c
- Fixing jit linkin error
  tpegolotti
  committed
  ce30875
- add support for different head sizes in fallback
  tpegolotti
  committed
  3309997
- Remove debug macros
  tpegolotti
  committed
  1693325
- Fallback output dtype corresponds to input dtype
  tpegolotti
  committed
  53655e6
- Addressing PR comments
  tpegolotti
  committed
  2f6f47f
- Add force_fused option to scaled_dot_product_attention (#4185)
  
  DT-HJ
  authored and
  tpegolotti
  committed
  2303711
- Add a fused full-attention path for head_dim 256 on NAX devices (#3842)
  
  authored and
  tpegolotti
  committed
  b7e9ed2
- use fallback now considers new head dimensions
  tpegolotti
  committed
  7b54f80
- Run formatter
  tpegolotti
  committed
  552d54e
- Addressing comments
  tpegolotti
  committed
  5ae919d
- Removed tensor ops header
  tpegolotti
  committed
  1b49f19

Commits

Commits on Sep 14, 2026