From b5f3f83abf6dba2e669bc0f57d14eeb292ab9dfe Mon Sep 17 00:00:00 2001 From: deano Date: Tue, 7 Jul 2026 19:50:39 +0300 Subject: [PATCH] perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The split-K draft top-K + logsumexp kernel (geometric_draft_topk_cuda.cu, the spec-decode draft path) added alongside the sampler in #478 was wired for the CUDA backend only, so on AMD the draft top-K fell back to the CPU reference (extract_draft_topk). This enables the GPU kernel on HIP so AMD (gfx1100/gfx1151/gfx1201) uses it too. * CMakeLists: compile geometric_draft_topk_cuda.cu as LANGUAGE HIP under the hip backend and define DFLASH27B_HAVE_DRAFT_TOPK_CUDA so callers take the GPU path; build test_draft_topk_cuda on hip too (CUDA-spelled source compiles as HIP via hip_compat/). Stacked on the sampler HIP PR, which adds the shared hip_compat pointer-attribute shims this kernel also uses. * test_draft_topk_cuda: add a DFLASH_TOPK_BENCH-gated CPU-vs-GPU microbench (mirrors test_gpu_sampler_cuda's DFLASH_SAMPLER_BENCH) so the win is measurable on any backend. The kernel needs no source or shim changes: it has NO warp-level intrinsics (a plain kBlock-wide __syncthreads reduction, so no wavefront-size assumption), and its only runtime calls (cudaMalloc / cudaMemcpy / cudaPointer*) resolve through the existing shim. Validated on the AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1): * test_draft_topk_cuda: 10/10 pass — top-K ids bit-for-bit vs the CPU reference (id_mismatch=0) across Qwen vocab (151936) and n/K/temp shapes, plus the K>kMaxK CPU-fallback contract. max log-prob err <= 3.2e-4. * CPU-vs-GPU microbench (DFLASH_TOPK_BENCH=1, K=8, 500 iters, logits resident on device = the decode steady state; GPU time = launch+sync+D2H of the n*K result): n=15 vocab=151936 CPU 695.2 us -> GPU 164.9 us = 4.22x n=1 vocab=151936 CPU 594.4 us -> GPU 153.6 us = 3.87x --- server/CMakeLists.txt | 27 ++++++++++-- server/test/test_draft_topk_cuda.cpp | 62 ++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index efa758f61..78fc83ca8 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -415,6 +415,16 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") # also compile their GPU dispatch branch. target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_GPU_SAMPLER=1) endif() + # Split-K GPU draft top-K + logsumexp (spec-decode draft path). Shares the + # CUDA source verbatim — compiled as HIP (no warp intrinsics; a plain + # kBlock-wide __syncthreads reduction, so no wavefront-size assumption). Its + # / cudaPointer* calls resolve through the hip_compat/ shim. + 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 consumers (e.g. test_dflash) also take the GPU draft top-K path + # instead of the CPU fallback. + target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1) if(DFLASH27B_HIP_SM80_EQUIV) find_path(DFLASH27B_ROCWMMA_INCLUDE_DIR rocwmma/rocwmma.hpp HINTS "${_dflash_rocm_root}/include" /opt/rocm/include @@ -663,12 +673,21 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/hip_compat) target_link_libraries(test_flashprefill_kernels PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) 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. - if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_draft_topk_cuda.cpp") + # GPU draft top-K kernel vs CPU reference (extract_draft_topk). Built on + # whichever backend compiled geometric_draft_topk_cuda.cu into dflash_common + # (cuda as CUDA, hip as HIP). On hip the CUDA-spelled test compiles as HIP via + # the hip_compat/ shim, like test_gpu_sampler_cuda / test_flashprefill_kernels. + if(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) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart) + else() + 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}/hip_compat) + target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common ${DFLASH27B_GGML_BACKEND_TARGET}) + endif() add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda) endif() # GPU port of the sample_logits chain vs the CPU reference. Built on whichever diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp index ba847ccc8..825c70495 100644 --- a/server/test/test_draft_topk_cuda.cpp +++ b/server/test/test_draft_topk_cuda.cpp @@ -19,9 +19,11 @@ #include +#include #include #include #include +#include #include #include @@ -115,6 +117,64 @@ bool run_case(const Case & c, unsigned seed) { return pass; } +// Per-call CPU-vs-GPU timing (gated by env DFLASH_TOPK_BENCH=1). Isolates the +// whole extract_draft_topk call the ddtree draft path makes each speculative +// step: the CPU reference vs the GPU kernel (launch + sync + D2H of the small +// n*K result). Logits stay resident on the device across iters — the same as +// decode, where the draft logits are produced on the GPU — so this measures the +// steady-state per-step cost, not a cold launch or a full-vocab H2D upload. +void draft_topk_microbench() { + auto now = []{ return std::chrono::steady_clock::now(); }; + auto us = [](auto a, auto b){ return std::chrono::duration(b - a).count(); }; + + struct BenchCase { int n, vocab, K; }; + const BenchCase cases[] = { + {15, 151936, 8}, // realistic ddtree draft batch (~15 positions), Qwen3 vocab + {1, 151936, 8}, // single draft position + }; + const int iters = 500; + + std::mt19937 rng(1234); + std::normal_distribution dist(0.f, 4.f); + + printf("[topk-bench] iters=%d (per call; >1.0x = GPU faster)\n", iters); + for (const auto & c : cases) { + const size_t n_logits = (size_t)c.n * c.vocab; + const size_t n_out = (size_t)c.n * c.K; + std::vector h_logits(n_logits); + for (auto & x : h_logits) x = dist(rng); + + float * d_logits = nullptr; + if (cudaMalloc(&d_logits, n_logits * sizeof(float)) != cudaSuccess) { + printf(" cudaMalloc failed\n"); return; + } + cudaMemcpy(d_logits, h_logits.data(), n_logits * sizeof(float), cudaMemcpyHostToDevice); + + std::vector lp(n_out); + std::vector ids(n_out); + + // Warmup both paths — the first GPU launch pays kernel-module load. + for (int i = 0; i < 20; i++) { + extract_draft_topk(h_logits.data(), c.n, c.vocab, c.K, lp.data(), ids.data(), 1.0f); + geometric_extract_draft_topk_cuda(d_logits, c.n, c.vocab, c.K, lp.data(), ids.data(), 1.0f); + } + + const auto t0 = now(); + for (int i = 0; i < iters; i++) + extract_draft_topk(h_logits.data(), c.n, c.vocab, c.K, lp.data(), ids.data(), 1.0f); + const auto t1 = now(); + for (int i = 0; i < iters; i++) + geometric_extract_draft_topk_cuda(d_logits, c.n, c.vocab, c.K, lp.data(), ids.data(), 1.0f); + const auto t2 = now(); + cudaFree(d_logits); + + const double cpu_us = us(t0, t1) / iters; + const double gpu_us = us(t1, t2) / iters; + printf(" n=%-3d vocab=%d K=%d CPU %8.2f us | GPU %8.2f us (%.2fx)\n", + c.n, c.vocab, c.K, cpu_us, gpu_us, gpu_us > 0 ? cpu_us / gpu_us : 0.0); + } +} + } // namespace int main() { @@ -124,6 +184,8 @@ int main() { return 0; } + if (std::getenv("DFLASH_TOPK_BENCH")) { draft_topk_microbench(); return 0; } + // The kernel supports K up to kMaxK (=8 in geometric_draft_topk_cuda.cu); larger K is // handled by a documented CPU fallback (returns false), checked separately. const Case cases[] = {