More attention kernels to its own folder (#368)

facebookresearch · Aug 25, 2022 · 096e8ef · 096e8ef
1 parent 5958b18
commit 096e8ef
Show file tree

Hide file tree

Showing 11 changed files with 3 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -138,7 +138,9 @@ def get_extensions():
 
     sources = main_file + source_cpu
 
-    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))
+    source_cuda = glob.glob(
+        os.path.join(extensions_dir, "cuda", "**", "*.cu"), recursive=True
+    )
 
     sputnik_dir = os.path.join(this_dir, "third_party", "sputnik")
     cutlass_dir = os.path.join(this_dir, "third_party", "cutlass", "include")

diff --git a/...mponents/attention/csrc/cuda/attention.cu → .../csrc/cuda/mem_eff_attention/attention.cu b/...mponents/attention/csrc/cuda/attention.cu → .../csrc/cuda/mem_eff_attention/attention.cu
diff --git a/...n/csrc/cuda/attention_backward_generic.cu → ...f_attention/attention_backward_generic.cu b/...n/csrc/cuda/attention_backward_generic.cu → ...f_attention/attention_backward_generic.cu
diff --git a/...on/csrc/cuda/attention_forward_generic.cu → ...ff_attention/attention_forward_generic.cu b/...on/csrc/cuda/attention_forward_generic.cu → ...ff_attention/attention_forward_generic.cu
diff --git a/...rc/cuda/attention_scaling_coefs_updater.h → ...tention/attention_scaling_coefs_updater.h b/...rc/cuda/attention_scaling_coefs_updater.h → ...tention/attention_scaling_coefs_updater.h
diff --git a/...ponents/attention/csrc/cuda/debug_utils.h → ...csrc/cuda/mem_eff_attention/debug_utils.h b/...ponents/attention/csrc/cuda/debug_utils.h → ...csrc/cuda/mem_eff_attention/debug_utils.h
diff --git a/...ntion/csrc/cuda/epilogue_rescale_output.h → ...m_eff_attention/epilogue_rescale_output.h b/...ntion/csrc/cuda/epilogue_rescale_output.h → ...m_eff_attention/epilogue_rescale_output.h
diff --git a/...rc/cuda/epilogue_thread_apply_logsumexp.h → ...tention/epilogue_thread_apply_logsumexp.h b/...rc/cuda/epilogue_thread_apply_logsumexp.h → ...tention/epilogue_thread_apply_logsumexp.h
diff --git a/...ts/attention/csrc/cuda/find_default_mma.h → ...cuda/mem_eff_attention/find_default_mma.h b/...ts/attention/csrc/cuda/find_default_mma.h → ...cuda/mem_eff_attention/find_default_mma.h
diff --git a/...s/attention/csrc/cuda/gemm_kernel_utils.h → ...uda/mem_eff_attention/gemm_kernel_utils.h b/...s/attention/csrc/cuda/gemm_kernel_utils.h → ...uda/mem_eff_attention/gemm_kernel_utils.h
diff --git a/...nents/attention/csrc/cuda/mma_from_smem.h → ...rc/cuda/mem_eff_attention/mma_from_smem.h b/...nents/attention/csrc/cuda/mma_from_smem.h → ...rc/cuda/mem_eff_attention/mma_from_smem.h