Commit 512c98ee authored by Tri Dao's avatar Tri Dao
Browse files

Add Cutlass as submodule

parent ad6c694b
[submodule "csrc/flash_attn/cutlass"]
path = csrc/flash_attn/cutlass
url = https://github.com/NVIDIA/cutlass.git
Subproject commit 319a389f42b776fae5701afcb943fc03be5b5c25
...@@ -111,6 +111,7 @@ if int(bare_metal_major) < 11: ...@@ -111,6 +111,7 @@ if int(bare_metal_major) < 11:
cc_flag.append("-gencode") cc_flag.append("-gencode")
cc_flag.append("arch=compute_80,code=sm_80") cc_flag.append("arch=compute_80,code=sm_80")
subprocess.run(["git", "submodule", "update", "--init", "csrc/flash_attn/cutlass"])
ext_modules.append( ext_modules.append(
CUDAExtension( CUDAExtension(
name="flash_attn_cuda", name="flash_attn_cuda",
...@@ -141,6 +142,7 @@ ext_modules.append( ...@@ -141,6 +142,7 @@ ext_modules.append(
include_dirs=[ include_dirs=[
Path(this_dir) / 'csrc' / 'flash_attn', Path(this_dir) / 'csrc' / 'flash_attn',
Path(this_dir) / 'csrc' / 'flash_attn' / 'src', Path(this_dir) / 'csrc' / 'flash_attn' / 'src',
Path(this_dir) / 'csrc' / 'flash_attn' / 'cutlass' / 'include',
], ],
) )
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment