Fast Multi-ahead Attention support on AMD ROCM

[qianfengz]

This PR adds three flash-attention implementation for AMD ROCM

1. Generic FMHA forward based on composable_kernel kernel components
2. decoder FMHA forward directly implemented in HIP kernel
3. triton FMHA forward operation based on triton

In more details, the following codes are added in this PR

1. Xformers Operator and its C++ implementation for Generic FMHA forward as well as the underlying composable_kernel_tiled submodule

   xformers/ops/fmha/ck.py
   xformers/csrc/attention/hip_fmha/
   thirty_party/composable_kernel_tiled/

2. Xformers Operator and its C++ implementation for decoder FMHA forward

   xformers/ops/fmha/ck_decoder.py, ck_splitk.py
   xformers/csrc/attention/hip_fmha/

3. Xformers Operator for triton FMHA forward

   xformers/ops/fmha/triton.py

The following scripts are used to verify the implementation

#> pytest tests/test_mem_eff_attention.py::test_forward
#> pytest tests/test_mem_eff_attention.py::test_mqa_forwrd
#> pytest tests/test_mem_eff_attention.py::test_decoder
#> pytest tests/test_mem_eff_attention.py::test_splitk_decoder
#> pytest  tests/test_mem_eff_attention.py::test_splitk_reference
#> pytest tests/test_mem_eff_attention.py::test_triton_splitk_decoder

The following scripts are used to benchmark the performance of the implementation

#> python xformers/benchmarks/benchmark_mem_eff_attention.py
#> python xformers/benchmarks/benchmark_mem_eff_attention_mqa.py
#> python xformers/benchmarks/benchmark_mem_eff_attn_decoder.py
#> python xformers/benchmarks/benchmark_attn_decoding.py

[xw285cornell]

rename this file?

[xw285cornell]

IIUC, this file is to implement the flash attention here rather than importing from third_party/flash-attention/flash_attn/flash_attn_triton.py, is that right? Is there any amd specific optimization here (btw it's great to import more masks than just the lowerTriangular) or it's similar?

[qianfengz]

Yes, this is from a patch PRed to http://github.com/ROCmSoftwarePlatform/xformers by [email protected].

But we did not keep the PR records because we just re-build the repo.

I will re-submit PR from http://github.com/ROCmSoftwarePlatform/xformers branch rather than from my personal repo.

Thanks

[HinaHyugaHime]

was this scrapped or something?