From 25a123d7a45c36b4c628bb2e43aaf72cb86a8df8 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Wed, 8 Jul 2026 19:59:15 +0530 Subject: [PATCH] perf(hip): compile DFlash GPU draft top-K kernel for HIP The GPU top-K + log-prob extraction (geometric_draft_topk_cuda.cu) was compiled only on CUDA, so DFLASH_GPU_DRAFT_TOPK was a no-op on ROCm and AMD paid a per-step vocab x n_tokens D2H + CPU heap extract every speculation step. The kernel body is plain arithmetic (no tensor cores, no CUDA-only intrinsics), so the same .cu compiles unchanged for HIP. Compile geometric_draft_topk_cuda.cu directly with LANGUAGE HIP through the hip_compat shim (the shared-.cu pattern already used by deepseek4_hc_cuda.cu), instead of adding a separate .hip.cu wrapper translation unit. Rename the backend guard macro DFLASH27B_HAVE_DRAFT_TOPK_CUDA -> DFLASH27B_HAVE_DRAFT_TOPK so the consumers (qwen35_dflash_target, test_dflash) are backend-neutral. test_draft_topk_cuda is now built on the HIP backend as well (CUDA spellings mapped by the same shim) and validated against the CPU reference on gfx1151 (Radeon 8060S, wave32, ROCm 6.4.4). Also qualify the README GPU-flag defaults to the server harness. --- README.md | 2 +- server/CMakeLists.txt | 36 ++++++++++++++++--- server/hip_compat/cuda_runtime.h | 9 +++++ server/src/common/geometric_draft_topk_cuda.h | 7 ++-- server/src/qwen35/qwen35_dflash_target.cpp | 2 +- server/test/test_dflash.cpp | 2 +- server/test/test_draft_topk_cuda.cpp | 12 ++++--- 7 files changed, 56 insertions(+), 14 deletions(-) diff --git a/README.md b/README.md index ed211ce1b..9cc8302f3 100644 --- a/README.md +++ b/README.md @@ -274,7 +274,7 @@ Requests that omit `temperature` use the model card's sampling (Qwen3.6: `temper **GPU draft top-K & verify-argmax (DFlash)** -The draft-token top-K extraction and the per-step verify argmax used to run on the CPU, each requiring a full `vocab × n_tokens` logits copy from device to host (D2H) every speculation step. These two env flags move both onto the GPU, reading the logits in place on the device buffer and skipping the bulk D2H. Both are **on by default** and only take effect on **CUDA builds** (the kernel is CUDA-only — on HIP/ROCm builds the flags are no-ops and the CPU path always runs). Each path validates its result and **falls back to the legacy CPU computation automatically** on any failure (e.g. an out-of-range index), so disabling them is only needed for debugging or A/B comparison. +The draft-token top-K extraction and the per-step verify argmax used to run on the CPU, each requiring a full `vocab × n_tokens` logits copy from device to host (D2H) every speculation step. These two env flags move both onto the GPU, reading the logits in place on the device buffer and skipping the bulk D2H. Both are **on by default in the server** (the `test_dflash` harness defaults `DFLASH_GPU_VERIFY_ARGMAX` to off, see the table below) and take effect on **both CUDA and HIP/ROCm builds**: the draft top-K uses a custom device kernel (`geometric_draft_topk_cuda.cu`, the same source compiled directly for HIP) and the verify argmax reads an in-graph `ggml_argmax` node, so neither depends on a CUDA-only path. Each path validates its result and **falls back to the legacy CPU computation automatically** on any failure (e.g. an out-of-range index), so disabling them is only needed for debugging or A/B comparison. | Env | Default | Effect | |---|---|---| diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 805d91c09..7e01f8c68 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -400,6 +400,19 @@ 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 draft top-K + log-prob kernel (DFlash). The shared body in + # geometric_draft_topk_cuda.cu compiles unchanged for HIP through the + # hip_compat shim, so it is added directly with LANGUAGE HIP + # (same shared-.cu pattern as deepseek4_hc_cuda.cu above — no separate HIP + # translation unit). This makes DFLASH_GPU_DRAFT_TOPK a real path on ROCm + # instead of a no-op, so AMD stops paying the per-step vocab x n_tokens D2H + + # CPU heap extract. The hip_compat include dir is already on dflash_common. + target_sources(dflash_common PRIVATE src/common/geometric_draft_topk_cuda.cu) + set_source_files_properties(src/common/geometric_draft_topk_cuda.cu + PROPERTIES LANGUAGE HIP) + # PUBLIC so test consumers (test_dflash / test_draft_topk_cuda) also take the + # GPU draft top-K path instead of the CPU fallback. + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) if(DFLASH27B_HIP_SM80_EQUIV) find_path(DFLASH27B_ROCWMMA_INCLUDE_DIR rocwmma/rocwmma.hpp HINTS "${_dflash_rocm_root}/include" /opt/rocm/include @@ -433,8 +446,9 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") src/flashprefill.cpp src/common/geometric_draft_topk_cuda.cu) # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro - # and take the GPU draft top-K path instead of the CPU fallback. - target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1) + # and take the GPU draft top-K path instead of the CPU fallback. Same macro + # name as the HIP branch above (backend-neutral). + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK=1) # GPU port of the sample_logits chain. Compiled in by default; the path is # then opted into at runtime via the DFLASH_GPU_SAMPLE env var. Turn the # whole thing off at configure time with -DDFLASH_GPU_SAMPLER=OFF. @@ -662,13 +676,27 @@ if(DFLASH27B_TESTS) target_link_libraries(test_rms_norm_hip PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) add_test(NAME rms_norm_hip COMMAND test_rms_norm_hip) endif() - # GPU draft top-K kernel vs CPU reference (extract_draft_topk). CUDA only: - # geometric_draft_topk_cuda.cu is compiled into dflash_common solely on the cuda backend. + # GPU draft top-K kernel vs CPU reference (extract_draft_topk). Built on both + # backends: geometric_draft_topk_cuda.cu is compiled into dflash_common on the + # cuda backend directly and on hip via LANGUAGE HIP + the hip_compat shim. if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp) target_include_directories(test_draft_topk_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart) add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda) + elseif(DFLASH27B_GPU_BACKEND STREQUAL "hip" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") + # HIP build of the same GPU-vs-CPU parity test. The test source uses CUDA + # spellings ( + cudaMalloc/cudaMemcpy/...); the hip_compat + # shim maps them onto HIP, exactly as the kernel TU does. Lets the draft + # top-K port be validated on AMD (gfx1100 / gfx1151 / gfx12xx). + add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp) + set_source_files_properties(test/test_draft_topk_cuda.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(test_draft_topk_cuda PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") + target_include_directories(test_draft_topk_cuda PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/src + ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) + target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) + 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. diff --git a/server/hip_compat/cuda_runtime.h b/server/hip_compat/cuda_runtime.h index dee4f0d92..4ca749899 100644 --- a/server/hip_compat/cuda_runtime.h +++ b/server/hip_compat/cuda_runtime.h @@ -91,3 +91,12 @@ using cudaDeviceProp = hipDeviceProp_t; // Device count #define cudaGetDeviceCount hipGetDeviceCount + +// Pointer attributes — used by geometric_draft_topk_cuda.cu to confirm the +// logits pointer is device memory and find which device it lives on. The +// hipPointerAttribute_t `.type` member (an hipMemoryType) is spelled `type` on +// ROCm >= 5.0; the AMD targets in scope (ROCm 6.4.x on the CI runner, 7.x per +// the maintainer's discrete-RDNA3 guidance) all use `type`. +using cudaPointerAttributes = hipPointerAttribute_t; +#define cudaPointerGetAttributes hipPointerGetAttributes +#define cudaMemoryTypeDevice hipMemoryTypeDevice diff --git a/server/src/common/geometric_draft_topk_cuda.h b/server/src/common/geometric_draft_topk_cuda.h index 19b410e68..b926dbefc 100644 --- a/server/src/common/geometric_draft_topk_cuda.h +++ b/server/src/common/geometric_draft_topk_cuda.h @@ -12,9 +12,12 @@ // out_log_probs = top-K scaled logits (desc), minus log_z // out_token_ids = matching vocab ids; ties broken toward the lower id // -// Returns false (caller must fall back to the CPU path) when CUDA is +// Returns false (caller must fall back to the CPU path) when the GPU runtime is // unavailable, the pointer is not device memory, K is out of range, or any -// CUDA call fails. Only compiled into CUDA builds; see CMakeLists.txt. +// device call fails. Compiled on both CUDA and HIP/ROCm builds — this same .cu +// is compiled directly with LANGUAGE HIP on ROCm, the cuda_runtime.h spellings +// mapped by the hip_compat shim; guarded by DFLASH27B_HAVE_DRAFT_TOPK. See +// CMakeLists.txt. #pragma once diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 4185b5d40..dc30ff7c4 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -692,7 +692,7 @@ bool Qwen35DFlashTarget::project_hidden_to_topk( top_log_probs.assign((size_t)n_tokens * K, 0.0f); top_token_ids.assign((size_t)n_tokens * K, 0); -#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA +#ifdef DFLASH27B_HAVE_DRAFT_TOPK // GPU path: top-K + logsumexp directly on the logits device buffer, skipping // the vocab×n_tokens D2H and the CPU heap extract. Falls back to the CPU path // on any failure. Escape hatch: DFLASH_GPU_DRAFT_TOPK=0. diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 407d23af1..7f5473995 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -3284,7 +3284,7 @@ int main(int argc, char ** argv) { } else { // DDTree K>1: need real log-probs for best-first tree scoring. bool topk_done = false; -#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA +#ifdef DFLASH27B_HAVE_DRAFT_TOPK // GPU path: top-K + logsumexp on the draft logits device buffer // (positions 1..q_len-1), no full-vocab D2H. Escape: DFLASH_GPU_DRAFT_TOPK=0. static const bool kGpuDraftTopk = [](){ diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp index ba847ccc8..bf734665e 100644 --- a/server/test/test_draft_topk_cuda.cpp +++ b/server/test/test_draft_topk_cuda.cpp @@ -1,8 +1,9 @@ // Correctness test for geometric_extract_draft_topk_cuda (GPU) vs extract_draft_topk (CPU). // -// The GPU kernel in src/common/geometric_draft_topk_cuda.cu is a drop-in replacement for -// the CPU top-K + online-logsumexp path in ddtree.cpp. This test feeds the same -// random logits to both and asserts the GPU results match the CPU reference: +// The GPU kernel in src/common/geometric_draft_topk_cuda.cu is a drop-in +// replacement for the CPU top-K + online-logsumexp path in ddtree.cpp. This test +// feeds the same random logits to both and asserts the GPU results match the CPU +// reference: // - token ids identical (rank by rank, per position) // - log-probs within a small bf16/float tolerance // @@ -11,8 +12,9 @@ // order the tied ids differently; we treat an id swap as OK when the matching // log-probs are equal within tolerance. // -// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS (CUDA only). -// Run: ./test_draft_topk_cuda (exit 0 = pass, non-zero = fail) +// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS for both the +// CUDA and HIP backends (the HIP build compiles this via the hip_compat +// shim). Run: ./test_draft_topk_cuda (0 = pass). #include "../src/common/geometric_draft_topk_cuda.h" #include "../src/common/ddtree.h"