Skip to content

Commit

Permalink
[TL] test flashattention script, remove commented code
Browse files Browse the repository at this point in the history
  • Loading branch information
tzj-fxz committed Sep 27, 2024
1 parent 6955697 commit 1b564a8
Showing 1 changed file with 0 additions and 23 deletions.
23 changes: 0 additions & 23 deletions testing/python/tilelang/test_tilelang_flash_atten.py
Original file line number Diff line number Diff line change
Expand Up @@ -257,26 +257,3 @@ def test_flashattn():

if __name__ == "__main__":
bitblas.testing.main()

# if __name__ == "__main__":
# parser = argparse.ArgumentParser()
# parser.add_argument("--batch", type=int, default=4, help="Batch size")
# parser.add_argument("--h", type=int, default=4, help="Number of heads")
# parser.add_argument("--n_ctx", type=int, default=256, help="Context size")
# parser.add_argument("--d_head", type=int, default=256, help="Head dimension")
# parser.add_argument("--casual", type=bool, default=True, help="Casual flag")
# args = parser.parse_args()
# BATCH, H, N_CTX, D_HEAD = args.batch, args.h, args.n_ctx, args.d_head
# casual = args.casual
# flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * D_HEAD
# total_flops = 2 * flops_per_matmul
# if casual:
# total_flops *= 0.5

# # best_latency, ref_latency = flashattn(BATCH, H, N_CTX, D_HEAD, casual)
# # best_config = None
# best_latency, best_config, ref_latency = flashattn_autotune(BATCH, H, N_CTX, D_HEAD, casual)
# print(f"Best latency: {best_latency}")
# print(f"Best TFlops: {total_flops / best_latency * 1e-9}")
# print(f"Best config: {best_config}")
# print(f"Ref TFlops: {total_flops / ref_latency * 1e-9}")

0 comments on commit 1b564a8

Please sign in to comment.