minor benchmark script

eth-easl · Sep 12, 2024 · e584b9c · e584b9c
1 parent 8d08401
commit e584b9c
Showing 1 changed file with 13 additions and 9 deletions.
diff --git a/benchmarks/bench_sbmm.py b/benchmarks/bench_sbmm.py
@@ -43,7 +43,9 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
         fp16_func,
         flops_func,
         kwargs={
+            "dist": distribution,
             "nr": nr,
+            "nm": nm,
             "m": m,
             "k": k,
             "x": x,
@@ -56,7 +58,9 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
         w4_2_4_forloop_func,
         flops_func,
         kwargs={
+            "dist": distribution,
             "nr": nr,
+            "nm": nm,
             "m": m,
             "k": k,
             "qweight": qweight,
@@ -71,7 +75,9 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
         w4_2_4_native_func,
         flops_func,
         kwargs={
+            "dist": distribution,
             "nr": nr,
+            "nm": nm,
             "m": m,
             "k": k,
             "qweight": qweight,
@@ -86,6 +92,7 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
         w4_2_4_multilaunch_func,
         flops_func,
         kwargs={
+            "dist": distribution,
             "nr": nr,
             "m": m,
             "k": k,
@@ -109,16 +116,13 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
 
 if __name__ == "__main__":
     results = []
-    nr = [16, 32, 64, 128]
+    nr = [100]
     nm = [
-        [2, 4, 8, 16],
-        [2, 4, 8, 16, 32],
-        [2, 4, 8, 16, 32, 64],
-        [2, 4, 8, 16, 32, 64, 128],
+        [2, 4, 8, 16, 32, 64, 100],
     ]
-    distributions = ["uniform", "zipf:1.5"]
-    ms = [4096, 8192]
-    ns = [4096, 8192]
+    distributions = ["zipf:2.0"]
+    ms = [2048]
+    ns = [2048]
     for distribution in distributions:
         for i in range(len(nr)):
             for j in range(len(nm[i])):
@@ -133,4 +137,4 @@ def w4_2_4_multilaunch_func(qweight, x, meta, scale, indices):
                                 f"Failed to benchmark sbmm nr={nr[i]},nm={nm[i][j]},m={m},n={n}"
                             )
                             print(e)
-    export_benchmark_results(results, ".local/sbmm_bench.json")
+    export_benchmark_results(results, ".local/sbmm_bench.json")