Add a fused full-attention path for head_dim 256 on NAX devices - #3842

Add a fused full-attention path for head_dim 256 on NAX devices#3842

zcbenz merged 6 commits into

ml-explore:mainml-explore/mlx:mainfrom

wyanzhao:nax-256-dsplitwyanzhao/mlx:nax-256-dsplitCopy head branch name to clipboard