[FA] clean up and make TMA, scheduling autotunable

[manman-ren]

Variants:

• triton_tutorial_flash_v2_opt no TMA, with computation pipelining
• triton_tutorial_flash_v2_tma TMA, with computation pipelining
• triton_tutorial_flash_v2_tma_ws: TMA, with computation pipelining and Warp Spec
• triton_tutorial_flash_v2_ws: no TMA, with computation pipelining and Warp Spec

Test plan:

CUDA_VISIBLE_DEVICES=5 TORCH_CUDA_ARCH_LIST=9.0a python run.py --op flash_attention --only triton_tutorial_flash_v2_opt,triton_tutorial_flash_v2_tma,triton_tutorial_flash_v2 --num-inputs 1 --seq-len 13 --metrics tflops --batch 8 --n-heads 16 --d-head 128

CUDA_VISIBLE_DEVICES=5 TORCH_CUDA_ARCH_LIST=9.0a python run.py --op flash_attention --only triton_tutorial_flash_v2_opt,triton_tutorial_flash_v2_tma,triton_tutorial_flash_v2 --num-inputs 1 --seq-len 13 --metrics accuracy --batch 8 --n-heads 16 --d-head 128 --baseline triton_tutorial_flash_v2

On compiler supporting WarpSpec:
CUDA_VISIBLE_DEVICES=5 TORCH_CUDA_ARCH_LIST=9.0a python run.py --op flash_attention --only triton_tutorial_flash_v2_ws,triton_tutorial_flash_v2_tma_ws,triton_tutorial_flash_v2 --num-inputs 1 --seq-len 13 --metrics tflops --batch 8 --n-heads 16 --d-head 128
CUDA_VISIBLE_DEVICES=5 TORCH_CUDA_ARCH_LIST=9.0a python run.py --op flash_attention --only triton_tutorial_flash_v2_ws,triton_tutorial_flash_v2_tma_ws,triton_tutorial_flash_v2 --num-inputs 1 --seq-len 13 --metrics accuracy --batch 8 --n-heads 16 --d-head 128 --baseline triton_tutorial_flash_v2

[facebook-github-bot]

@manman-ren has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@manman-ren merged this pull request in c0a8479.