Luce-Org · geometric · Jul 8, 2026 · Jul 9, 2026 · Jul 9, 2026 · Jul 9, 2026
diff --git 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
@@ -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.

diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu
@@ -6,6 +6,12 @@
 #        define STRIDED_ITERATOR_AVAILABLE
 #    endif
 using namespace cub;
+#elif defined(GGML_CUDA_USE_HIPCUB)
+// hipCUB exposes the CUB device-sort API (DeviceRadixSort / DeviceSegmentedSort /
+// DeviceSegmentedRadixSort) over rocPRIM. No strided-iterator / CCCL support, so
+// STRIDED_ITERATOR_AVAILABLE stays undefined and the init_offsets path is used.
+#    include <hipcub/hipcub.hpp>
+using namespace hipcub;
 #endif  // GGML_CUDA_USE_CUB
 
 static __global__ void init_indices(int * indices, const int ncols, const int nrows) {
@@ -26,7 +32,7 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr
 }
 #endif  // STRIDED_ITERATOR_AVAILABLE
 
-#ifdef GGML_CUDA_USE_CUB
+#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB)
 void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
                               const float *    x,
                               int *            dst,
@@ -138,7 +144,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
         }
     }
 }
-#endif  // GGML_CUDA_USE_CUB
+#endif  // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB
 
 // Bitonic sort implementation
 template<typename T>
@@ -148,6 +154,16 @@ static inline __device__ void ggml_cuda_swap(T & a, T & b) {
     b = tmp;
 }
 
+// Warp-local xor shuffle. For the bitonic inner stages where the stride j is
+// smaller than the wavefront width, both partners of a compare-exchange live in
+// the same wave, so the exchange can be done register-to-register via shuffle —
+// no shared-memory round-trip and no __syncthreads() barrier.
+#if defined(GGML_USE_HIP) || defined(__HIP_PLATFORM_AMD__)
+#    define GGML_ARGSORT_SHFL_XOR(v, mask) __shfl_xor((v), (mask))
+#else
+#    define GGML_ARGSORT_SHFL_XOR(v, mask) __shfl_xor_sync(0xffffffffu, (v), (mask))
+#endif
+
 template<ggml_sort_order order>
 static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) {
     // bitonic sort
@@ -159,37 +175,95 @@ static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int n
     }
 
     const float * x_row = x + row * ncols;
-    extern __shared__ int dst_row[];
-
-    // initialize indices
+    // Shared layout: [ncols_pad] indices followed by [ncols_pad] cached key
+    // values. The bitonic network re-reads two keys per comparison across
+    // ~log2(ncols_pad)^2 stages; caching the value that travels with each
+    // index removes the repeated indirect global gathers from the inner loop.
+    extern __shared__ int smem[];
+    int *   dst_row = smem;
+    float * val_row = (float *) (smem + ncols_pad);
+
+    // initialize indices and cache the key each index points at (padding lanes
+    // keep an untouched value; they are handled by the index-based checks below)
     dst_row[col] = col;
+    if (col < ncols) {
+        val_row[col] = x_row[col];
+    }
 
     __syncthreads();
 
     for (int k = 2; k <= ncols_pad; k *= 2) {
-        for (int j = k / 2; j > 0; j /= 2) {
+        int j = k / 2;
+
+        // Cross-wave stages (stride >= warpSize): partners live in different
+        // waves, so the exchange must go through shared memory with a full
+        // block barrier between stages.
+        for (; j >= warpSize; j /= 2) {
             int ixj = col ^ j;
             if (ixj > col) {
                 if ((col & k) == 0) {
                     if (dst_row[col] >= ncols ||
                         (dst_row[ixj] < ncols && (order == GGML_SORT_ORDER_ASC ?
-                            x_row[dst_row[col]] > x_row[dst_row[ixj]] :
-                            x_row[dst_row[col]] < x_row[dst_row[ixj]]))
+                            val_row[col] > val_row[ixj] :
+                            val_row[col] < val_row[ixj]))
                     ) {
                         ggml_cuda_swap(dst_row[col], dst_row[ixj]);
+                        ggml_cuda_swap(val_row[col], val_row[ixj]);
                     }
                 } else {
                     if (dst_row[ixj] >= ncols ||
                         (dst_row[col] < ncols && (order == GGML_SORT_ORDER_ASC ?
-                            x_row[dst_row[col]] < x_row[dst_row[ixj]] :
-                            x_row[dst_row[col]] > x_row[dst_row[ixj]]))
+                            val_row[col] < val_row[ixj] :
+                            val_row[col] > val_row[ixj]))
                     ) {
                         ggml_cuda_swap(dst_row[col], dst_row[ixj]);
+                        ggml_cuda_swap(val_row[col], val_row[ixj]);
                     }
                 }
             }
             __syncthreads();
         }
