L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix int16 overflow in SDPA NAX mask indexing for KV sequences > 32K (#3361)

Browse Source 
Co-authored-by: Angelos Katharopoulos <[email protected]>
This commit is contained in:
Clydingus
2026-04-10 00:01:47 -07:00
committed by GitHub
co-authored by Angelos Katharopoulos
parent d6d9b24801
commit a33b791615
3 changed files with 102 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
benchmarks/python/sdpa_bench.py
+6
View File
@@ -176,6 +176,8 @@ if __name__ == "__main__":
( 1, 1024, 1024, 64, 32, 8),
( 1, 2048, 2048, 64, 32, 8),
( 1, 4096, 4096, 64, 32, 8),
( 1, 4096, 5000, 64, 32, 8),
( 1, 2048, 32121, 64, 32, 8),
)
shapes_80 = (
@@ -183,6 +185,8 @@ if __name__ == "__main__":
( 1, 1024, 1024, 80, 32, 8),
( 1, 2048, 2048, 80, 32, 8),
( 1, 4096, 4096, 80, 32, 8),
( 1, 4096, 5000, 80, 32, 8),
( 1, 2048, 32121, 80, 32, 8),
)
shapes_128 = (
@@ -190,6 +194,8 @@ if __name__ == "__main__":
( 1, 1024, 1024, 128, 32, 8),
( 1, 2048, 2048, 128, 32, 8),
( 1, 4096, 4096, 128, 32, 8),
( 1, 4096, 5000, 128, 32, 8),
( 1, 2048, 32121, 128, 32, 8),
)
# fmt: on