Skip to content

Commit 524b322

Browse files
charlifuAlvant
authored andcommitted
[Kernel][Hardware][Amd]Custom paged attention kernel for rocm (vllm-project#8310)
Signed-off-by: Alvant <alvasian@yandex.ru>
1 parent f0f3e3e commit 524b322

File tree

8 files changed

+1371
-16
lines changed

8 files changed

+1371
-16
lines changed

CMakeLists.txt

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,12 +324,35 @@ define_gpu_extension_target(
324324
WITH_SOABI)
325325

326326

327+
if(VLLM_GPU_LANG STREQUAL "HIP")
328+
#
329+
# _rocm_C extension
330+
#
331+
set(VLLM_ROCM_EXT_SRC
332+
"csrc/rocm/torch_bindings.cpp"
333+
"csrc/rocm/attention.cu")
334+
335+
define_gpu_extension_target(
336+
_rocm_C
337+
DESTINATION vllm
338+
LANGUAGE ${VLLM_GPU_LANG}
339+
SOURCES ${VLLM_ROCM_EXT_SRC}
340+
COMPILE_FLAGS ${VLLM_GPU_FLAGS}
341+
ARCHITECTURES ${VLLM_GPU_ARCHES}
342+
USE_SABI 3
343+
WITH_SOABI)
344+
endif()
345+
327346

328347
if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
329348
message(STATUS "Enabling C extension.")
330349
add_dependencies(default _C)
331350

332351
message(STATUS "Enabling moe extension.")
333352
add_dependencies(default _moe_C)
353+
endif()
334354

355+
if(VLLM_GPU_LANG STREQUAL "HIP")
356+
message(STATUS "Enabling rocm extension.")
357+
add_dependencies(default _rocm_C)
335358
endif()

0 commit comments

Comments
 (0)