Fa cmake #29
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,221 @@ | ||
| # Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
|
|
||
| # Adapted from https://github.com/NVIDIA/apex/blob/master/setup.py | ||
| import ast | ||
| import logging | ||
| import os | ||
| import platform | ||
| import re | ||
| import shutil | ||
| import subprocess | ||
| import sys | ||
| import warnings | ||
| from pathlib import Path | ||
|
|
||
| from packaging.version import parse | ||
| from setuptools import find_packages, setup | ||
| from setuptools.command.install import install as _install | ||
| from wheel.bdist_wheel import bdist_wheel as _bdist_wheel | ||
|
|
||
| import paddle | ||
| from paddle.utils.cpp_extension.extension_utils import find_cuda_home | ||
|
|
||
| version_detail = sys.version_info | ||
| python_version = platform.python_version() | ||
| version = version_detail[0] + version_detail[1] / 10 | ||
| env_version = os.getenv("PY_VERSION") | ||
|
|
||
| if version < 3.7: | ||
|
There was a problem hiding this comment. 这个判断比较粗糙。 |
||
| raise RuntimeError( | ||
| f"Paddle only supports Python version >= 3.7 now," | ||
| f"you are using Python {python_version}" | ||
| ) | ||
| elif env_version is None: | ||
| print(f"export PY_VERSION = { python_version }") | ||
| os.environ["PY_VERSION"] = python_version | ||
|
|
||
| elif env_version != version: | ||
| warnings.warn( | ||
| f"You set PY_VERSION={env_version}, but" | ||
| f"your current python environment is {version}" | ||
| f"we will use your current python version to execute" | ||
| ) | ||
| os.environ["PY_VERSION"] = python_version | ||
|
|
||
| paddle_include_path = paddle.sysconfig.get_include() | ||
| paddle_lib_path = paddle.sysconfig.get_lib() | ||
|
There was a problem hiding this comment. 为了将安装的libflash_attn_advanced.so拷贝到paddle路径下 There was a problem hiding this comment. 这种方式我不太确定，请@sneaxiy 也看下。我理解：
There was a problem hiding this comment. FA应该作为Paddle侧的一个外部算子？类似于xformers跟PyTorch的关系，而不是新的so直接去替换Paddle里的FA的so？ |
||
|
|
||
| print("Paddle Include Path:", paddle_include_path) | ||
| print("Paddle Lib Path:", paddle_lib_path) | ||
|
|
||
| # preparing parameters for setup() | ||
| paddle_version = paddle.version.full_version | ||
| cuda_version = paddle.version.cuda_version | ||
|
|
||
|
|
||
| with open("../../README.md", "r", encoding="utf-8") as fh: | ||
| long_description = fh.read() | ||
|
|
||
|
|
||
| # ninja build does not work unless include_dirs are abs path | ||
| this_dir = os.path.dirname(os.path.abspath(__file__)) | ||
| CUDA_HOME = find_cuda_home() | ||
| PACKAGE_NAME = "paddle_flash_attn" | ||
|
|
||
|
|
||
| def get_platform(): | ||
| """ | ||
| Returns the platform name as used in wheel filenames. | ||
| """ | ||
| if sys.platform.startswith('linux'): | ||
| return 'linux_x86_64' | ||
| elif sys.platform == 'darwin': | ||
| mac_version = '.'.join(platform.mac_ver()[0].split('.')[:2]) | ||
| return f'macosx_{mac_version}_x86_64' | ||
| elif sys.platform == 'win32': | ||
| return 'win_amd64' | ||
| else: | ||
| raise ValueError(f'Unsupported platform: {sys.platform}') | ||
|
|
||
|
|
||
| def get_cuda_bare_metal_version(cuda_dir): | ||
| raw_output = subprocess.check_output( | ||
| [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True | ||
| ) | ||
| output = raw_output.split() | ||
| release_idx = output.index("release") + 1 | ||
| bare_metal_version = parse(output[release_idx].split(",")[0]) | ||
|
|
||
| return raw_output, bare_metal_version | ||
|
|
||
|
|
||
| def _is_cuda_available(): | ||
| """ | ||
| Check whether CUDA is available. | ||
| """ | ||
| try: | ||
| assert len(paddle.static.cuda_places()) > 0 | ||
|
|
||
| return True | ||
| except Exception as e: | ||
| logging.warning( | ||
| "You are using GPU version PaddlePaddle, but there is no GPU " | ||
| "detected on your machine. Maybe CUDA devices is not set properly." | ||
| f"\n Original Error is {e}" | ||
| ) | ||
| return False | ||
|
|
||
|
|
||
| check = _is_cuda_available() | ||
| cmdclass = {} | ||
|
|
||
|
|
||
| def get_package_version(): | ||
| with open(Path(this_dir) / "../flash_attn" / "__init__.py", "r") as f: | ||
| version_match = re.search( | ||
| r"^__version__\s*=\s*(.*)$", f.read(), re.MULTILINE | ||
| ) | ||
| public_version = ast.literal_eval(version_match.group(1)) | ||
| local_version = os.environ.get("FLASH_ATTN_LOCAL_VERSION") | ||
| if local_version: | ||
| return f"{public_version}+{local_version}" | ||
| else: | ||
| return str(public_version) | ||
|
|
||
|
|
||
| def get_data_files(): | ||
| data_files = [] | ||
| #source_lib_path = 'libflashattn.so' | ||
| #data_files.append((".", [source_lib_path])) | ||
| data_files.append((".", ['libflashattn_advanced.so'])) | ||
| return data_files | ||
|
|
||
|
|
||
| class CustomWheelsCommand(_bdist_wheel): | ||
| """ | ||
| The CachedWheelsCommand plugs into the default bdist wheel, which is ran by pip when it cannot | ||
| find an existing wheel (which is currently the case for all flash attention installs). We use | ||
| the environment parameters to detect whether there is already a pre-built version of a compatible | ||
| wheel available and short-circuits the standard full build pipeline. | ||
| """ | ||
|
|
||
| def run(self): | ||
| self.run_command('build_ext') | ||
| super().run() | ||
| # Determine the version numbers that will be used to determine the correct wheel | ||
| # We're using the CUDA version used to build paddle, not the one currently installed | ||
| # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME) | ||
| python_version = f"cp{sys.version_info.major}{sys.version_info.minor}" | ||
| platform_name = get_platform() | ||
| flash_version = get_package_version() | ||
| cxx11_abi = "" # str(paddle._C.-D_GLIBCXX_USE_CXX11_ABI).upper() | ||
|
|
||
| # Determine wheel URL based on CUDA version, paddle version, python version and OS | ||
| wheel_filename = f'{PACKAGE_NAME}-{flash_version}-cu{cuda_version}-paddle{paddle_version}-{python_version}-{python_version}-{platform_name}.whl' | ||
|
There was a problem hiding this comment. whl包里面不需要加paddle版本吧？本身flashattn对paddle版本并没有依赖，是paddle对flashattn版本存在依赖 There was a problem hiding this comment. 好的这个后续会去掉，现在的是默认版本：paddle_flash_attn-2.0.8-cp37-none-any.whl |
||
| impl_tag, abi_tag, plat_tag = self.get_tag() | ||
| original_wheel_name = ( | ||
| f"{self.wheel_dist_name}-{impl_tag}-{abi_tag}-{plat_tag}" | ||
| ) | ||
|
|
||
| # new_wheel_name = wheel_filename | ||
| new_wheel_name = ( | ||
| f"{self.wheel_dist_name}-{python_version}-{abi_tag}-{plat_tag}" | ||
| ) | ||
| shutil.move( | ||
| f"{self.dist_dir}/{original_wheel_name}.whl", | ||
| f"{self.dist_dir}/{new_wheel_name}.whl", | ||
| ) | ||
|
|
||
|
|
||
| class CustomInstallCommand(_install): | ||
| def run(self): | ||
| _install.run(self) | ||
| install_path = self.install_lib | ||
| source_lib_path = os.path.abspath('libflashattn_advanced.so') | ||
| destination_lib_path = os.path.join(paddle_lib_path, 'libflashattn_advanced.so') | ||
| shutil.copy(f"{source_lib_path}", f"{destination_lib_path}") | ||
|
|
||
|
|
||
| setup( | ||
| name=PACKAGE_NAME, | ||
| version=get_package_version(), | ||
| packages=find_packages(), | ||
| data_files=get_data_files(), | ||
| package_data={PACKAGE_NAME: ['build/libflashattn.so']}, | ||
| author_email="[email protected]", | ||
| description="Flash Attention: Fast and Memory-Efficient Exact Attention", | ||
| long_description=long_description, | ||
| long_description_content_type="text/markdown", | ||
| url="https://github.com/PaddlePaddle/flash-attention", | ||
| classifiers=[ | ||
| "Programming Language :: Python :: 37", | ||
| "License :: OSI Approved :: BSD License", | ||
| "Operating System :: Unix", | ||
| ], | ||
| cmdclass={ | ||
| 'bdist_wheel': CustomWheelsCommand, | ||
| 'install': CustomInstallCommand, | ||
| }, | ||
| python_requires=">=3.7", | ||
| install_requires=[ | ||
| "common", | ||
| "dual", | ||
| "tight>=0.1.0", | ||
| "data", | ||
| "prox", | ||
| "ninja", # Put ninja before paddle if paddle depends on it | ||
| "einops", | ||
| "packaging", | ||
| ], | ||
| ) | ||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
这个写法在外部设置了
-DNVCC_ARCH_BIN=...的情况下，取值会是多少，是80还是外部设置的值？There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
外部设置的值， 已经做过实验会拿到外部设置的