Add flash-attn

recipes/flash-attn/conda-forge.yml

@@ -0,0 +1,2 @@
+azure:
+  timeout_minutes: 360

recipes/flash-attn/meta.yaml

@@ -0,0 +1,66 @@
+{% set name = "flash-attn" %}
+{% set version = "2.5.8" %}
+
+package:
+  name: {{ name|lower }}
+  version: {{ version }}
+
+source:
+  - url: https://pypi.io/packages/source/{{ name[0] }}/{{ name }}/flash_attn-{{ version }}.tar.gz
+    sha256: 2e5b2bcff6d5cff40d494af91ecd1eb3c5b4520a6ce7a0a8b1f9c1ed129fb402
+  # Overwrite with a simpler build script that doesn't try to revend pre-compiled binaries
+  - path: pyproject.toml
+  - path: setup.py
+
+build:
+  number: 0
+  script: {{ PYTHON }} -m pip install . -vvv --no-deps --no-build-isolation
+  script_env:
+    - MAX_JOBS=$CPU_COUNT
+    - TORCH_CUDA_ARCH_LIST=8.0;8.6;8.9;9.0+PTX
+  skip: true           # [cuda_compiler_version in (undefined, "None")]
+  skip: true           # [not linux]
+  rpaths:
+    - lib/
+  # PyTorch libs are in site-packages instead of with other shared objects
+    - {{ SP_DIR }}/torch/lib/
+
+requirements:
+  build:
+    - {{ compiler('c') }}
+    - {{ compiler('cxx') }}
+    - {{ compiler('cuda') }}
+    - {{ stdlib('c') }}
+    - ninja
+  host:
+    - cuda-version {{ cuda_compiler_version }}  # same cuda for host and build
+    - cuda-cudart-dev  # [(cuda_compiler_version or "").startswith("12")]
+    - libtorch         # required until pytorch run_exports libtorch
+    - pip
+    - python
+    - pytorch
+    - pytorch =*=cuda*
+    - setuptools
+  run:
+    - einops
+    - python
+    - pytorch =*=cuda*
+
+test:
+  imports:
+    - flash_attn
+  commands:
+    - pip check
+  requires:
+    - pip
+
+about:
+  home: https://github.com/Dao-AILab/flash-attention
+  summary: 'Flash Attention: Fast and Memory-Efficient Exact Attention'
+  license: BSD-3-Clause
+  license_file: LICENSE
+
+extra:
+  recipe-maintainers:
+    - carterbox
+    - weiji14

recipes/flash-attn/pyproject.toml

@@ -0,0 +1,28 @@
+[build-system]
+requires = ["setuptools>=62", "torch", "ninja"]
+build-backend = "setuptools.build_meta"
+
+[project]
+dynamic = ["version"]
+name = "flash_attn"
+authors = [
+    {name = "Tri Dao", email = "trid@cs.stanford.edu"},
+]
+description="Flash Attention: Fast and Memory-Efficient Exact Attention"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: Unix",
+    "License :: OSI Approved :: BSD License",
+]
+readme = "README.md"
+license = {file = "LICENSE"}
+dependencies = [
+    "torch",
+    "einops",
+]
+
+[project.urls]
+Homepage = "https://github.com/Dao-AILab/flash-attention"
+
+[tool.setuptools.dynamic]
+version = {attr = "flash_attn.__version__"}

recipes/flash-attn/setup.py

@@ -0,0 +1,109 @@
+# Copyright (c) 2023, Tri Dao.
+# Copyright (c) 2024, Conda-forge Contributors.
+
+"""Since this package is a pytorch extension, this setup file uses the custom
+CUDAExtension build system from pytorch. This ensures that compatible compiler
+args, headers, etc for pytorch.
+
+Read more at the pytorch docs:
+https://pytorch.org/docs/stable/cpp_extension.html#torch.utils.cpp_extension.CUDAExtension
+"""
+
+import pathlib
+
+from setuptools import setup, find_packages
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+_this_dir = pathlib.Path(__file__).parent.absolute()
+
+setup(
+    packages=find_packages(
+        include=["flash_attn*"],
+    ),
+    ext_modules=[
+        CUDAExtension(
+            name="flash_attn_2_cuda",
+            sources=[
+                "csrc/flash_attn/flash_api.cpp",
+                "csrc/flash_attn/src/flash_fwd_hdim32_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim32_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim64_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim96_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim96_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim128_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim128_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim160_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim160_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim192_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim192_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim224_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim224_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim256_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_hdim256_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim32_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim32_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim64_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim64_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim96_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim96_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim128_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim128_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim160_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim160_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim192_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim192_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim224_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim224_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim256_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_bwd_hdim256_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim32_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim32_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim64_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim96_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim96_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim128_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim160_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim160_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim192_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim192_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim224_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim224_bf16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim256_fp16_sm80.cu",
+                "csrc/flash_attn/src/flash_fwd_split_hdim256_bf16_sm80.cu",
+            ],
+            extra_compile_args={
+                "cxx": [
+                    "-std=c++17",
+                ],
+                "nvcc": [
+                    "-std=c++17",
+                    "-U__CUDA_NO_HALF_OPERATORS__",
+                    "-U__CUDA_NO_HALF_CONVERSIONS__",
+                    "-U__CUDA_NO_HALF2_OPERATORS__",
+                    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+                    "--expt-relaxed-constexpr",
+                    "--expt-extended-lambda",
+                    "--use_fast_math",
+                    # "--ptxas-options=-v",
+                    # "--ptxas-options=-O2",
+                    # "-lineinfo",
+                    # "-DFLASHATTENTION_DISABLE_BACKWARD",
+                    # "-DFLASHATTENTION_DISABLE_DROPOUT",
+                    # "-DFLASHATTENTION_DISABLE_ALIBI",
+                    # "-DFLASHATTENTION_DISABLE_UNEVEN_K",
+                    # "-DFLASHATTENTION_DISABLE_LOCAL",
+                ],
+            },
+            include_dirs=[
+                _this_dir / "csrc" / "flash_attn",
+                _this_dir / "csrc" / "flash_attn" / "src",
+                _this_dir / "csrc" / "cutlass" / "include",
+            ],
+        ),
+    ],
+    cmdclass={"build_ext": BuildExtension},
+    zip_safe=False,
+)