pre-commit

Author: Edenzzzz

benchmarks/bench_blackwell_attention.py

@@ -75,9 +75,10 @@ def flops(ms):
         else:
             return batch_size * qkv_len * qkv_len * num_heads * head_dim * 4 / ms / 1e9
 
-    print(
-        f"bench_fmha_blackwell (batch_size={batch_size}, qkv_len={qkv_len}, num_heads={num_heads}, head_dim={head_dim}, causal={causal}), flops: {flops(ms):.3f} TFLOPs/s"
-    )
+    msg = f"bench_fmha_blackwell (batch_size={batch_size}, qkv_len={qkv_len}, num_heads={num_heads}, head_dim={head_dim}, causal={causal}), flops: {flops(ms):.3f} TFLOPs/s"
+    print(msg)
+    with open("bench_fmha_blackwell.txt", "a") as f:
+        f.write(msg + "\n")
 
 
 if __name__ == "__main__":

benchmarks/bench_deepgemm_blackwell.py

@@ -67,11 +67,15 @@ def bench_deepgemm_grouped_fp8_blackwell(batch_size, m, n, k, in_dtype, out_dtyp
         * 1e-9
         / ms
     )
-    print(
-        f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} "
-        f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s"
-        f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
-    )
+    # print(
+    #     f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} "
+    #     f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s"
+    #     f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
+    # )
+    msg = f"deepgemm_grouped_fp8_blackwell batch_size={batch_size} m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
+    print(msg)
+    with open("bench_deepgemm_fp8_blackwell.txt", "a") as f:
+        f.write(msg + "\n")
 
     return tflops_per_second
 
@@ -118,11 +122,15 @@ def bench_deepgemm_batch_fp8_blackwell(batch_size, m, n, k, in_dtype, out_dtype)
         * 1e-9
         / ms
     )
-    print(
-        f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} "
-        f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s"
-        f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
-    )
+    msg = f"bench_deepgemm_fp8_blackwell batch_size={batch_size} m={m} n={n} k={k} in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
+    print(msg)
+    with open("bench_deepgemm_fp8_blackwell.txt", "a") as f:
+        f.write(msg + "\n")
+    # print(
+    #     f"group_deepgemm_fp8_nt_groupwise batch_size={batch_size} m={m} n={n} k={k} "
+    #     f"in_dtype={in_dtype} out_dtype={out_dtype}: {tflops_per_second:.2f} TFLOPs/s"
+    #     f"memory_bandwidth: {memory_bandwidth_per_second:.2f} TB/s"
+    # )
 
     return tflops_per_second
 

benchmarks/bench_groupwise_gemm_fp8_blackwell.py

@@ -170,7 +170,9 @@ def bench_groupwise_gemm_fp8_blackwell(m, n, k, in_dtype, out_dtype):
     gemm_fp8_nt_groupwise(a, b, a_scale, b_scale, out=out, scale_major_mode="MN")
 
     measurements = bench_gpu_time(
-        lambda: gemm_fp8_nt_groupwise(a, b, a_scale, b_scale, out=out, scale_major_mode="MN")
+        lambda: gemm_fp8_nt_groupwise(
+            a, b, a_scale, b_scale, out=out, scale_major_mode="MN"
+        )
     )
     ms = np.median(measurements)
     tflops_per_second = 2 * m * n * k * 1e-9 / ms

benchmarks/bench_tgv_gemm.py

@@ -99,7 +99,9 @@ def test_tgv_gemm_bf16_sm100_perf():
         torch.cuda.synchronize()
         end_time = time.time()
         cublas_avg_time = (end_time - start_time) / 100
-        print(f"CUBLAS average time: {cublas_avg_time * 1000:.6f} ms, {flops/cublas_avg_time:.3f} TFLOPS")
+        print(
+            f"CUBLAS average time: {cublas_avg_time * 1000:.6f} ms, {flops / cublas_avg_time:.3f} TFLOPS"
+        )
 
         # Warmup
         with autotune(tune_mode=True):
@@ -128,7 +130,7 @@ def test_tgv_gemm_bf16_sm100_perf():
 
         tgv_avg_time = (end_time - start_time) / 100
         print(
-            f"TGV average time: {tgv_avg_time * 1000:.6f} ms, {flops/tgv_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / tgv_avg_time:.2f}x"
+            f"TGV average time: {tgv_avg_time * 1000:.6f} ms, {flops / tgv_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / tgv_avg_time:.2f}x"
         )
 
         # Test with PDL
@@ -151,7 +153,7 @@ def test_tgv_gemm_bf16_sm100_perf():
 
         pdl_avg_time = (end_time - start_time) / 100
         print(
-            f"PDL average time: {pdl_avg_time * 1000:.6f} ms, {flops/pdl_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / pdl_avg_time:.2f}x"
+            f"PDL average time: {pdl_avg_time * 1000:.6f} ms, {flops / pdl_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / pdl_avg_time:.2f}x"
         )
 
         # Store results for CSV