perf(hip): compile DFlash GPU draft top-K kernel for HIP

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 |
 |---|---|---|

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

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

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
 

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.

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 = [](){

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
+//        <cuda_runtime.h> shim). Run: ./test_draft_topk_cuda (0 = pass).
 
 #include "../src/common/geometric_draft_topk_cuda.h"
 #include "../src/common/ddtree.h"