Use unroll_count(4) for the NAX attention Q@K.T loop#3843

wyanzhao wants to merge 1 commit into

ml-explore:mainml-explore/mlx:mainfrom

wyanzhao:nax-128-qk-unrollwyanzhao/mlx:nax-128-qk-unrollCopy head branch name to clipboard