support fp16

2025-06-26 18:15:54 +00:00 · 2025-02-24 01:58:53 -08:00
parent 15a82b81b8
commit 65fb7732fc
7 changed files with 139 additions and 91 deletions
--- a/setup.py
+++ b/setup.py
@@ -37,6 +37,8 @@ ext_modules.append(
        sources=[
            "csrc/flash_api.cpp",
            "csrc/flash_fwd_mla_bf16_sm90.cu",
+            "csrc/flash_fwd_mla_fp16_sm90.cu",
+            "csrc/flash_fwd_mla_metadata.cu",
        ],
        extra_compile_args={
            "cxx": cxx_args,