Merge branch 'main' into will_fp8_mr

setup.py

@@ -11,12 +11,35 @@ from torch.utils.cpp_extension import (
     IS_WINDOWS,
 )
 
+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",
+        "csrc/flash_fwd_mla_bf16_sm90.cu",
+        "csrc/flash_fwd_mla_fp8_sm90.cu",
+        "csrc/flash_fwd_mla_metadata.cu",
+    ]
+
+    if not DISABLE_FP16:
+        sources.append("csrc/flash_fwd_mla_fp16_sm90.cu")
+
+    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 = []
@@ -34,14 +57,9 @@ ext_modules = []
 ext_modules.append(
     CUDAExtension(
         name="flash_mla_cuda",
-        sources=[
-            "csrc/flash_api.cpp",
-            "csrc/flash_mla_utils.cu",
-            "csrc/flash_fwd_mla_bf16_sm90.cu",
-            "csrc/flash_fwd_mla_fp8_sm90.cu",
-        ],
+        sources=get_sources(),
         extra_compile_args={
-            "cxx": cxx_args,
+            "cxx": cxx_args + get_features_args(),
             "nvcc": append_nvcc_threads(
                 [
                     "-O3",
@@ -60,7 +78,7 @@ ext_modules.append(
                     "--ftemplate-backtrace-limit=0"
                 ]
                 + cc_flag
-            ),
+            ) + get_features_args(),
         },
         include_dirs=[
             Path(this_dir) / "csrc",