perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend

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
+    # <cuda_runtime.h> / 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

server/test/test_draft_topk_cuda.cpp

@@ -19,9 +19,11 @@
 
 #include <cuda_runtime.h>
 
+#include <chrono>
 #include <cmath>
 #include <cstdint>
 #include <cstdio>
+#include <cstdlib>
 #include <random>
 #include <vector>
 
@@ -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<double, std::micro>(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<float> 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<float> 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<float>   lp(n_out);
+        std::vector<int32_t> 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[] = {