https://github.com/meta-pytorch/tritonbench/blob/494153381172daf96d703e455364afb7739a6b81/tritonbench/kernels/proton_fused_attention.py