[Kernel][Hardware][Amd]Custom paged attention kernel for rocm (#8310)

2024-09-13 19:01:11 -05:00
parent 851725202a
commit 1ef0d2efd0
8 changed files with 1372 additions and 17 deletions
--- a/setup.py
+++ b/setup.py
@@ -462,6 +462,9 @@ if _build_core_ext():
 if _is_cuda() or _is_hip():
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

+if _is_hip():
+    ext_modules.append(CMakeExtension(name="vllm._rocm_C"))
+
 if _build_custom_ops():
    ext_modules.append(CMakeExtension(name="vllm._C"))