()
| 831 | |
| 832 | |
| 833 | def main(): |
| 834 | parser = argparse.ArgumentParser() |
| 835 | parser.add_argument("-K", type=int, required=False, default=512) |
| 836 | parser.add_argument("--K_range", type=int, nargs=2) |
| 837 | parser.add_argument("--K_step", type=int, default=512) |
| 838 | parser.add_argument("--prec", type=str, choices=["fp8", "fp16"], default="fp16") |
| 839 | args = parser.parse_args() |
| 840 | |
| 841 | if args.prec == "fp8" and (not hasattr(torch, "float8_e4m3fn") or not is_cuda()): |
| 842 | print("This example requires CUDA/HIP with fp8 support.") |
| 843 | else: |
| 844 | if args.K and args.K_range is None: |
| 845 | args.K_range = [args.K, args.K] |
| 846 | args.K_step = 1 # doesn't matter as long as it's not 0 |
| 847 | |
| 848 | torch.manual_seed(0) |
| 849 | |
| 850 | kwargs = {"M": 256, "N": 256_000, "K": 8_192, "dtype": torch.float16} |
| 851 | |
| 852 | validate(**kwargs) |
| 853 | |
| 854 | proton.start("matmul", hook="triton") |
| 855 | proton.deactivate() |
| 856 | bench(**kwargs) |
| 857 | proton.finalize() |
| 858 | show_profile(args.prec, "matmul") |
| 859 | |
| 860 | |
| 861 | if __name__ == "__main__": |
no test coverage detected