Add format

.gitignore

@@ -14,8 +14,5 @@ __pycache__/
 .ipynb_checkpoints/
 .idea
 *.egg-info/
-<<<<<<< HEAD
 torch_compile_debug/
-=======
 build/
->>>>>>> 1a642f1 (Add flash_attention_benchmark)

benchmarks/flash_attention_bench/run.py

@@ -12,13 +12,31 @@
 
 
 def run():
-    args = ["--batch", "4", "--seq-len", "16384", "--n-heads", "32", "--d-head", "64", "--precision", "fp16", "--bwd", "--only", "triton_tutorial_flash_v2", "--causal", "--metrics", "tflops"]
+    args = [
+        "--batch",
+        "4",
+        "--seq-len",
+        "16384",
+        "--n-heads",
+        "32",
+        "--d-head",
+        "64",
+        "--precision",
+        "fp16",
+        "--bwd",
+        "--only",
+        "triton_tutorial_flash_v2",
+        "--causal",
+        "--metrics",
+        "tflops",
+    ]
     flash_attn_op = tritonbench.load_opbench_by_name("flash_attention")
     parser = get_parser()
     args, extra_args = parser.parse_known_args(args)
     flash_attn_bench = flash_attn_op(args, extra_args)
     flash_attn_bench.run()
     print(flash_attn_bench.output)
 
+
 if __name__ == "__main__":
     run()

benchmarks/gemm_bench/run.py

@@ -13,13 +13,27 @@
 
 
 def run():
-    args = ["--m", "4096", "--n", "4096", "--k", "4096", "--precision", "fp16", "--only", "triton_tutorial_matmul", "--metrics", "tflops"]
+    args = [
+        "--m",
+        "4096",
+        "--n",
+        "4096",
+        "--k",
+        "4096",
+        "--precision",
+        "fp16",
+        "--only",
+        "triton_tutorial_matmul",
+        "--metrics",
+        "tflops",
+    ]
     gemm_op = tritonbench.load_opbench_by_name("gemm")
     parser = get_parser()
     args, extra_args = parser.parse_known_args(args)
     gemm_bench = gemm_op(args, extra_args)
     gemm_bench.run()
     print(gemm_bench.output)
 
+
 if __name__ == "__main__":
     run()

tritonbench/operators/flash_attention/operator.py

@@ -150,7 +150,9 @@ def parse_op_args(args: List[str]):
         action="store_true",
         help="enable causal (always true on backward)",
     )
-    parser.add_argument("--additional-inputs", action="store_true", help="enable additional inputs")
+    parser.add_argument(
+        "--additional-inputs", action="store_true", help="enable additional inputs"
+    )
     return parser.parse_args(args)
 
 
@@ -481,6 +483,7 @@ def get_bwd_fn(self, fwd_fn: Callable) -> Callable:
 
     def get_input_iter(self) -> Generator:
         import math
+
         D_HEAD = self.D_HEAD
         BATCH = self.BATCH
         H = self.H