setup.py

import os
from pathlib import Path
from datetime import datetime
import subprocess

from setuptools import setup, find_packages

from torch.utils.cpp_extension import (
    BuildExtension,
    CUDAExtension,
    IS_WINDOWS,
)



def get_sources():
    sources = [
        "csrc/flash_api.cpp",
    ]
    return sources

def get_ixformer_lib_path():
    import ixformer
    path = os.path.dirname(ixformer.__file__)
    return path

cc_flag = []
cc_flag.append("-gencode")

this_dir = os.path.dirname(os.path.abspath(__file__))

if IS_WINDOWS:
    raise NotImplementedError()
else:
    cxx_args = ["-O3", "-std=c++17", "-DNDEBUG", "-Wno-deprecated-declarations"]

ext_modules = []
ext_modules.append(
    CUDAExtension(
        name="flash_mla_cuda",
        sources=get_sources(),
        extra_compile_args={
            "cxx": cxx_args,
            "nvcc": []
        },
        extra_link_args=[f"-L{get_ixformer_lib_path()}", '-lixformer']
    )
)

def get_local_version() -> str:
    LOCAL_VERSION_IDENTIFIER = os.getenv("LOCAL_VERSION_IDENTIFIER", "")
    if LOCAL_VERSION_IDENTIFIER:
        LOCAL_VERSION_IDENTIFIER = "+" + LOCAL_VERSION_IDENTIFIER
    return LOCAL_VERSION_IDENTIFIER

setup(
    name="flash_mla",
    version="1.0.0" + get_local_version(),
    packages=find_packages(include=['flash_mla']),
    ext_modules=ext_modules,
    cmdclass={"build_ext": BuildExtension},
)