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

sdpa specialization for head dim 256 (#2007)

Browse Source
This commit is contained in:
Awni Hannun
2025-03-27 19:31:25 -07:00
committed by GitHub
parent a6b5d6e759
commit bc62932984
2 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
mlx/backend/metal/kernels/scaled_dot_product_attention.metal
View File

@@ -32,9 +32,11 @@ using namespace metal;
instantiate_sdpa_vector(type, 64, 64) \
instantiate_sdpa_vector(type, 96, 96) \
instantiate_sdpa_vector(type, 128, 128) \
instantiate_sdpa_vector(type, 256, 256) \
instantiate_sdpa_vector_aggregation(type, 64) \
instantiate_sdpa_vector_aggregation(type, 96) \
instantiate_sdpa_vector_aggregation(type, 128)
instantiate_sdpa_vector_aggregation(type, 128) \
instantiate_sdpa_vector_aggregation(type, 256)
instantiate_sdpa_vector_heads(float)
instantiate_sdpa_vector_heads(bfloat16_t)