zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-05 00:31:15 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Update sdpa_benchmarks

Browse Source
This commit is contained in:
Jagrit Digani 2024-11-20 15:19:42 -08:00
parent f1d87a2d3e
commit c9ab537b9a
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
benchmarks/python/sdpa_bench.py
View File

@ -145,6 +145,8 @@ if __name__ == "__main__":
"float32",
)
transposes = (False,)
# clang-format off
shapes = (
# ( B, qsl, ksl, head_dim, n_qh, n_kvh)
(1, 32, 32, 64, 32, 32),
@ -156,6 +158,7 @@ if __name__ == "__main__":
(1, 2048, 2048, 64, 32, 32),
(1, 4096, 4096, 64, 32, 32),
)
# clang-format on
print(" B, qsl, ksl, hdim, n_qh, n_kvh, tpose, dtype, t_unfs, t_fuse, diff%")