[ROCm] Add support for Punica kernels on AMD GPUs (#3140)

Co-authored-by: miloice <jeffaw99@hotmail.com>
This commit is contained in:
kliuae
2024-05-10 00:19:50 +08:00
committed by GitHub
parent 0ee535b294
commit ff5abcd746
10 changed files with 287 additions and 26 deletions

View File

@@ -385,12 +385,12 @@ ext_modules = []
if _is_cuda():
ext_modules.append(CMakeExtension(name="vllm._moe_C"))
if _install_punica():
ext_modules.append(CMakeExtension(name="vllm._punica_C"))
if not _is_neuron():
ext_modules.append(CMakeExtension(name="vllm._C"))
if _install_punica():
ext_modules.append(CMakeExtension(name="vllm._punica_C"))
package_data = {
"vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}