fix compile

2025-06-26 18:15:54 +00:00 · 2025-02-25 21:30:39 +08:00
parent 7409203f44
commit c50d29d170
3 changed files with 86 additions and 76 deletions
--- a/setup.py
+++ b/setup.py
@@ -36,6 +36,7 @@ ext_modules.append(
        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",
        ],