format

Author: yyttt6

examples/flash_attention/test_example_flash_attention.py

@@ -34,32 +34,35 @@ def test_example_gqa_bwd_wgmma_pipelined():
 @tilelang.testing.requires_cuda
 def test_example_mha_bwd():
     example_mha_bwd.main(
-    BATCH = 1,
-    H = 16,
-    N_CTX = 512,
-    D_HEAD = 64,
-    causal = False,)
+        BATCH=1,
+        H=16,
+        N_CTX=512,
+        D_HEAD=64,
+        causal=False,
+    )
 
 
 @tilelang.testing.requires_cuda
 def test_example_mha_bwd_bhsd():
     example_mha_bwd_bhsd.main(
-    BATCH = 1,
-    H = 16,
-    N_CTX = 512,
-    D_HEAD = 64,
-    causal = False,)
+        BATCH=1,
+        H=16,
+        N_CTX=512,
+        D_HEAD=64,
+        causal=False,
+    )
 
 
 @tilelang.testing.requires_cuda
 @tilelang.testing.requires_cuda_compute_version_ge(9, 0)
 def test_example_mha_bwd_wgmma_pipelined():
     example_mha_bwd_wgmma_pipelined.main(
-    BATCH = 1,
-    H = 16,
-    N_CTX = 512,
-    D_HEAD = 64,
-    causal = False,)
+        BATCH=1,
+        H=16,
+        N_CTX=512,
+        D_HEAD=64,
+        causal=False,
+    )
 
 
 @tilelang.testing.requires_cuda
@@ -99,7 +102,7 @@ def test_example_mha_fwd_bshd():
 
 @tilelang.testing.requires_cuda
 def test_example_mha_fwd_varlen():
-    example_mha_fwd_varlen.main(batch = 4, heads = 16, seq_len = 512, dim = 64)
+    example_mha_fwd_varlen.main(batch=4, heads=16, seq_len=512, dim=64)
 
 
 if __name__ == "__main__":

examples/flash_decoding/example_mha_inference.py

@@ -302,7 +302,7 @@ def flash_split_ref(Q, K, V, causal):
                                              3), gacc_o.to(torch.float16).permute(1, 2, 3, 0, 4)
 
 
-def main(BATCH=1,  H=32, Q_CTX=128, KV_CTX=8192, D_HEAD=128, causal=False):
+def main(BATCH=1, H=32, Q_CTX=128, KV_CTX=8192, D_HEAD=128, causal=False):
     flops_per_matmul = 2.0 * BATCH * H * Q_CTX * KV_CTX * D_HEAD
     total_flops = 2 * flops_per_matmul
     if causal:

examples/flash_decoding/test_example_flash_decoding.py

@@ -12,7 +12,7 @@ def test_example_example_gqa_decode():
 
 
 def test_example_example_mha_inference():
-    example_mha_inference.main(BATCH=1,  H=32, Q_CTX=128, KV_CTX=2048, D_HEAD=128, causal=False)
+    example_mha_inference.main(BATCH=1, H=32, Q_CTX=128, KV_CTX=2048, D_HEAD=128, causal=False)
 
 
 if __name__ == "__main__":