[FA] Unifying TMA Kernels with Warp Specialization Flag

[codingwithsurya]

Summary:
This PR consolidates redundant TMA attention kernels into a unified implementation. Previously, _attn_fwd_tma and _attn_fwd_tma_ws contained duplicate code (mainly the TMA descriptors) and didn't leverage the existing ENABLE_WS flag.

I've merged the redundant kernels into a single _attn_fwd_tma_unified kernel. We now use the ENABLE_WS flag to toggle between regular and warp-specialized execution.

Changes:

• Merged both kernels into _attn_fwd_tma_unified kernel with handling of regular and warp-specialized paths
• Utilized existing ENABLE_WS parameter to control warp specialization
• Unified TMA descriptor creation logic

Test Plan:
Unit Tests and Benchmarking

WITH_TMA=1  python run.py --op flash_attention --only triton_tutorial_flash_v2_tma_ws --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128
WITH_TMA=1  python run.py --op flash_attention --only triton_tutorial_flash_v2_tma --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128
python -m unittest test/test_gpu/main.py -k test_gpu_tritonbench_flash_attention

The performance metrics before and after the code change are identical (4.39805e+12 FLOPS).

Follow Up PR for Base + Opt Kernels: #233

Differential Revision: D75307125 and D75308966 (follow-up diff)

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D75308966

[facebook-github-bot]

This pull request was exported from Phabricator. Differential Revision: D75308966

[manman-ren]

Thanks for working on this! The patch looks good. We have been focusing on non-causal, so maybe testing

WITH_TMA=1 CUDA_VISIBLE_DEVICES=5 python run.py --op flash_attention --only triton_tutorial_flash_v2_tma_ws --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128
WITH_TMA=1 CUDA_VISIBLE_DEVICES=5 python run.py --op flash_attention --only triton_tutorial_flash_v2_tma --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128

We need WITH_TMA=1 to actually enable tma.
If you haven't imported the diff to fbsource yet, please do it, it will trigger tests with our internal Triton.

[codingwithsurya]

Thanks for working on this! The patch looks good. We have been focusing on non-causal, so maybe testing

WITH_TMA=1 CUDA_VISIBLE_DEVICES=5 python run.py --op flash_attention --only triton_tutorial_flash_v2_tma_ws --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128
WITH_TMA=1 CUDA_VISIBLE_DEVICES=5 python run.py --op flash_attention --only triton_tutorial_flash_v2_tma --num-inputs 1 --seq-len 8192 --metrics tflops --batch 8 --n-heads 16 --d-head 128

We need WITH_TMA=1 to actually enable tma. If you haven't imported the diff to fbsource yet, please do it, it will trigger tests with our internal Triton.

Thanks for letting me know! I have tested it with WITH_TMA=1 and it works. I have updated the test plan.

I exported it from fbsource. For reference, the diff in fbsource is here (this specific PR is the first two diffs in the stack).

[mandroid6]

LGTM!