add rope benchmarks

Author: kvignesh1420

benchmark/benchmark_geglu.py

@@ -90,7 +90,7 @@ def full():
             full, quantiles=QUANTILES, grad_to_none=[x], rep=10
         )
 
-    return ms, max_ms, min_ms
+    return ms, min_ms, max_ms
 
 
 def benchmark_speed_geglu_wrapper():

benchmark/benchmark_rope.py

@@ -109,7 +109,7 @@ def full():
         ms, min_ms, max_ms = triton.testing.do_bench(
             full, quantiles=QUANTILES, grad_to_none=[q, k], rep=400
         )
-    return ms, max_ms, min_ms
+    return ms, min_ms, max_ms
 
 
 def benchmark_speed_rope_wrapper():
@@ -158,8 +158,9 @@ def full():
             (q_out, k_out), (q, k), (dq, dk), allow_unused=True, retain_graph=True
         )
 
-    mem = _test_memory(full)
-    return mem / 2**20
+    mem, min_mem, max_mem = _test_memory(full, quantiles=QUANTILES)
+    return (mem / 2**20, min_mem / 2**20, max_mem / 2**20)
+
 
 
 def benchmark_memory_rope_wrapper():

benchmark/rope_memory/rope-full-memory-benchmark-seq-2048.png

@@ -1,4 +1,4 @@
 total_hidden_size,Liger,Hugging Face
-512.000000,0.046912,0.153600
-2048.000000,0.021344,0.189792
-8192.000000,0.058720,0.619296
+512.000000,0.048880,0.173376
+2048.000000,0.050368,0.190304
+8192.000000,0.058784,0.621280

benchmark/rope_memory/rope-full-memory-benchmark-total_dim_8192.png

@@ -1,6 +1,6 @@
 seq_len,Liger,Hugging Face
-1024.000000,0.033824,0.328864
-2048.000000,0.059200,0.619296
-4096.000000,0.109568,1.186832
-8192.000000,0.209344,2.317760
-16384.000000,0.410304,4.547232
+1024.000000,0.033760,0.329920
+2048.000000,0.058784,0.621312
+4096.000000,0.109664,1.189088
+8192.000000,0.210016,2.320608
+16384.000000,0.409984,4.549056

benchmark/rope_speed/rope-backward-speed-benchmark-seq-2048.csv

@@ -1,4 +1,4 @@
 total_hidden_size,Liger,Hugging Face
-512.000000,0.010272,0.078816
-2048.000000,0.021184,0.156160
-8192.000000,0.059232,0.512704
+512.000000,0.011360,0.078816
+2048.000000,0.020576,0.157344
+8192.000000,0.058720,0.516448

benchmark/rope_speed/rope-backward-speed-benchmark-seq-2048.png

@@ -1,6 +1,6 @@
 seq_len,Liger,Hugging Face
-1024.000000,0.033536,0.277696
-2048.000000,0.058528,0.512928
-4096.000000,0.109504,0.986304
-8192.000000,0.209920,1.920208
-16384.000000,0.409840,3.790624
+1024.000000,0.034240,0.280288
+2048.000000,0.059296,0.516512
+4096.000000,0.109056,0.993184
+8192.000000,0.209376,1.933456
+16384.000000,0.410464,3.816672

benchmark/rope_speed/rope-backward-speed-benchmark-total_dim_8192.csv

@@ -1,4 +1,4 @@
 total_hidden_size,Liger,Hugging Face
-512.000000,0.120384,0.350336
-2048.000000,0.122208,0.341728
-8192.000000,0.115936,1.128096
+512.000000,0.159264,0.429248
+2048.000000,0.158368,0.342784
+8192.000000,0.114720,1.134144

benchmark/rope_speed/rope-backward-speed-benchmark-total_dim_8192.png

@@ -1,6 +1,6 @@
 seq_len,Liger,Hugging Face
-1024.000000,0.120064,0.601600
-2048.000000,0.114944,1.127872
-4096.000000,0.214656,2.167280
-8192.000000,0.415552,4.232288
-16384.000000,0.816928,8.331616
+1024.000000,0.109088,0.605088
+2048.000000,0.114336,1.134048
+4096.000000,0.214880,2.176800
+8192.000000,0.415680,4.246112
+16384.000000,0.816704,8.356367