From 0693aac53feb7e6972b4590176a90830f29c04c9 Mon Sep 17 00:00:00 2001 From: deano Date: Tue, 7 Jul 2026 19:34:29 +0300 Subject: [PATCH] perf(hip): enable the fused GPU sampler on the HIP backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fused GPU sampler (sample_logits chain: rep/freq/pres penalties -> softmax(temp) -> greedy/multinomial draw) added in #478 — its kernel (geometric_sampler_cuda.cu), header, the sampler.cpp / qwen35_backend.cpp dispatch, and test_gpu_sampler_cuda — was wired for the CUDA backend only. This enables it on HIP so AMD (gfx1100/gfx1151/gfx1201) samples on-device instead of falling back to the CPU chain. No kernel or dispatch changes; the whole diff is the HIP build wiring + shim: * CMakeLists: compile geometric_sampler_cuda.cu as LANGUAGE HIP under the hip backend (option DFLASH_GPU_SAMPLER, ON) and define DFLASH27B_HAVE_GPU_SAMPLER so the backends take the GPU dispatch branch; build test_gpu_sampler_cuda on hip too (the CUDA-spelled test compiles as HIP via hip_compat/, like test_flashprefill_kernels). * hip_compat/cuda_runtime.h: add cudaPointerAttributes / cudaPointerGetAttributes / cudaMemoryTypeDevice, and map CUDA's 32-bit-mask __shfl_xor_sync onto HIP's mask-less __shfl_xor (HIP's _sync form static_asserts against the implicit 32->64-bit mask promotion; the all-lanes-active wave32 reductions here don't need the mask, matching rms_norm_hip.cu). The kernel's 32-lane warp assumption holds on RDNA wave32. Validated on the AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1): * test_gpu_sampler_cuda: 22/22 pass (greedy/penalties match CPU argmax bit-for-bit; temp + top_p-assist draws match analytic softmax/nucleus). * per-call microbench (vocab=151936, 1000 iters), GPU sampling from the device logits tensor (the production decode path, logits_on_device=true): greedy (temp=0) 54.1 us vs CPU 183.4 us = 3.39x temp=0.8 (full vocab) 92.1 us vs CPU 218.8 us = 2.38x temp=0.8 rep_pen=1.2 181.4 us vs CPU 303.7 us = 1.67x (GPU+H2D is ~1.0x — copying 152k logits H2D per token eats the win, which is why the decode loop samples in place.) * dflash_server e2e (Qwen3.6-27B Q4_K_M + draft, temp=0.8): DFLASH_GPU_SAMPLE=1 vs =0 both coherent, no regression (21.9 vs 21.7 tok/s), no HIP errors. --- server/CMakeLists.txt | 32 ++++++++++++++++++++++++++++---- server/hip_compat/cuda_runtime.h | 11 +++++++++++ 2 files changed, 39 insertions(+), 4 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index b693e1041..efa758f61 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -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 + # 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 @@ -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") diff --git a/server/hip_compat/cuda_runtime.h b/server/hip_compat/cuda_runtime.h index dee4f0d92..e0d33315c 100644 --- a/server/hip_compat/cuda_runtime.h +++ b/server/hip_compat/cuda_runtime.h @@ -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 @@ -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 @@ -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