Merge branch 'cmake'.

Switches the build system to cmake, which is more flexible, and fixes up the compatibility for jax 0.5.0--0.7.0.

Author: nshepperd

.clangd

@@ -6,6 +6,9 @@ CompileFlags:
    - --use_fast_math
    - --threads
    - -gencode
+   - -forward-unknown-to-host-compiler
+   - --generate-code=*
+   - -Xcompiler=*
   Add:
     - --no-cuda-version-check
 ---

.github/workflows/publish.yml

@@ -41,8 +41,8 @@ jobs:
       fail-fast: false
       matrix:
           os: [ubuntu-20.04]
-          python-version: ['cp39', 'cp310', 'cp311', 'cp312']
-          cuda-version: ['11.8', '12.3']
+          python-version: ['cp311', 'cp312']
+          cuda-version: ['12.8']
 
     steps:
       - name: Checkout
@@ -51,7 +51,7 @@ jobs:
       - name: Set up python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
 
       - name: Set CUDA and PyTorch versions
         run: |
@@ -76,7 +76,7 @@ jobs:
         uses: pypa/[email protected]
         env:
           CIBW_BUILD: ${{ matrix.python-version }}-manylinux_x86_64
-          CIBW_MANYLINUX_X86_64_IMAGE: sameli/manylinux2014_x86_64_cuda_${{ matrix.cuda-version }}
+          CIBW_BEFORE_ALL: bash scripts/install-cuda-linux.sh ${{ matrix.cuda-version }}
           CIBW_BUILD_VERBOSITY: 1
 
       - name: Log Built Wheels
@@ -128,17 +128,15 @@ jobs:
 
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
 
       - name: Install dependencies
         run: |
-          pip install setuptools==68.0.0
-          pip install git+https://github.com/nshepperd/setuptools-cuda-cpp
-          pip install ninja packaging wheel pybind11
+          pip install uv
 
       - name: Build core package
         run: |
-          CUDA_HOME=/ python setup.py sdist --dist-dir=dist
+          uv build --sdist
 
       - name: Retrieve release distributions
         uses: actions/download-artifact@v4

CMakeLists.txt

