DeepSeek/FlashMLA
1
0
Fork 0
mirror of https://github.com/deepseek-ai/FlashMLA synced 2025-06-26 18:15:54 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

cuda12.8 recommendation

Browse Source
This commit is contained in:
ljss
2025-02-26 00:05:57 +08:00
parent b549289fb4
commit 4edea86f9e
2 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
setup.py
View File

@@ -13,10 +13,12 @@ from torch.utils.cpp_extension import (
DISABLE_FP16 = os.getenv("FLASH_MLA_DISABLE_FP16", "FALSE") == "TRUE"
def append_nvcc_threads(nvcc_extra_args):
nvcc_threads = os.getenv("NVCC_THREADS") or "32"
return nvcc_extra_args + ["--threads", nvcc_threads]
def get_sources():
sources = [
"csrc/flash_api.cpp",
@@ -29,12 +31,14 @@ def get_sources():
return sources
def get_features_args():
features_args = []
if DISABLE_FP16:
features_args.append("-DFLASH_MLA_DISABLE_FP16")
return features_args
subprocess.run(["git", "submodule", "update", "--init", "csrc/cutlass"])
cc_flag = []