Skip to content

perf(hip): enable the fused GPU sampler on the HIP backend #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DeanoC wants to merge 1 commit into
base: main
Choose a base branch
from
+39 −4
Open

perf(hip): enable the fused GPU sampler on the HIP backend #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 28 additions & 4 deletions server/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -400,6 +400,21 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip")
# rms_norm_hip.cu is needed by the HIP chunk-B graph path regardless of SM80_EQUIV.
target_sources(dflash_common PRIVATE src/rms_norm_hip.cu)
set_source_files_properties(src/rms_norm_hip.cu PROPERTIES LANGUAGE HIP)
# GPU port of the sample_logits chain (fused penalty+softmax+draw). Shares
# the CUDA source verbatim — compiled as HIP, with its <cuda_runtime.h>
# resolving to the hip_compat/ shim (already on dflash_common's include path).
# The kernel's warp reductions assume a 32-lane wavefront, which holds on
# RDNA (gfx11xx/gfx12xx wave32). Opted into at runtime via DFLASH_GPU_SAMPLE;
# turn off at configure time with -DDFLASH_GPU_SAMPLER=OFF.
option(DFLASH_GPU_SAMPLER "Build the GPU sample_logits path" ON)
if(DFLASH_GPU_SAMPLER)
target_sources(dflash_common PRIVATE src/common/geometric_sampler_cuda.cu)
set_source_files_properties(src/common/geometric_sampler_cuda.cu
PROPERTIES LANGUAGE HIP)
# PUBLIC so the backends and test executables that call sample_logits()
# also compile their GPU dispatch branch.
target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_GPU_SAMPLER=1)
endif()
if(DFLASH27B_HIP_SM80_EQUIV)
find_path(DFLASH27B_ROCWMMA_INCLUDE_DIR rocwmma/rocwmma.hpp
HINTS "${_dflash_rocm_root}/include" /opt/rocm/include
Expand Down Expand Up @@ -656,12 +671,21 @@ if(DFLASH27B_TESTS)
target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart)
add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda)
endif()
# GPU port of the sample_logits chain vs the CPU reference. CUDA only:
# geometric_sampler_cuda.cu is compiled into dflash_common solely on the cuda backend.
if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gpu_sampler_cuda.cpp")
# GPU port of the sample_logits chain vs the CPU reference. Built on whichever
# backend compiled geometric_sampler_cuda.cu into dflash_common (cuda as CUDA,
# hip as HIP). On hip the CUDA-spelled test source compiles as HIP via the
# hip_compat/ shim, exactly like test_flashprefill_kernels above.
if(DFLASH_GPU_SAMPLER AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gpu_sampler_cuda.cpp")
add_executable(test_gpu_sampler_cuda test/test_gpu_sampler_cuda.cpp)
target_include_directories(test_gpu_sampler_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
target_link_libraries(test_gpu_sampler_cuda PRIVATE dflash_common CUDA::cudart)
if(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
target_link_libraries(test_gpu_sampler_cuda PRIVATE dflash_common CUDA::cudart)
else()
set_source_files_properties(test/test_gpu_sampler_cuda.cpp PROPERTIES LANGUAGE HIP)
set_target_properties(test_gpu_sampler_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}")
target_include_directories(test_gpu_sampler_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat)
target_link_libraries(test_gpu_sampler_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET})
endif()
add_test(NAME gpu_sampler_cuda COMMAND test_gpu_sampler_cuda)
endif()
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kv_quant.cpp")
Expand Down
11 changes: 11 additions & 0 deletions server/hip_compat/cuda_runtime.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ using cudaEvent_t = hipEvent_t;
using cudaError_t = hipError_t;
using cudaMemcpyKind = hipMemcpyKind;
using cudaDeviceProp = hipDeviceProp_t;
using cudaPointerAttributes = hipPointerAttribute_t; // .type / .device match CUDA's

// Memcpy kind constants
#define cudaMemcpyHostToHost hipMemcpyHostToHost
Expand Down Expand Up @@ -77,6 +78,12 @@ using cudaDeviceProp = hipDeviceProp_t;
// Launch bounds
#define __launch_bounds__ __launch_bounds__

// Warp shuffle: CUDA's *_sync intrinsics take a 32-bit lane mask; HIP's require
// a 64-bit mask (wave64 ISAs) and static_assert against the implicit 32-bit
// promotion. The mask-less __shfl_xor covers the all-lanes-active warp
// reductions used here (RDNA wave32), matching rms_norm_hip.cu's idiom.
#define __shfl_xor_sync(mask, var, laneMask, width) __shfl_xor(var, laneMask, width)

// Stream capture status (added CUDA 10.0 — ROCm compat headers may omit this)
#define cudaStreamCaptureStatus hipStreamCaptureStatus
#define cudaStreamCaptureStatusNone hipStreamCaptureStatusNone
Expand All @@ -91,3 +98,7 @@ using cudaDeviceProp = hipDeviceProp_t;

// Device count
#define cudaGetDeviceCount hipGetDeviceCount

// Pointer attributes (used by the GPU sampler to run where device logits live)
#define cudaPointerGetAttributes hipPointerGetAttributes
#define cudaMemoryTypeDevice hipMemoryTypeDevice
Loading