fix: add include path

setup.py

@@ -101,6 +101,7 @@ if __name__ == "__main__":
         "third_party/mio/include",
         "third_party/spdlog/include",
         "third_party/Block-Sparse-Attention/csrc/block_sparse_attn",
+        "third_party/flash_c_api/include",
     ]
 
     INCLUDE_DIRS = [os.path.join(ROOT_DIR, dir) for dir in INCLUDE_DIRS]
@@ -124,13 +125,15 @@ if __name__ == "__main__":
 
     #assert len(sm_targets) > 0, "No SM targets found"
 
-    GCC_FLAGS = ["-DENABLE_BF16=1", "-DBUILD_NUNCHAKU=1", "-fvisibility=hidden", "-g", "-std=c++2a", "-UNDEBUG", "-Og"]
+    GCC_FLAGS = ["-DENABLE_BF16=1", "-DBUILD_NUNCHAKU=1", "-fvisibility=hidden", "-g", "-std=c++2a", "-UNDEBUG", "-O1"]
     MSVC_FLAGS = ["/DENABLE_BF16=1", "/DBUILD_NUNCHAKU=1", "/std:c++2a", "/UNDEBUG", "/Zc:__cplusplus", "/FS"]
     NVCC_FLAGS = [
         "-DDCU_ASM",
+        "-DUSE_ROCM",
         "-DENABLE_BF16=1",
         "-DBUILD_NUNCHAKU=1",
         "-g",
+        "-O1",
         "-std=c++2a",
         "-UNDEBUG",
         "-mllvm",
@@ -185,28 +188,30 @@ if __name__ == "__main__":
             # *ncond(
             #     "third_party/Block-Sparse-Attention/csrc/block_sparse_attn/src/flash_fwd_block_hdim128_bf16_sm80.cu"
             # ),
-            # "src/kernels/activation_kernels.cu",
-            # "src/kernels/layernorm_kernels.cu",
-            # "src/kernels/misc_kernels.cu",
-            # "src/kernels/zgemm/gemm_w4a4.cu",
-            # "src/kernels/zgemm/gemm_w4a4_test.cu",
-            # "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4.cu",
-            # "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4_fasteri2f.cu",
-            # "src/kernels/zgemm/gemm_w4a4_launch_fp16_fp4.cu",
-            # "src/kernels/zgemm/gemm_w4a4_launch_bf16_int4.cu",
-            # "src/kernels/zgemm/gemm_w4a4_launch_bf16_fp4.cu",
-            # "src/kernels/zgemm/gemm_w8a8.cu",
-            # "src/kernels/zgemm/attention.cu",
-            # "src/kernels/dwconv.cu",
-            # "src/kernels/gemm_batched.cu",
-            # "src/kernels/gemm_f16.cu",
-            # "src/kernels/awq/gemm_awq.cu",
+            "src/kernels/activation_kernels.cu",
+            "src/kernels/layernorm_kernels.cu",
+            "src/kernels/misc_kernels.cu",
+            "src/kernels/zgemm/gemm_w4a4.cu",
+            "src/kernels/zgemm/gemm_w4a4_test.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_int4_fasteri2f.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_fp16_fp4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_bf16_int4.cu",
+            "src/kernels/zgemm/gemm_w4a4_launch_bf16_fp4.cu",
+            "src/kernels/zgemm/gemm_w8a8.cu",
+            "src/kernels/zgemm/attention.cu",
+            "src/kernels/dwconv.cu",
+            "src/kernels/gemm_batched.cu",
+            "src/kernels/gemm_f16.cu",
+            "src/kernels/awq/gemm_awq.cu",
             "src/kernels/awq/gemv_awq.cu",
-            *ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/flash_api.cpp"),
-            *ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/flash_api_adapter.cpp"),
+            #*ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/flash_api.cpp"),
+            #*ncond("third_party/Block-Sparse-Attention/csrc/block_sparse_attn/flash_api_adapter.cpp"),
         ],
         extra_compile_args={"gcc": GCC_FLAGS, "msvc": MSVC_FLAGS, "nvcc": NVCC_FLAGS, "nvcc_msvc": NVCC_MSVC_FLAGS},
         include_dirs=INCLUDE_DIRS,
+        libraries=["flash_atten_c"],
+        library_dirs=["third_party/flash_c_api/lib"]
     )
 
     setuptools.setup(