tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

benchmark script for matmuls in bert (#9752)

Browse Source 
2 main matmuls in the bert layers. getting these to be fast makes bert fast
This commit is contained in:
chenyu
2025-04-06 19:34:25 +08:00
committed by GitHub
parent 139999c6d7
commit 8a585dc5c1
1 changed files with 14 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/external/external_benchmark_bert_matmuls.py
Vendored
+14
View File
@@ -0,0 +1,14 @@
from tinygrad import Tensor
if __name__ == "__main__":
# matmuls in bert layers
BS=96
shapes = [
((BS//6, 16, 512, 512), (BS//6, 16, 512, 64)), # x48. 27 TFLOPS on 4090 with BEAM=4
((BS//6, 16, 512, 64), (BS//6, 16, 64, 512)), # x48. 47 TFLOPS on 4090 with BEAM=4
]
for s0, s1 in shapes:
t0 = Tensor.empty(s0, dtype="half")
t1 = Tensor.empty(s1, dtype="half")
for _ in range(5):
t0.matmul(t1, dtype="half").realize()