@@ -0,0 +1,111 @@
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+cmake_minimum_required(VERSION 3.18)
+
+project(flash_attn LANGUAGES CXX CUDA)
+
+set(CMAKE_JOB_POOLS cuda=6)
+set(CMAKE_INSTALL_RPATH "$ORIGIN/nvidia/cuda_runtime/lib")
+# Make sure RPATH is used instead of RUNPATH
+set(CMAKE_INSTALL_RPATH_USE_LINK_PATH FALSE)
+
+# == Find dependencies ==
+find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
+
+execute_process(
+    COMMAND ${Python_EXECUTABLE} -m pybind11 --cmakedir
+    OUTPUT_VARIABLE pybind11_DIR
+    OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+
+find_package(pybind11 CONFIG REQUIRED)
+
+# == Setup CUDA ==
+string(REGEX REPLACE "--generate-code=arch=compute_[0-9]+,code=\\[?compute_[0-9]+,sm_[0-9]+\\]?" "" 
+       CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
+string(REGEX REPLACE "-gencode arch=compute_[0-9]+,code=sm_[0-9]+" "" 
+       CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
+
+message(WARNING "CMAKE_CUDA_FLAGS: ${CMAKE_CUDA_FLAGS}")
+
+# Set up ccache
+find_program(CCACHE_PROGRAM ccache)
+if(CCACHE_PROGRAM)
+    set(CMAKE_CUDA_COMPILER_LAUNCHER "${CCACHE_PROGRAM}")
+    message(STATUS "Using ccache: ${CCACHE_PROGRAM}")
+endif()
+
+# CUDA handling
+# Get CUDA architectures from environment or use default
+if(DEFINED ENV{FLASH_ATTN_CUDA_ARCHS})
+    set(CMAKE_CUDA_ARCHITECTURES $ENV{FLASH_ATTN_CUDA_ARCHS})
+else()
+    # set(CMAKE_CUDA_ARCHITECTURES "80;90;100;120")
+    set(CMAKE_CUDA_ARCHITECTURES "80") 
+endif()
+
+find_package(CUDAToolkit REQUIRED)
+
+
+# CUDA flags
+set(CUDA_FLAGS
+    -O3
+    -std=c++20
+    --use_fast_math
+    --expt-relaxed-constexpr
+    --expt-extended-lambda
+    -U__CUDA_NO_HALF_OPERATORS__
+    -U__CUDA_NO_HALF_CONVERSIONS__
+    -U__CUDA_NO_HALF2_OPERATORS__
+    -U__CUDA_NO_BFLOAT16_CONVERSIONS__
+    -DFLASHATTENTION_DISABLE_DROPOUT=1
+    -DFLASHATTENTION_DISABLE_ALIBI=1
+)
+
+# Collect source files
+file(GLOB CUDA_SOURCES 
+    "csrc/flash_attn/src/flash_fwd_hdim*.cu"
+    "csrc/flash_attn/src/flash_bwd_hdim*.cu"
+    "csrc/flash_attn/src/flash_fwd_split_hdim*.cu"
+)
+
+file(GLOB CC_SOURCES 
+    "csrc/flash_attn/*.cpp"
+)
+
+# Create CUDA extension
+pybind11_add_module(flash_api
+    ${CC_SOURCES}
+    ${CUDA_SOURCES}
+)
+
+set_property(TARGET flash_api PROPERTY JOB_POOL_COMPILE cuda)
+
+target_compile_options(flash_api PRIVATE
+    $<$<COMPILE_LANGUAGE:CUDA>:${CUDA_FLAGS}>
+)
+
+target_include_directories(flash_api PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/flash_attn
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/flash_attn/src
+    ${CMAKE_CURRENT_SOURCE_DIR}/csrc/cutlass/include
+)
+
+target_link_libraries(flash_api PRIVATE
+    CUDA::cudart
+)
+
+if(FLASH_ATTENTION_FORCE_CXX11_ABI)
+    target_compile_definitions(flash_api PRIVATE
+        _GLIBCXX_USE_CXX11_ABI=1
+    )
+endif()
+
+# Installation
+install(TARGETS flash_api
+        DESTINATION ${SKBUILD_PLATLIB_DIR}/flash_attn_jax_lib
+)
+
+install(DIRECTORY src/flash_attn_jax/
+        DESTINATION ${SKBUILD_PLATLIB_DIR}/flash_attn_jax
+        FILES_MATCHING PATTERN "*.py"
+)

README.md

@@ -1,19 +1,20 @@
 # FlashAttention JAX
 This repository provides a jax binding to <https://github.com/Dao-AILab/flash-attention>. To avoid depending on pytorch, since torch and jax installations often conflict, this is a fork of the official repo.
 
-Please see [Tri Dao's repo](https://github.com/Dao-AILab/flash-attention) for more information about flash attention.
+Please see [Tri Dao's repo](https://github.com/Dao-AILab/flash-attention) for more information about flash attention. Also check there for how to cite the authors if you used flash attention in your work.
 
 FlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).
 Please cite (see below) and credit FlashAttention if you use it.
 
 ## Installation
 
 Requirements:
-- CUDA 11.8 and above.
+- CUDA 12.8 and above.
 - Linux. Same story as with the pytorch repo. I haven't tested compilation of the jax bindings on windows.
-- JAX >=`0.4.24`. The custom sharding used for ring attention requires some somewhat advanced features.
+- JAX >= `0.5.*`. The custom call api changed in this version.
 
-To install: `pip install flash-attn-jax` will get the latest release from pypi. This gives you the cuda 12.3 build. If you want to use the cuda 11.8 build, you can install from the releases page (but according to jax's documentation, 11.8 will stop being supported for newer versions of jax).
+To install: `pip install flash-attn-jax` will get the latest release from pypi. This gives you the cuda 12.8
+build. CUDA 11 isn't supported any more (since jax stopped supporting it).
 
 ### Installing from source
 
@@ -25,7 +26,7 @@ cd flash-attn-jax
 cibuildwheel --only cp312-manylinux_x86_64 # I think cibuildwheel needs superuser privileges on some systems because of docker reasons?
 ```
 
-This will create a wheel in the `wheelhouse` directory. You can then install it with `pip install wheelhouse/flash_attn_jax_0.2.0-cp312-cp312-manylinux_x86_64.whl`. Or you could use setup.py to build the wheel and install it. You need cuda toolkit installed in that case.
+This will create a wheel in the `wheelhouse` directory. You can then install it with `pip install wheelhouse/flash_attn_jax_*.whl`. Or you could build it without docker using `uv build --wheel`. You need cuda installed in that case.
 
 ## Usage
 
@@ -45,15 +46,16 @@ This supports multi-query and grouped-query attention (when hk != h). The `softm
 Use jax.Array and shard your tensors along the length dimension, and flash_mha will automatically use the ring attention algorithm (forward and backward).
 
 ```py
-os.environ["XLA_FLAGS"] = '--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_async_collectives=true'
+os.environ["XLA_FLAGS"] = '--xla_gpu_enable_latency_hiding_scheduler=true'
 #...
 with Mesh(devices, axis_names=('len',)) as mesh:
         sharding = NamedSharding(mesh, P(None,'len')) # n l
         tokens = jax.device_put(tokens, sharding)
         # invoke your jax.jit'd transformer.forward
 ```
 
-It's not entirely reliable at hiding the communication latency though, depending on the whims of the xla optimizer. I'm waiting https://github.com/google/jax/issues/20864 to be fixed, then I can make it better.
+The latency hiding seems to be reliable now that some bugs have been fixed, as long as you enable the 
+latency hiding scheduler as above.
 
 ### GPU support
 
@@ -63,19 +65,3 @@ FlashAttention-2 currently supports:
    GPUs for now.
 2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).
 3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100/A800 or H100/H800~~. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.
-
-## Citation
-If you use this codebase, or otherwise found our work valuable, please cite:
-```
-@inproceedings{dao2022flashattention,
-  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
-  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
-  booktitle={Advances in Neural Information Processing Systems},
-  year={2022}
-}
-@article{dao2023flashattention2,
-  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
-  author={Dao, Tri},
-  year={2023}
-}
-```

Tupfile

@@ -10,6 +10,7 @@
 #include "static_switch.h"
 #include "flash.h"
 #include "flash_fwd_kernel.h"
+#include "kernel_traits.h"
 
 template<typename Kernel_traits, bool Is_dropout, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Return_softmax>
 __global__ void flash_fwd_kernel(__grid_constant__ const Flash_fwd_params params) {

csrc/flash_attn/src/flash_fwd_launch_template.h

@@ -4,11 +4,9 @@
 
 #pragma once
 
-#include "cute/algorithm/copy.hpp"
-
-#include "cutlass/cutlass.h"
-#include "cutlass/layout/layout.h"
-#include <cutlass/numeric_types.h>
+#include "cute/tensor.hpp"
+#include "cute/atom/mma_atom.hpp"
+#include "cute/atom/copy_atom.hpp"
 
 using namespace cute;
 

csrc/flash_attn/src/kernel_traits.h

@@ -0,0 +1,22 @@
+import os, sys
+from subprocess import Popen, PIPE
+import json
+import re
+
+Popen(["cmake", ".", "-B", "build"]).wait()
+with open("build/compile_commands.json", "r") as f:
+    compile_commands = json.load(f)
+
+# --options-file CMakeFiles/flash_attn_2_cuda.dir/includes_CUDA.rsp
+re_options = re.compile(r"--options-file ([A-Za-z0-9/\._]*)")
+
+for command in compile_commands:
+    if re_options.search(command["command"]):
+        m = re_options.search(command["command"])
+        options_file = m.group(1)
+        with open(os.path.join('build', options_file), "r") as f:
+            options = f.read()
+        command["command"] = command["command"].replace(m.group(0), options)
+
+with open("compile_commands.json", "w") as f:
+    json.dump(compile_commands, f, indent=2)

make_compile_commands.py

@@ -1,5 +1,70 @@
 [build-system]
-requires = ["setuptools", "wheel", "setuptools-cuda-cpp @ git+https://github.com/nshepperd/setuptools-cuda-cpp", "packaging", "pybind11"]
+requires = [
+    "scikit-build-core>=0.8.0",
+    "cmake>=3.18",
+    "ninja>=1.10",
+    "packaging",
+    "psutil",
+    "pybind11>=2.11.0",
+]
+build-backend = "scikit_build_core.build"
+
+[project]
+name = "flash_attn_jax"
+dynamic = ["version"]
+description = "Flash Attention port for JAX"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "BSD-3-Clause" }
+authors = [
+    { name = "Tri Dao", email = "[email protected]" },
+    { name = "Emily Shepperd", email = "[email protected]" }
+]
+dependencies = [
+    "jax>=0.5.0, <0.8.0"
+]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: BSD License",
+    "Operating System :: Unix",
+]
+
+[dependency-groups]
+test = [
+    "pytest>=7.0.0",
+    "einops",
+    "jax[cuda12]",
+]
+[project.urls]
+Homepage = "https://github.com/nshepperd/flash_attn_jax"
+
+[tool.scikit-build]
+wheel.expand-macos-universal-tags = false
+cmake.version = ">=3.26.1"
+ninja.version = ">=1.11"
+build.verbose = true
+cmake.build-type = "Release"
+cmake.args = []
+
+[tool.scikit-build.cmake.define]
+SKBUILD = "ON"
+FLASH_ATTENTION_FORCE_BUILD = { env = "FLASH_ATTENTION_FORCE_BUILD" }
+FLASH_ATTENTION_SKIP_CUDA_BUILD = { env = "FLASH_ATTENTION_SKIP_CUDA_BUILD" }
+FLASH_ATTENTION_FORCE_CXX11_ABI = { env = "FLASH_ATTENTION_FORCE_CXX11_ABI" }
+FLASH_ATTENTION_TRITON_AMD_ENABLE = { env = "FLASH_ATTENTION_TRITON_AMD_ENABLE" }
+FLASH_ATTN_CUDA_ARCHS = { env = "FLASH_ATTN_CUDA_ARCHS" }
+CMAKE_VERBOSE_MAKEFILE = "ON"
+
+[tool.scikit-build.metadata.version]
+provider = "scikit_build_core.metadata.regex"
+input = "src/flash_attn_jax/__init__.py"
 
 [tool.cibuildwheel]
-manylinux-x86_64-image = "sameli/manylinux_2_28_x86_64_cuda_12.3"
+# manylinux-x86_64-image = "quay.io/pypa/manylinux_2_28_x86_64:latest"
+before-all = "bash scripts/install-cuda-linux.sh"
+build = "cp312-manylinux_x86_64"
+repair-wheel-command = "auditwheel repair --exclude=libcudart.so* -w {dest_dir} {wheel}"
+
+[tool.cibuildwheel.environment]
+PATH="/opt/rh/gcc-toolset-13/root/usr/bin:/usr/local/cuda/bin:$PATH"
+CUDA_HOME="/usr/local/cuda"