Flash attention for rocm

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "csrc/flash_attn/cutlass"]
 	path = csrc/flash_attn/cutlass
 	url = https://github.com/NVIDIA/cutlass.git
+[submodule "csrc/flash_attn_rocm/composable_kernel"]
+	path = csrc/flash_attn_rocm/composable_kernel
+	url = https://github.com/ROCmSoftwarePlatform/composable_kernel

Dockerfile.rocm

@@ -0,0 +1,18 @@
+# BSD 3 Clause
+# Copyright 2023 Advanced Micro Devices, Inc.
+# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+FROM rocm/pytorch:rocm5.3.2_ubuntu20.04_py3.7_pytorch_1.12.1
+
+WORKDIR /workspace
+USER root
+
+RUN pip install ninja
+COPY ./ /workspace/flash-attention_private/
+RUN cd /workspace/flash-attention_private \
+       && patch /opt/conda/lib/python3.7/site-packages/torch/utils/hipify/hipify_python.py hipify_patch.patch \
+       && python setup.py install

MANIFEST.in

@@ -7,3 +7,8 @@ recursive-include flash_attn *.cu
 recursive-include flash_attn *.h
 recursive-include flash_attn *.cuh
 recursive-include flash_attn *.cpp
+
+recursive-include flash_attn_rocm *.cu
+recursive-include flash_attn_rocm *.h
+recursive-include flash_attn_rocm *.cuh
+recursive-include flash_attn_rocm *.cpp

csrc/flash_attn_rocm/CMakeLists.txt

@@ -0,0 +1,154 @@
+# BSD 3 Clause
+# Copyright 2023 Advanced Micro Devices, Inc.
+# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
+project(fmha_api)
+
+IF(NOT DEFINED ENV{ROCM_PATH})
+  SET(ROCM_PATH /opt/rocm)
+ELSE()
+  SET(ROCM_PATH $ENV{ROCM_PATH})
+ENDIF()
+if(NOT DEFINED ENV{ROCM_INCLUDE_DIRS})
+  set(ROCM_INCLUDE_DIRS ${ROCM_PATH}/include)
+else()
+  set(ROCM_INCLUDE_DIRS $ENV{ROCM_INCLUDE_DIRS})
+endif()
+# HIP_PATH
+IF(NOT DEFINED ENV{HIP_PATH})
+  SET(HIP_PATH ${ROCM_PATH}/hip)
+ELSE()
+  SET(HIP_PATH $ENV{HIP_PATH})
+ENDIF()
+
+
+
+IF(NOT EXISTS ${HIP_PATH})
+  return()
+ENDIF()
+
+
+
+# HCC_PATH
+IF(NOT DEFINED ENV{HCC_PATH})
+  SET(HCC_PATH ${ROCM_PATH}/hcc)
+ELSE()
+  SET(HCC_PATH $ENV{HCC_PATH})
+ENDIF()
+
+
+
+# HSA_PATH
+IF(NOT DEFINED ENV{HSA_PATH})
+  SET(HSA_PATH ${ROCM_PATH}/hsa)
+ELSE()
+  SET(HSA_PATH $ENV{HSA_PATH})
+ENDIF()
+
+
+
+# ROCBLAS_PATH
+IF(NOT DEFINED ENV{ROCBLAS_PATH})
+  SET(ROCBLAS_PATH ${ROCM_PATH}/rocblas)
+ELSE()
+  SET(ROCBLAS_PATH $ENV{ROCBLAS_PATH})
+ENDIF()
+
+
+
+# ROCSPARSE_PATH
+IF(NOT DEFINED ENV{ROCSPARSE_PATH})
+  SET(ROCSPARSE_PATH ${ROCM_PATH}/rocsparse)
+ELSE()
+  SET(ROCSPARSE_PATH $ENV{ROCSPARSE_PATH})
+ENDIF()
+
+
+
+# ROCFFT_PATH
+IF(NOT DEFINED ENV{ROCFFT_PATH})
+  SET(ROCFFT_PATH ${ROCM_PATH}/rocfft)
+ELSE()
+  SET(ROCFFT_PATH $ENV{ROCFFT_PATH})
+ENDIF()
+
+
+
+# HIPSPARSE_PATH
+IF(NOT DEFINED ENV{HIPSPARSE_PATH})
+  SET(HIPSPARSE_PATH ${ROCM_PATH}/hipsparse)
+ELSE()
+  SET(HIPSPARSE_PATH $ENV{HIPSPARSE_PATH})
+ENDIF()
+
+
+
+# THRUST_PATH
+IF(NOT DEFINED ENV{THRUST_PATH})
+  SET(THRUST_PATH ${ROCM_PATH}/include)
+ELSE()
+  SET(THRUST_PATH $ENV{THRUST_PATH})
+ENDIF()
+
+
+
+# HIPRAND_PATH
+IF(NOT DEFINED ENV{HIPRAND_PATH})
+  SET(HIPRAND_PATH ${ROCM_PATH}/hiprand)
+ELSE()
+  SET(HIPRAND_PATH $ENV{HIPRAND_PATH})
+ENDIF()
+
+
+
+# ROCRAND_PATH
+IF(NOT DEFINED ENV{ROCRAND_PATH})
+  SET(ROCRAND_PATH ${ROCM_PATH}/rocrand)
+ELSE()
+  SET(ROCRAND_PATH $ENV{ROCRAND_PATH})
+ENDIF()
+
+
+
+# MIOPEN_PATH
+IF(NOT DEFINED ENV{MIOPEN_PATH})
+  SET(MIOPEN_PATH ${ROCM_PATH}/miopen)
+ELSE()
+  SET(MIOPEN_PATH $ENV{MIOPEN_PATH})
+ENDIF()
+
+
+
+# Add HIP to the CMAKE Module Path
+set(CMAKE_MODULE_PATH ${HIP_PATH}/cmake ${CMAKE_MODULE_PATH})
+
+find_package(HIP)
+
+set(CMAKE_CXX_COMPILER /opt/rocm/hip/bin/hipcc)
+set(CMAKE_CXX_STANDARD 20)
+
+list(APPEND CMAKE_PREFIX_PATH "/opt/conda/lib/python3.8/site-packages/torch/share/cmake")
+find_package(Torch REQUIRED)
+
+find_package(rocblas)
+find_package(hipfft)
+find_package(hiprand)
+find_package(hipsparse)
+
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/src)
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/composable_kernel/include)
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/composable_kernel/library/include)
+include_directories(/opt/conda/include/python3.8)
+
+aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/src FLA_SRCS)
+aux_source_directory(${CMAKE_CURRENT_SOURCE_DIR}/composable_kernel/library/src/utility CK_SRCS)
+
+add_executable(fmha_api fmha_api.cpp ${FLA_SRCS} ${CK_SRCS})
+target_link_libraries(fmha_api "${TORCH_LIBRARIES}")
+
+message("${TORCH_LIBRARIES}")

csrc/flash_attn_rocm/Dockerfile

@@ -0,0 +1,31 @@
+# BSD 3 Clause
+# Copyright 2023 Advanced Micro Devices, Inc.
+# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+FROM rocm/pytorch:rocm5.3.2_ubuntu20.04_py3.7_pytorch_1.12.1
+WORKDIR /flash_attn
+
+USER root
+
+ENV DEBIAN_FRONTEND noninteractive
+ENV TZ "Asia/Shanghai"
+
+RUN apt-get update \
+    && apt install -y git-all \
+    && git clone https://<your github username>:<your github token>@github.com/ROCmSoftwarePlatform/flash-attention_private \
+    && cd /flash_attn/flash-attention_private \
+    && git checkout flash_attention_for_rocm \
+    && cd /flash_attn/flash-attention_private/csrc/flash_attn_rocm/composable_kernel \
+    && git submodule init \
+    && git submodule update \
+    && cd /flash_attn/flash-attention_private/csrc/flash_attn_rocm \
+    && mkdir build \
+    && cd build \
+    && cmake .. \
+    && cd /flash_attn/flash-attention_private/csrc/flash_attn_rocm/build \
+    && make -j64
+

csrc/flash_attn_rocm/README.md

@@ -0,0 +1,49 @@
+Here is the folder for APIs on rocm, which the backend code is from composable kernel.
+
+Below is the introduction to the files.
+
+"src/fmha.h" is the header file for the C++ APIs, in which declared the  function "run_fmha_fp16_bf16_gfx90a".
+
+"fmha_api.cpp" is the c++ file that defined the API function "mha_fwd", this function will call function "run_fmha_fp16_bf16_gfx90a". This function also contains a main function to test with the API.
+
+"src/fmha_fprop_fp16_bf16_kernel.gfx90a" is the interface that link API in fmha_api.cpp and the CK backend, which defined function "run_fmha_fp16_bf16_gfx90a". In this function, it will use parameters conveyed from "mha_fwd" to choose proper instance parameters for CK function. Function "run_fmha_fp16_bf16_gfx90a_loop_" will use parameters from "run_fmha_fp16_bf16_gfx90a" to initialize instance in CK and call CK function. 
+
+"CMakeList.txt" is a cmake file to compile the example above.
+
+Useage for "CMakeLists.txt": 
+```
+$mkdir build
+$cd build
+$cmake ..
+$make
+```
+
+My docker is from https://hub.docker.com/layers/rocm/pytorch/rocm5.3.2_ubuntu20.04_py3.7_pytorch_1.12.1/images/sha256-387b2538d14cfd55a9510b7ea07049f1e71b7e755413080153b997c798fe5099?context=explore
+
+If you choose another docker or you install pytorch by yourself.
+
+Please change line 8 in CMakeLists.txt file with your own path.
+
+You can use command
+``` 
+python -c 'import torch;print(torch.utils.cmake_prefix_path)'
+```
+to find your path.
+
+Way to build with docker file:
+
+Change the github username and tocken with that of yourself in line https://github.com/ROCmSoftwarePlatform/flash-attention_private/blob/41ddb2fb3884085ee5318d30f8e919944ee18745/csrc/flash_attn_rocm/Dockerfile#L11 firstly.
+
+Then
+```
+sudo docker build -t flash_attention:rocm5.3.2 .
+```
+
+If you want to test the performance, you can set the parameter “time_kernel” as true. And then the kernel will run 10 times and give out the average running time. You can find the parameter in this line: https://github.com/ROCmSoftwarePlatform/flash-attention_private/blob/fb47a607682a873a3e0b17ae220849cc11a34d8b/csrc/flash_attn_rocm/src/fmha_fprop_fp16_bf16_kernel.gfx90a.cpp#L142
+
+If you want to verify the results, you can set the parameter “do_verification” in this line https://github.com/ROCmSoftwarePlatform/flash-attention_private/blob/fb47a607682a873a3e0b17ae220849cc11a34d8b/csrc/flash_attn_rocm/fmha_api.cpp#L271 . And then the code can do the same computation on cpu and compare with the results from device and show whether device results are right.
+
+
+
+
+