setup.py

from setuptools import setup, find_packages
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os

setup(
    name='flash_attn_mma',
    version='0.1.0',
    packages=find_packages(),
    ext_modules=[
        CUDAExtension(
            name='flash_attn_mma._C',
            sources=[
                'src/flash_attn_api.cu',
                'src/flash_fwd_kernel.cu',
                'src/flash_fwd_splitkv.cu',
                'src/flash_fwd_splitd.cu',
                'src/flash_bwd_kernel.cu',
            ],
            include_dirs=[os.path.join(os.path.abspath(os.path.dirname(__file__)), 'include')],
            extra_compile_args={'cxx': ['-O3'], 'nvcc': ['-O3', '-lineinfo', '--ptxas-options=-v']}
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)