Skip to content

GitLab

Commit a1eef562 authored by shenzhe's avatar shenzhe Committed by zhanghj2
Browse files

Add DSA MLS sparse prefill dispatch

parent 4e0bdf6e
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
setup.py
View file @ a1eef562
...@@ -100,6 +100,7 @@ ext_modules.append( ...@@ -100,6 +100,7 @@ ext_modules.append(
"csrc/gfx93/prefill/sparse/instantiations/phase1_k512_topklen.cu", "csrc/gfx93/prefill/sparse/instantiations/phase1_k512_topklen.cu",
"csrc/gfx93/prefill/sparse/instantiations/phase1_k576.cu", "csrc/gfx93/prefill/sparse/instantiations/phase1_k576.cu",
"csrc/gfx93/prefill/sparse/instantiations/phase1_k576_topklen.cu", "csrc/gfx93/prefill/sparse/instantiations/phase1_k576_topklen.cu",
"csrc/gfx93/prefill/sparse/dsa_mls/fwd.cu",
"csrc/extension/flash_fwd_mla_bf16_gfx936.cu", "csrc/extension/flash_fwd_mla_bf16_gfx936.cu",
"csrc/extension/flash_fwd_mla_fp16_gfx936.cu", "csrc/extension/flash_fwd_mla_fp16_gfx936.cu",
...@@ -130,7 +131,8 @@ ext_modules.append( ...@@ -130,7 +131,8 @@ ext_modules.append(
Path(this_dir) / "csrc", Path(this_dir) / "csrc",
Path(this_dir) / "csrc" / "kerutils" / "include", # TODO Remove me Path(this_dir) / "csrc" / "kerutils" / "include", # TODO Remove me
Path(this_dir) / "csrc" / "gfx93", Path(this_dir) / "csrc" / "gfx93",
Path(this_dir) / "csrc" / "cutlass" / "include", Path(this_dir) / "csrc" / "cutlass" / "cutlass_3.2.1" / "include",
Path(this_dir) / "csrc" / "gfx93" / "prefill" / "sparse" / "dsa_mls" / "legacy" / "include",
], ],
) )
) )
...@@ -167,4 +169,3 @@ setup( ...@@ -167,4 +169,3 @@ setup(
cmdclass={"build_ext": BuildExtension}, cmdclass={"build_ext": BuildExtension},
install_requires=[pytorch_dep], install_requires=[pytorch_dep],
) )
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