+
+        // Intra-wave tail (stride < warpSize): partner is in the same wave.
+        // Pull (idx, val) into registers and drive the remaining stages with
+        // xor shuffles — no LDS traffic, no barriers. Both lanes of a pair
+        // reconstruct the same (low, high) view and compute an identical swap
+        // decision, so the exchange stays consistent. Semantics match the
+        // shared-memory path above byte-for-byte (col == low lane, ixj == high).
+        if (j > 0) {
+            int   my_idx = dst_row[col];
+            float my_val = val_row[col];
+            for (; j > 0; j /= 2) {
+                const float p_val = GGML_ARGSORT_SHFL_XOR(my_val, j);
+                const int   p_idx = GGML_ARGSORT_SHFL_XOR(my_idx, j);
+
+                const bool  low      = (col & j) == 0;
+                const int   low_idx  = low ? my_idx : p_idx;
+                const int   high_idx = low ? p_idx : my_idx;
+                const float low_val  = low ? my_val : p_val;
+                const float high_val = low ? p_val : my_val;
+
+                bool swap;
+                if ((col & k) == 0) {
+                    swap = low_idx >= ncols ||
+                           (high_idx < ncols && (order == GGML_SORT_ORDER_ASC ?
+                                low_val > high_val :
+                                low_val < high_val));
+                } else {
+                    swap = high_idx >= ncols ||
+                           (low_idx < ncols && (order == GGML_SORT_ORDER_ASC ?
+                                low_val < high_val :
+                                low_val > high_val));
+                }
+                if (swap) {
+                    my_idx = p_idx;
+                    my_val = p_val;
+                }
+            }
+            dst_row[col] = my_idx;
+            val_row[col] = my_val;
+            __syncthreads();
+        }
     }
 
     // copy the result to dst without the padding
@@ -217,7 +291,8 @@ void argsort_f32_i32_cuda_bitonic(const float *   x,
 
     const dim3 block_dims(ncols_pad, 1, 1);
     const dim3 block_nums(nrows, 1, 1);
-    const size_t shared_mem = ncols_pad * sizeof(int);
+    // indices (int) + cached key values (float), both ncols_pad entries
+    const size_t shared_mem = ncols_pad * (sizeof(int) + sizeof(float));
 
     // FIXME: this limit could be raised by ~2-4x on Ampere or newer
     GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb);
@@ -248,7 +323,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
 
     enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0];
 
-#ifdef GGML_CUDA_USE_CUB
+#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB)
     const int    ncols_pad      = next_power_of_2(ncols);
     const size_t shared_mem     = ncols_pad * sizeof(int);
     const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb;

diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh
@@ -2,15 +2,15 @@
 
 void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
-#ifdef GGML_CUDA_USE_CUB
+#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB)
 void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool,
                               const float *    x,
                               int *            dst,
                               const int        ncols,
                               const int        nrows,
                               ggml_sort_order  order,
                               cudaStream_t     stream);
-#endif  // GGML_CUDA_USE_CUB
+#endif  // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB
 void argsort_f32_i32_cuda_bitonic(const float *   x,
                                   int *           dst,
                                   const int       ncols,

diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh
@@ -107,6 +107,13 @@
 #    define GGML_CUDA_USE_CUB
 #endif  // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070
 
+// HIP has no CUB, but rocPRIM ships a drop-in CUB-API layer via hipCUB. Enable it
+// so argsort / top-k support ncols > 1024 (the single-block bitonic path caps at
+// 1024 threads/block); without it those ops fall back to the CPU reference on ROCm.
+#if defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+#    define GGML_CUDA_USE_HIPCUB
+#endif  // defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+
 #ifdef __CUDA_ARCH_LIST__
 constexpr bool ggml_cuda_has_arch_impl(int) {
     return false;

diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -5276,10 +5276,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
             return ggml_is_contiguous_rows(op->src[0]);
         case GGML_OP_TOP_K:
         case GGML_OP_ARGSORT:
-#ifndef GGML_CUDA_USE_CUB
-            return op->src[0]->ne[0] <= 1024;
-#else
+#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB)
             return true;
+#else
+            // No device-wide sort backend: the single-block bitonic path caps at
+            // 1024 threads/block, so only ncols <= 1024 is supported on GPU.
+            return op->src[0]->ne[0] <= 1024;
 #endif
         case GGML_OP_SUM_ROWS:
         case GGML_OP_MEAN: