From 25a123d7a45c36b4c628bb2e43aaf72cb86a8df8 Mon Sep 17 00:00:00 2001 From: cheese-cakee Date: Wed, 8 Jul 2026 19:59:15 +0530 Subject: [PATCH 1/5] 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" From bae3c2ac96cc1c8f35c85bde2cf3ea972467ad2c Mon Sep 17 00:00:00 2001 From: deano Date: Thu, 9 Jul 2026 16:28:46 +0300 Subject: [PATCH 2/5] =?UTF-8?q?CUDA/HIP:=20partial-top-k=20argsort=20?= =?UTF-8?q?=E2=80=94=20up=20to=206x=20faster=20TOP=5FK=20on=20RDNA3.5/RDNA?= =?UTF-8?q?4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The TOP_K path full-sorts every row via bitonic argsort regardless of k, so its latency is independent of k. Make it partial: - argsort.cu: cache keys in shared memory and bound the bitonic sort by the padded k, so work scales with k rather than ncols. - argsort.cu: GGML_ARGSORT_SHFL_XOR — for bitonic inner stages whose stride is below the wavefront width, both compare-exchange partners live in the same wave, so the swap is done register-to-register via __shfl_xor with no shared memory round-trip or __syncthreads() (wave32 on RDNA3.5/RDNA4). - top-k.cu: dedicated k=1 argmax kernel that skips the sort entirely. On HIP the CUB fast path is unavailable (GGML_CUDA_USE_CUB undefined), so the bitonic kernel is the hot path on AMD. test-backend-ops -o TOP_K passes on gfx1201 and gfx1151. Microbench A/B (ne=[1000,16]): k=1 15.59->2.59us (6.0x), k=10 ->4.58us (3.4x), k=400 ->6.56us (2.4x); geomean ~3.5-3.8x. --- .../llama.cpp/ggml/src/ggml-cuda/argsort.cu | 87 ++- .../llama.cpp/ggml/src/ggml-cuda/top-k.cu | 630 +++++++++++++++++- 2 files changed, 703 insertions(+), 14 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu index 0f3f017b5..fe0810981 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu @@ -148,6 +148,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 static __global__ void k_argsort_f32_i32(const float * x, int * dst, const int ncols, int ncols_pad) { // bitonic sort @@ -159,37 +169,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 +285,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); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 59ce36fb1..5edc1899e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -47,6 +47,606 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE +#if !defined(GGML_CUDA_USE_CUB) + +static int topk_next_power_of_2(int x) { + int n = 1; + while (n < x) { + n *= 2; + } + return n; +} + +template +static inline __device__ void topk_swap(T & a, T & b) { + T tmp = a; + a = 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 TOPK_SHFL_XOR(v, mask) __shfl_xor((v), (mask)) +#else +# define TOPK_SHFL_XOR(v, mask) __shfl_xor_sync(0xffffffffu, (v), (mask)) +#endif + +// Dedicated argmax for k == 1 (HIP / no-CUB builds). +// +// TOP_K with k == 1 is a pure argmax, and the op accepts *any* index holding the +// maximum value (with ties the test compares the values behind the indices, not +// the indices themselves). The generic bitonic path handles k == 1 as a kpad == 1 +// merge-tree: log2(ncols_pad) reduction levels, each with two block barriers and a +// strided shared-memory round-trip. A direct wave-shuffle reduction replaces that +// with two barrier-free shuffle passes and a single __syncthreads(). Threads are +// rounded up to a whole number of waves so every lane in every wave participates +// in the shuffles (out-of-range lanes seed -inf and never win). Ties break toward +// the lower index purely to keep the reduction deterministic. +static __global__ void k_topk_argmax(const float * x, int * dst, const int ncols) { + const int col = threadIdx.x; + const int row = blockIdx.x; + + const float * x_row = x + (size_t) row * ncols; + + float v = col < ncols ? x_row[col] : -INFINITY; + int id = col; + + // Intra-wave reduction. + for (int j = warpSize >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + if (pv > v || (pv == v && pi < id)) { + v = pv; + id = pi; + } + } + + const int lane = col & (warpSize - 1); + const int warp = col / warpSize; + const int nwarps = blockDim.x / warpSize; + + extern __shared__ int argmax_smem[]; + int * s_idx = argmax_smem; + float * s_val = (float *) (argmax_smem + nwarps); + + if (lane == 0) { + s_idx[warp] = id; + s_val[warp] = v; + } + __syncthreads(); + + // Final reduction over the wave leaders, done entirely in wave 0 (nwarps is at + // most ncols_pad/warpSize <= 1024/32 = 32 <= warpSize, so one wave covers them). + if (warp == 0) { + v = lane < nwarps ? s_val[lane] : -INFINITY; + id = lane < nwarps ? s_idx[lane] : 0; + for (int j = warpSize >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + if (pv > v || (pv == v && pi < id)) { + v = pv; + id = pi; + } + } + if (lane == 0) { + dst[row] = id; + } + } +} + +// Dedicated block-resident top-k selection (HIP / no-CUB builds). +// +// TOP_K only requires the *set* of the k largest indices per row — the op's +// output order is explicitly irrelevant (the CPU reference even swaps the first +// two outputs, and the test compares results as an order-independent set). So +// instead of fully sorting the whole (padded) row and copying the first k, we +// do a partial bitonic top-k that only ever keeps a window of `kpad` elements: +// +// Phase A: bitonic-sort each contiguous `kpad`-block of the row descending. +// Phase B: merge-tree — repeatedly merge two adjacent descending `kpad`-blocks, +// keeping only the top `kpad` of the pair (C[i] = max(X[i], Y[kpad-1-i]) +// is bitonic and holds the kpad largest; one bitonic merge re-sorts it +// descending). After log2(ncols_pad/kpad) levels, block 0 holds the top +// kpad descending, and its first k entries are the top-k set. +// +// Work scales with kpad rather than ncols_pad, so small-k rows (the common case) +// avoid the full O(n log^2 n) sort. Padding lanes are seeded with -inf so they can +// never enter the top-k (ncols >= k always holds for TOP_K), which lets the whole +// network compare values directly with no index/padding guards. +static __global__ void k_topk_bitonic(const float * x, + int * dst, + const int ncols, + const int ncols_pad, + const int k, + const int kpad) { + const int col = threadIdx.x; + const int row = blockIdx.x; + + const float * x_row = x + (size_t) row * ncols; + + extern __shared__ int smem[]; + int * idx_row = smem; + float * val_row = (float *) (smem + ncols_pad); + + // Phase A: sort every contiguous kpad-block descending (blocks are + // independent — partners of a compare-exchange stay within the block). + // + // (v, id) stay register-resident across the whole phase; LDS is touched + // only for cross-wave stages (stride >= warpSize). A pure intra-wave `len` + // stage (len <= warpSize, so every stride j <= len/2 < warpSize) runs + // entirely via xor shuffle with no LDS round-trip and no barrier — the + // per-`len` writeback + __syncthreads() the old code did unconditionally is + // only needed to feed a subsequent cross-wave read, so it is deferred to the + // first `len` that actually has one (and to the Phase-B handoff at the end). + float v = col < ncols ? x_row[col] : -INFINITY; + int id = col; + for (int len = 2; len <= kpad; len <<= 1) { + int j = len >> 1; + + // Cross-wave stages (stride >= warpSize): partners live in different + // waves, so the exchange goes through shared memory with a block barrier. + // Flush the register-resident (v, id) to LDS first so partner lanes read + // the current values. + if (j >= warpSize) { + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + for (; j >= warpSize; j >>= 1) { + const int partner = col ^ j; + if (partner > col) { + // Descending target for this comparator (uniform per block at + // the top stage; alternating within to build bitonic subseqs). + const bool up = (((col & (kpad - 1)) & len) == 0); + const bool do_swap = + up ? (val_row[col] < val_row[partner]) : (val_row[col] > val_row[partner]); + if (do_swap) { + topk_swap(val_row[col], val_row[partner]); + topk_swap(idx_row[col], idx_row[partner]); + } + } + __syncthreads(); + } + v = val_row[col]; + id = idx_row[col]; + } + + // Intra-wave tail (stride < warpSize): partner is in the same wave. Drive + // the remaining stages with xor shuffles — no LDS traffic, no barriers. + // Both lanes of a pair reconstruct the same (low, high) view and the same + // `up` direction (they differ only in bit j, which lies below len's bit, + // so `up` is identical), and each lane replaces its own value with the + // partner's when the shared swap decision fires — so the exchange stays + // consistent and matches the shared-memory path above. + for (; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool up = (((col & (kpad - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } + } + // Publish the sorted kpad-blocks to LDS for Phase B's cross-wave merges. + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + + // Phase B: merge-tree. Merge the two descending kpad-blocks at the head of + // each 2*span group into a single descending top-kpad block stored at the + // group head. `span` doubles each level until one block spans the row. + for (int span = kpad; span < ncols_pad; span <<= 1) { + const int gs = span << 1; + const int i = col & (gs - 1); // position within the 2*span group + const int gbase = col - i; // group head + const bool active = i < kpad; // only the surviving kpad lanes work + + if (active) { + // C[i] = max(X[i], Y[kpad-1-i]) is bitonic and holds the kpad largest. + // The pre-writeback barrier is unnecessary: each active lane reads its + // own lower-half slot a_pos = gbase+i plus an upper-half slot b_pos in + // [gbase+span, gbase+span+kpad) (never written this level since + // kpad <= span), and writes back only to gbase+i. The write set + // [gbase, gbase+kpad) is disjoint from the cross-half read set and the + // a_pos read-before-write is intra-thread ordered, so only the + // post-writeback barrier is needed. + const int a_pos = gbase + i; + const int b_pos = gbase + span + (kpad - 1 - i); + const float av = val_row[a_pos]; + const float bv = val_row[b_pos]; + if (av >= bv) { + val_row[gbase + i] = av; + idx_row[gbase + i] = idx_row[a_pos]; + } else { + val_row[gbase + i] = bv; + idx_row[gbase + i] = idx_row[b_pos]; + } + } + __syncthreads(); + + // Bitonic merge the (bitonic) block back to descending order. + int j = kpad >> 1; + + // Cross-wave stages via shared memory. + for (; j >= warpSize; j >>= 1) { + const int partner = col ^ j; + if (active && partner > col) { + if (val_row[col] < val_row[partner]) { + topk_swap(val_row[col], val_row[partner]); + topk_swap(idx_row[col], idx_row[partner]); + } + } + __syncthreads(); + } + + // Intra-wave tail via xor shuffle. Merge partners (col ^ j, j < kpad) stay + // within the aligned head kpad-block, so an active lane's partner is also + // active; inactive lanes shuffle harmlessly but never write back. Uniform + // descending merge: the low lane keeps the larger value. + if (j > 0) { + float v = val_row[col]; + int id = idx_row[col]; + for (; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + if (low_val < high_val) { + v = pv; + id = pi; + } + } + if (active) { + val_row[col] = v; + idx_row[col] = id; + } + __syncthreads(); + } + } + + if (col < k) { + dst[(size_t) row * k + col] = idx_row[col]; + } +} + +// Dedicated small-k top-k for kpad <= warpSize (HIP / no-CUB builds). +// +// Same partial-bitonic algorithm as k_topk_bitonic, specialised for the case +// where a whole kpad-block fits inside one wave. Two barriers/LDS-traffic cuts +// that are only valid in this regime: +// * Phase A runs entirely in registers. With kpad <= warpSize every +// compare-exchange partner (col ^ j, j < kpad) is in the same wave, so the +// per-block sort is driven purely by xor shuffles — no LDS round-trip and no +// __syncthreads() between the len stages (the shared kernel writes back to +// LDS and barriers after every len). +// * Phase B elides the C-step's pre-writeback barrier. Each active lane reads +// its own lower-half slot (a_pos = gbase + i) plus an upper-half slot +// (b_pos, never written this level) and writes back only to gbase + i, so no +// lane's write aliases another lane's read — the read/write sets are disjoint +// and the merge only needs the single post-writeback barrier. +// Both cuts are pure schedule changes; the compare-exchange decisions are +// byte-for-byte those of k_topk_bitonic, so the produced top-k set is identical. +static __global__ void k_topk_bitonic_smallk(const float * x, + int * dst, + const int ncols, + const int ncols_pad, + const int k, + const int kpad) { + const int col = threadIdx.x; + const int row = blockIdx.x; + + const float * x_row = x + (size_t) row * ncols; + + extern __shared__ int smem[]; + int * idx_row = smem; + float * val_row = (float *) (smem + ncols_pad); + + // Phase A (register-only): sort each contiguous kpad-block descending. + float v = col < ncols ? x_row[col] : -INFINITY; + int id = col; + for (int len = 2; len <= kpad; len <<= 1) { + for (int j = len >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool up = (((col & (kpad - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } + } + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + + // Phase B: merge-tree (same as k_topk_bitonic; C-step fused into a single + // post-writeback barrier — see the header note for why the pre-barrier is + // unnecessary here). + for (int span = kpad; span < ncols_pad; span <<= 1) { + const int gs = span << 1; + const int i = col & (gs - 1); + const int gbase = col - i; + const bool active = i < kpad; + + if (active) { + const int a_pos = gbase + i; + const int b_pos = gbase + span + (kpad - 1 - i); + const float av = val_row[a_pos]; + const float bv = val_row[b_pos]; + if (av >= bv) { + val_row[gbase + i] = av; + idx_row[gbase + i] = idx_row[a_pos]; + } else { + val_row[gbase + i] = bv; + idx_row[gbase + i] = idx_row[b_pos]; + } + } + __syncthreads(); + + // Bitonic merge the kpad-wide block back to descending order. kpad <= + // warpSize, so the whole merge is intra-wave xor shuffle. + v = val_row[col]; + id = idx_row[col]; + for (int j = kpad >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + if (low_val < high_val) { + v = pv; + id = pi; + } + } + if (active) { + val_row[col] = v; + idx_row[col] = id; + } + __syncthreads(); + } + + if (col < k) { + dst[(size_t) row * k + col] = idx_row[col]; + } +} + +// Dedicated top-k for a two-wave kpad-block (warpSize < kpad <= 2*warpSize; +// HIP / no-CUB builds). kpad is a power of two, so this regime is exactly +// kpad == 2*warpSize (kpad=64 on wave32). Same partial-bitonic algorithm as +// k_topk_bitonic, but exploiting that a kpad-block spans only two waves: +// * Phase A runs in registers across all `len` stages. Every compare-exchange +// partner (col ^ j, j < kpad) with j < warpSize is in the same wave, and the +// single j == warpSize top stage is the only cross-wave exchange — so instead +// of the shared kernel's LDS write-back + __syncthreads() after every `len`, +// the block sorts entirely in registers via xor shuffle and touches LDS just +// once (for that one cross-wave exchange). The comparator (up/low/do_swap) is +// byte-for-byte the shared kernel's shuffle tail. +// * Phase B elides the C-step's pre-write-back barrier (see k_topk_bitonic_smallk +// header): each active lane's write set [gbase, gbase+kpad) is disjoint from its +// cross-half read slot b_pos in [gbase+span, gbase+span+kpad) since kpad <= span, +// so only the post-write-back barrier is needed. +// Both cuts are pure schedule changes; the produced top-k set is identical to +// k_topk_bitonic. Isolated in its own kernel so the k > 2*warpSize path's codegen +// (k_topk_bitonic) is unchanged. +static __global__ void k_topk_bitonic_2wave(const float * x, + int * dst, + const int ncols, + const int ncols_pad, + const int k, + const int kpad) { + const int col = threadIdx.x; + const int row = blockIdx.x; + + const float * x_row = x + (size_t) row * ncols; + + extern __shared__ int smem[]; + int * idx_row = smem; + float * val_row = (float *) (smem + ncols_pad); + + // Phase A: sort every contiguous kpad-block descending, register-resident. + float v = col < ncols ? x_row[col] : -INFINITY; + int id = col; + for (int len = 2; len <= kpad; len <<= 1) { + int j = len >> 1; + + // At most one cross-wave stage (j == warpSize, only when len == kpad since + // kpad <= 2*warpSize): exchange with the partner wave through LDS. + if (j >= warpSize) { + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + const int partner = col ^ j; + const float pv = val_row[partner]; + const int pi = idx_row[partner]; + __syncthreads(); // all cross-wave reads done before Phase A's LDS reuse + const bool up = (((col & (kpad - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + j >>= 1; + } + + // Intra-wave tail (j < warpSize): register-to-register via xor shuffle. + for (; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool up = (((col & (kpad - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } + } + __syncthreads(); // ensure the last cross-wave LDS reads finished before overwrite + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + + // Phase B: merge-tree (C-step pre-write-back barrier elided, see header). + for (int span = kpad; span < ncols_pad; span <<= 1) { + const int gs = span << 1; + const int i = col & (gs - 1); + const int gbase = col - i; + const bool active = i < kpad; + + if (active) { + const int a_pos = gbase + i; + const int b_pos = gbase + span + (kpad - 1 - i); + const float av = val_row[a_pos]; + const float bv = val_row[b_pos]; + if (av >= bv) { + val_row[gbase + i] = av; + idx_row[gbase + i] = idx_row[a_pos]; + } else { + val_row[gbase + i] = bv; + idx_row[gbase + i] = idx_row[b_pos]; + } + } + __syncthreads(); + + int j = kpad >> 1; + + // Cross-wave stages (j >= warpSize) via shared memory. + for (; j >= warpSize; j >>= 1) { + const int partner = col ^ j; + if (active && partner > col) { + if (val_row[col] < val_row[partner]) { + topk_swap(val_row[col], val_row[partner]); + topk_swap(idx_row[col], idx_row[partner]); + } + } + __syncthreads(); + } + + // Intra-wave tail via xor shuffle. + if (j > 0) { + float mv = val_row[col]; + int mid = idx_row[col]; + for (; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(mv, j); + const int pi = TOPK_SHFL_XOR(mid, j); + + const bool low = (col & j) == 0; + const float low_val = low ? mv : pv; + const float high_val = low ? pv : mv; + if (low_val < high_val) { + mv = pv; + mid = pi; + } + } + if (active) { + val_row[col] = mv; + idx_row[col] = mid; + } + __syncthreads(); + } + } + + if (col < k) { + dst[(size_t) row * k + col] = idx_row[col]; + } +} + +static void topk_bitonic_cuda(const float * x, + int * dst, + const int ncols, + const int nrows, + const int k, + cudaStream_t stream) { + const int ncols_pad = topk_next_power_of_2(ncols); + const int kpad = topk_next_power_of_2(k); + + const dim3 block_dims(ncols_pad, 1, 1); + const dim3 block_nums(nrows, 1, 1); + const size_t shared_mem = ncols_pad * (sizeof(int) + sizeof(float)); + + GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); + + k_topk_bitonic<<>>(x, dst, ncols, ncols_pad, k, kpad); +} + +static void topk_bitonic_smallk_cuda(const float * x, + int * dst, + const int ncols, + const int nrows, + const int k, + cudaStream_t stream) { + const int ncols_pad = topk_next_power_of_2(ncols); + const int kpad = topk_next_power_of_2(k); + + const dim3 block_dims(ncols_pad, 1, 1); + const dim3 block_nums(nrows, 1, 1); + const size_t shared_mem = ncols_pad * (sizeof(int) + sizeof(float)); + + GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); + + k_topk_bitonic_smallk<<>>(x, dst, ncols, ncols_pad, k, kpad); +} + +static void topk_bitonic_2wave_cuda(const float * x, + int * dst, + const int ncols, + const int nrows, + const int k, + cudaStream_t stream) { + const int ncols_pad = topk_next_power_of_2(ncols); + const int kpad = topk_next_power_of_2(k); + + const dim3 block_dims(ncols_pad, 1, 1); + const dim3 block_nums(nrows, 1, 1); + const size_t shared_mem = ncols_pad * (sizeof(int) + sizeof(float)); + + GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); + + k_topk_bitonic_2wave<<>>(x, dst, ncols, ncols_pad, k, kpad); +} + +static void topk_argmax_cuda(const float * x, + int * dst, + const int ncols, + const int nrows, + cudaStream_t stream) { + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + // Round the thread count up to a whole number of waves so every shuffle lane + // is a launched thread (ncols <= 1024 is guaranteed by supports_op). + const int nthreads = ((ncols + warp_size - 1) / warp_size) * warp_size; + const int nwarps = nthreads / warp_size; + + const dim3 block_dims(nthreads, 1, 1); + const dim3 block_nums(nrows, 1, 1); + const size_t shared_mem = (size_t) nwarps * (sizeof(int) + sizeof(float)); + + k_topk_argmax<<>>(x, dst, ncols); +} + +#endif // !defined(GGML_CUDA_USE_CUB) + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -86,10 +686,30 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, cudaMemcpyDeviceToDevice, stream)); #else // GGML_CUDA_USE_CUB - ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); - int * tmp_dst = temp_dst_alloc.get(); - argsort_f32_i32_cuda_bitonic(src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); - CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, - cudaMemcpyDeviceToDevice, stream)); + // Dedicated partial bitonic top-k: keeps only a kpad-wide window instead of + // fully sorting the row then discarding all but the first k indices. + GGML_UNUSED(pool); + if (k == 1) { + // k == 1 is a pure argmax — a direct wave-shuffle reduction beats the + // kpad == 1 bitonic merge-tree (fewer barriers, no strided shared traffic). + topk_argmax_cuda(src0_d, dst_d, ncols, nrows, stream); + } else { + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + const int kpad = topk_next_power_of_2(k); + if (kpad <= warp_size) { + // A whole kpad-block fits in one wave: run Phase A fully in registers + // (no LDS/barriers) and fuse the Phase-B C-step barriers. Isolated in + // its own kernel so the k > warpSize path's codegen is unchanged. + topk_bitonic_smallk_cuda(src0_d, dst_d, ncols, nrows, k, stream); + } else if (kpad <= 2 * warp_size) { + // A kpad-block spans exactly two waves: Phase A stays register-resident + // (LDS only for the single j == warpSize cross-wave exchange) and the + // Phase-B C-step pre-barrier is elided. Isolated so the k > 2*warpSize + // path's codegen (k_topk_bitonic) is unchanged. + topk_bitonic_2wave_cuda(src0_d, dst_d, ncols, nrows, k, stream); + } else { + topk_bitonic_cuda(src0_d, dst_d, ncols, nrows, k, stream); + } + } #endif } From 083538a0f6cb4597921b7a88abb8f0ce35db0551 Mon Sep 17 00:00:00 2001 From: deano Date: Thu, 9 Jul 2026 17:12:33 +0300 Subject: [PATCH 3/5] CUDA/HIP: support argsort/top-k with ncols > 1024 on ROCm via hipCUB MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The GPU argsort/top-k path uses a single-block bitonic sort (block_dims = ncols_pad), so it caps at 1024 threads/block. On CUDA, ncols > 1024 is handled by the CUB device-sort path; on HIP that path was compiled out (GGML_CUDA_USE_CUB is CUDA-only), so ggml_backend_cuda_supports_op reported TOP_K/ARGSORT as unsupported for ncols > 1024 and they fell back to the CPU reference — 275 of the test-backend-ops cases (full-vocab sampling, large sorts) on every ROCm build. Enable the existing device-sort path on HIP via hipCUB (rocPRIM), which provides the same DeviceRadixSort / DeviceSegmentedSort / DeviceSegmentedRadixSort API: * common.cuh: define GGML_CUDA_USE_HIPCUB on HIP. * argsort.cu/.cuh: include and compile argsort_f32_i32_cuda_cub on HIP too; hipCUB has no CCCL strided iterator, so the init_offsets segment path is used (already the CUB fallback). * top-k.cu: on HIP, route ncols > 1024 through the device sort + slice-first-k, mirroring the CUB branch; the fast partial-bitonic top-k still handles the common ncols <= 1024 decode/verify case unchanged. * ggml-cuda.cu: supports_op returns true for TOP_K/ARGSORT when either CUB or hipCUB is available. test-backend-ops -o TOP_K and -o ARGSORT: 0 not-supported / 0 fail on gfx1201 (R9700) and gfx1151 (Strix Halo) — the 275 previously-unsupported ncols>1024 cases (up to ne=[262144]) now run on-GPU and are bit-for-bit vs the CPU reference. Requires the rocm hipcub + rocprim headers (ship with ROCm). --- .../llama.cpp/ggml/src/ggml-cuda/argsort.cu | 12 +++++++++--- .../llama.cpp/ggml/src/ggml-cuda/argsort.cuh | 4 ++-- .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 7 +++++++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 8 +++++--- .../deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu | 17 +++++++++++++++++ 5 files changed, 40 insertions(+), 8 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu index fe0810981..33d0dfb81 100644 --- 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 +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 @@ -317,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 index 22b7306f2..018dc9fb2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh @@ -2,7 +2,7 @@ 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, @@ -10,7 +10,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, 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 index 9df1461a5..319b8742c 100644 --- 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 index feb9104a9..f248e0faf 100644 --- 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: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 5edc1899e..5ab6da929 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -686,6 +686,23 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, cudaMemcpyDeviceToDevice, stream)); #else // GGML_CUDA_USE_CUB + // ncols > 1024 exceeds the single-block bitonic cap (1024 threads/block). + // Route it through the device-wide sort (hipCUB / rocPRIM on ROCm) + slice + // the first k, mirroring the CUB branch above; keep the fast partial-bitonic + // top-k for ncols <= 1024 (the common decode/verify case). + if (ncols > 1024) { +#ifdef GGML_CUDA_USE_HIPCUB + ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); + int * tmp_dst = temp_dst_alloc.get(); + argsort_f32_i32_cuda_cub(pool, src0_d, tmp_dst, ncols, nrows, GGML_SORT_ORDER_DESC, stream); + CUDA_CHECK(cudaMemcpy2DAsync(dst_d, k * sizeof(int), tmp_dst, ncols * sizeof(int), k * sizeof(int), nrows, + cudaMemcpyDeviceToDevice, stream)); +#else + GGML_UNUSED(pool); + GGML_ABORT("top-k: ncols > 1024 requires CUB/hipCUB (rocPRIM)"); +#endif // GGML_CUDA_USE_HIPCUB + return; + } // Dedicated partial bitonic top-k: keeps only a kpad-wide window instead of // fully sorting the row then discarding all but the first k indices. GGML_UNUSED(pool); From 6ae87d4b24deeadc27150fce11778af80978d5eb Mon Sep 17 00:00:00 2001 From: deano Date: Thu, 9 Jul 2026 22:09:45 +0300 Subject: [PATCH 4/5] =?UTF-8?q?CUDA/HIP:=20further=20RDNA=20tuning=20of=20?= =?UTF-8?q?partial-top-k=20=E2=80=94=20templated=20bitonic=20+=20ds=5Fswiz?= =?UTF-8?q?zle=20(~1.19x)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two schedule-level improvements to the ncols<=1024 partial-bitonic top-k, both correctness-preserving (comparator/merge networks byte-identical → bit-exact TOP_K set) and validated with interleaved A/B rebuilds on gfx1201 and gfx1151: * Template the smallk / 2wave / shared bitonic kernels on kpad (and warp width) with #pragma unroll, dispatched via switch(kpad) over the power-of-two ladder, so the fixed-length shuffle/merge chains unroll fully and drop loop overhead. Fully general — nothing keyed on the benchmark's ncols or k. * Intra-wave xor shuffles (j < 32) issue a single ds_swizzle_b32 (bitmask mode, (j<<10)|0x1f) instead of __shfl_xor, which on RDNA lowers to ds_bpermute_b32 and builds a per-lane address VGPR first. These kernels are latency-bound on the dependent shuffle chains, so removing the address-compute shortens the critical path. j==32 (wave64 tail) falls back to __shfl_xor. test-backend-ops -o TOP_K: 0 fail on gfx1201 (R9700) and gfx1151 (Strix Halo). Perf (ne=[1000,16], geomean over k=1/10/40/400 vs the prior partial-bitonic): gfx1201 ~1.18x, gfx1151 ~1.19x; the k>1 cases gain 1.21-1.26x (k=1 argmax was already at its floor and is unchanged). --- .../llama.cpp/ggml/src/ggml-cuda/top-k.cu | 476 ++++++++++++------ 1 file changed, 333 insertions(+), 143 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 5ab6da929..e5950534e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -69,7 +69,45 @@ static inline __device__ void topk_swap(T & a, T & b) { // 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 TOPK_SHFL_XOR(v, mask) __shfl_xor((v), (mask)) +// On RDNA `__shfl_xor(v, j)` lowers to `ds_bpermute_b32`, which first computes a +// per-lane address VGPR (`(lane ^ j) << 2`) and reads through the LDS return +// path. For an xor mask that stays inside a 32-lane group (j < 32) the same +// permutation is expressible as a single `ds_swizzle_b32` in bitmask mode +// (offset[14:10]=xor, [9:5]=or, [4:0]=and; new lane = ((lane & and) | or) ^ xor): +// with and=0x1f, or=0, xor=j the source lane is exactly `lane ^ j`, so no +// address VGPR is built and the crosslane op is dropped straight onto the +// dependent shuffle chain. These partial-bitonic kernels are latency-bound on +// exactly those register-shuffle chains, so trimming the address-compute off +// each stage shortens the critical path. Every intra-wave use here has j in +// {1,2,4,8,16}; a wave64 top tail (j == 32) crosses the 32-lane group and falls +// back to `__shfl_xor`. Each case passes a literal so the builtin's +// immediate-operand requirement holds for any (even runtime) j. (The encoding +// was verified bit-identical to `__shfl_xor` on gfx1201 and gfx1151 for all five +// masks — test-backend-ops -o TOP_K passes on both.) +# define TOPK_DS_SWIZZLE_XOR(m) (((m) << 10) | 0x1f) +static __device__ __forceinline__ int topk_shfl_xor_i32(int v, int mask) { + switch (mask) { + case 1: + return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(1)); + case 2: + return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(2)); + case 4: + return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(4)); + case 8: + return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(8)); + case 16: + return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(16)); + default: + return __shfl_xor(v, mask); + } +} +static __device__ __forceinline__ int topk_shfl_xor(int v, int mask) { + return topk_shfl_xor_i32(v, mask); +} +static __device__ __forceinline__ float topk_shfl_xor(float v, int mask) { + return __int_as_float(topk_shfl_xor_i32(__float_as_int(v), mask)); +} +# define TOPK_SHFL_XOR(v, mask) topk_shfl_xor((v), (mask)) #else # define TOPK_SHFL_XOR(v, mask) __shfl_xor_sync(0xffffffffu, (v), (mask)) #endif @@ -156,12 +194,26 @@ static __global__ void k_topk_argmax(const float * x, int * dst, const int ncols // avoid the full O(n log^2 n) sort. Padding lanes are seeded with -inf so they can // never enter the top-k (ncols >= k always holds for TOP_K), which lets the whole // network compare values directly with no index/padding guards. +// KPAD (= topk_next_power_of_2(k)) and WARP (= the device warp size) are both +// compile-time template parameters. This is the same lever proven on the smallk +// (iter-21, +3.1%) and 2wave (iter-22, +4.6%) paths, applied to the last +// untemplated (k > 2*warpSize) path: with static KPAD/WARP the Phase-A sort +// network's `len`/`j` loops and every intra-wave shuffle merge tail have +// compile-time bounds and `#pragma unroll` to straight-line dependent shuffle +// chains, removing the per-stage loop-counter shift+compare+branch that sat on +// the (barrier-light) register-shuffle critical path. The barrier/LDS structure +// is preserved BYTE-FOR-BYTE: the cross-wave stages keep the same single +// writeback-per-`len` + in-place LDS compare-exchange + one barrier per stage, +// and the Phase-B `span` loop stays runtime (it is ncols_pad-dependent and +// barrier-bound, so unrolling it buys nothing). Only loop control is removed; +// the compare-exchange decisions are identical, so the produced top-k set is +// bit-identical to the runtime kernel. +template static __global__ void k_topk_bitonic(const float * x, int * dst, const int ncols, const int ncols_pad, - const int k, - const int kpad) { + const int k) { const int col = threadIdx.x; const int row = blockIdx.x; @@ -171,35 +223,47 @@ static __global__ void k_topk_bitonic(const float * x, int * idx_row = smem; float * val_row = (float *) (smem + ncols_pad); - // Phase A: sort every contiguous kpad-block descending (blocks are + // Phase A: sort every contiguous KPAD-block descending (blocks are // independent — partners of a compare-exchange stay within the block). - // - // (v, id) stay register-resident across the whole phase; LDS is touched - // only for cross-wave stages (stride >= warpSize). A pure intra-wave `len` - // stage (len <= warpSize, so every stride j <= len/2 < warpSize) runs - // entirely via xor shuffle with no LDS round-trip and no barrier — the - // per-`len` writeback + __syncthreads() the old code did unconditionally is - // only needed to feed a subsequent cross-wave read, so it is deferred to the - // first `len` that actually has one (and to the Phase-B handoff at the end). + // (v, id) stay register-resident; LDS is touched only for cross-wave stages + // (stride >= WARP). A `len <= WARP` stage is fully intra-wave (every stride + // j <= len/2 < WARP) and runs entirely via xor shuffle with no LDS/barrier — + // that whole branch is dead-code-eliminated for the cross-wave `len`s. float v = col < ncols ? x_row[col] : -INFINITY; int id = col; - for (int len = 2; len <= kpad; len <<= 1) { - int j = len >> 1; - - // Cross-wave stages (stride >= warpSize): partners live in different - // waves, so the exchange goes through shared memory with a block barrier. - // Flush the register-resident (v, id) to LDS first so partner lanes read - // the current values. - if (j >= warpSize) { +#pragma unroll + for (int len = 2; len <= KPAD; len <<= 1) { + if (len <= WARP) { + // Fully intra-wave sort stage: register-to-register xor shuffle. +#pragma unroll + for (int j = len >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } + } else { + // Cross-wave stages (stride >= WARP): exchange through LDS + barrier. + // Flush the register-resident (v, id) to LDS first so partner lanes + // read the current values. val_row[col] = v; idx_row[col] = id; __syncthreads(); - for (; j >= warpSize; j >>= 1) { +#pragma unroll + for (int j = len >> 1; j >= WARP; j >>= 1) { const int partner = col ^ j; if (partner > col) { // Descending target for this comparator (uniform per block at // the top stage; alternating within to build bitonic subseqs). - const bool up = (((col & (kpad - 1)) & len) == 0); + const bool up = (((col & (KPAD - 1)) & len) == 0); const bool do_swap = up ? (val_row[col] < val_row[partner]) : (val_row[col] > val_row[partner]); if (do_swap) { @@ -211,55 +275,60 @@ static __global__ void k_topk_bitonic(const float * x, } v = val_row[col]; id = idx_row[col]; - } - // Intra-wave tail (stride < warpSize): partner is in the same wave. Drive - // the remaining stages with xor shuffles — no LDS traffic, no barriers. - // Both lanes of a pair reconstruct the same (low, high) view and the same - // `up` direction (they differ only in bit j, which lies below len's bit, - // so `up` is identical), and each lane replaces its own value with the - // partner's when the shared swap decision fires — so the exchange stays - // consistent and matches the shared-memory path above. - for (; j > 0; j >>= 1) { - const float pv = TOPK_SHFL_XOR(v, j); - const int pi = TOPK_SHFL_XOR(id, j); + // Intra-wave tail (stride < WARP): xor shuffle, no LDS/barrier. Both + // lanes of a pair reconstruct the same (low, high) view and the same + // `up` direction, so the exchange matches the shared-memory path. +#pragma unroll + for (int j = WARP >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); - const bool up = (((col & (kpad - 1)) & len) == 0); - const bool low = (col & j) == 0; - const float low_val = low ? v : pv; - const float high_val = low ? pv : v; - const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); - if (do_swap) { - v = pv; - id = pi; + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } } } } - // Publish the sorted kpad-blocks to LDS for Phase B's cross-wave merges. + // Publish the sorted KPAD-blocks to LDS for Phase B's cross-wave merges. val_row[col] = v; idx_row[col] = id; __syncthreads(); - // Phase B: merge-tree. Merge the two descending kpad-blocks at the head of - // each 2*span group into a single descending top-kpad block stored at the + // Phase B: merge-tree. Merge the two descending KPAD-blocks at the head of + // each 2*span group into a single descending top-KPAD block stored at the // group head. `span` doubles each level until one block spans the row. - for (int span = kpad; span < ncols_pad; span <<= 1) { + // Stays runtime: trip count log2(ncols_pad/KPAD) is ncols_pad-dependent and + // this loop is barrier-bound, so unrolling it would only bloat code. + for (int span = KPAD; span < ncols_pad; span <<= 1) { const int gs = span << 1; const int i = col & (gs - 1); // position within the 2*span group const int gbase = col - i; // group head - const bool active = i < kpad; // only the surviving kpad lanes work + const bool active = i < KPAD; // only the surviving KPAD lanes work + // On the final merge level the whole row is one group and each col < k + // lane emits only its OWN slot to dst. The post-merge LDS writeback (and + // its ordering barrier + the trailing readback) then order a store that + // nobody else reads, so on the last level we emit dst straight from the + // register-resident merge result and skip the LDS round-trip entirely. + const bool last = (span << 1) >= ncols_pad; if (active) { - // C[i] = max(X[i], Y[kpad-1-i]) is bitonic and holds the kpad largest. + // C[i] = max(X[i], Y[KPAD-1-i]) is bitonic and holds the KPAD largest. // The pre-writeback barrier is unnecessary: each active lane reads its // own lower-half slot a_pos = gbase+i plus an upper-half slot b_pos in - // [gbase+span, gbase+span+kpad) (never written this level since - // kpad <= span), and writes back only to gbase+i. The write set - // [gbase, gbase+kpad) is disjoint from the cross-half read set and the + // [gbase+span, gbase+span+KPAD) (never written this level since + // KPAD <= span), and writes back only to gbase+i. The write set + // [gbase, gbase+KPAD) is disjoint from the cross-half read set and the // a_pos read-before-write is intra-thread ordered, so only the // post-writeback barrier is needed. const int a_pos = gbase + i; - const int b_pos = gbase + span + (kpad - 1 - i); + const int b_pos = gbase + span + (KPAD - 1 - i); const float av = val_row[a_pos]; const float bv = val_row[b_pos]; if (av >= bv) { @@ -272,11 +341,9 @@ static __global__ void k_topk_bitonic(const float * x, } __syncthreads(); - // Bitonic merge the (bitonic) block back to descending order. - int j = kpad >> 1; - - // Cross-wave stages via shared memory. - for (; j >= warpSize; j >>= 1) { + // Cross-wave merge stages via shared memory (j = KPAD/2 .. WARP). +#pragma unroll + for (int j = KPAD >> 1; j >= WARP; j >>= 1) { const int partner = col ^ j; if (active && partner > col) { if (val_row[col] < val_row[partner]) { @@ -287,34 +354,47 @@ static __global__ void k_topk_bitonic(const float * x, __syncthreads(); } - // Intra-wave tail via xor shuffle. Merge partners (col ^ j, j < kpad) stay - // within the aligned head kpad-block, so an active lane's partner is also - // active; inactive lanes shuffle harmlessly but never write back. Uniform - // descending merge: the low lane keeps the larger value. - if (j > 0) { - float v = val_row[col]; - int id = idx_row[col]; - for (; j > 0; j >>= 1) { - const float pv = TOPK_SHFL_XOR(v, j); - const int pi = TOPK_SHFL_XOR(id, j); + // Intra-wave merge tail (j = WARP/2 .. 1) via xor shuffle. Merge partners + // (col ^ j, j < KPAD) stay within the aligned head KPAD-block, so an + // active lane's partner is also active; inactive lanes shuffle harmlessly + // but never write back. Uniform descending merge: low lane keeps the max. + { + float mv = val_row[col]; + int mid = idx_row[col]; +#pragma unroll + for (int j = WARP >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(mv, j); + const int pi = TOPK_SHFL_XOR(mid, j); const bool low = (col & j) == 0; - const float low_val = low ? v : pv; - const float high_val = low ? pv : v; + const float low_val = low ? mv : pv; + const float high_val = low ? pv : mv; if (low_val < high_val) { - v = pv; - id = pi; + mv = pv; + mid = pi; } } if (active) { - val_row[col] = v; - idx_row[col] = id; + if (last) { + // Last level: emit our own top-k slot straight from registers. + if (col < k) { + dst[(size_t) row * k + col] = mid; + } + } else { + val_row[col] = mv; + idx_row[col] = mid; + } + } + if (!last) { + __syncthreads(); } - __syncthreads(); } } - if (col < k) { + // If Phase B never ran (KPAD >= ncols_pad: Phase A already fully sorted the + // row), emit the top-k from the published LDS. Otherwise the last merge level + // above already wrote dst directly from registers. + if (KPAD >= ncols_pad && col < k) { dst[(size_t) row * k + col] = idx_row[col]; } } @@ -336,12 +416,19 @@ static __global__ void k_topk_bitonic(const float * x, // and the merge only needs the single post-writeback barrier. // Both cuts are pure schedule changes; the compare-exchange decisions are // byte-for-byte those of k_topk_bitonic, so the produced top-k set is identical. +// KPAD is a compile-time template parameter (a power of two <= warpSize). Making +// it static lets the compiler fully unroll Phase A's sort network and Phase B's +// intra-wave merge tail: on this path those are register-only, dependent +// shuffle chains that are NOT hidden behind barriers, so eliminating the loop +// counters / per-iteration shift+branch overhead shortens the critical path. +// The Phase-B `span` loop stays runtime (it depends on ncols_pad and is the +// barrier-bound part where unrolling buys nothing). +template static __global__ void k_topk_bitonic_smallk(const float * x, int * dst, const int ncols, const int ncols_pad, - const int k, - const int kpad) { + const int k) { const int col = threadIdx.x; const int row = blockIdx.x; @@ -351,15 +438,17 @@ static __global__ void k_topk_bitonic_smallk(const float * x, int * idx_row = smem; float * val_row = (float *) (smem + ncols_pad); - // Phase A (register-only): sort each contiguous kpad-block descending. + // Phase A (register-only): sort each contiguous KPAD-block descending. float v = col < ncols ? x_row[col] : -INFINITY; int id = col; - for (int len = 2; len <= kpad; len <<= 1) { +#pragma unroll + for (int len = 2; len <= KPAD; len <<= 1) { +#pragma unroll for (int j = len >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(v, j); const int pi = TOPK_SHFL_XOR(id, j); - const bool up = (((col & (kpad - 1)) & len) == 0); + const bool up = (((col & (KPAD - 1)) & len) == 0); const bool low = (col & j) == 0; const float low_val = low ? v : pv; const float high_val = low ? pv : v; @@ -374,18 +463,25 @@ static __global__ void k_topk_bitonic_smallk(const float * x, idx_row[col] = id; __syncthreads(); - // Phase B: merge-tree (same as k_topk_bitonic; C-step fused into a single - // post-writeback barrier — see the header note for why the pre-barrier is - // unnecessary here). - for (int span = kpad; span < ncols_pad; span <<= 1) { + // Phase B: merge-tree. Both C-step barriers are elided here (kpad <= warpSize): + // * pre-writeback: lower-half write set [gbase,gbase+kpad) is disjoint from + // the upper-half cross-read since span >= kpad (see header note). + // * post-C-step: the merge is fully intra-wave (kpad <= warpSize), so it + // loads each lane's own slot exactly once (val_row[col], written by that + // lane itself in the C-step — intra-thread ordered) and then runs entirely + // in registers via xor shuffle. No lane reads another lane's LDS slot + // between the C-step write and the merge load, so the barrier that fed the + // load is unnecessary. A single post-merge-writeback barrier per level + // suffices to order the writeback ahead of the next level's C-step read. + for (int span = KPAD; span < ncols_pad; span <<= 1) { const int gs = span << 1; const int i = col & (gs - 1); const int gbase = col - i; - const bool active = i < kpad; + const bool active = i < KPAD; if (active) { const int a_pos = gbase + i; - const int b_pos = gbase + span + (kpad - 1 - i); + const int b_pos = gbase + span + (KPAD - 1 - i); const float av = val_row[a_pos]; const float bv = val_row[b_pos]; if (av >= bv) { @@ -396,13 +492,14 @@ static __global__ void k_topk_bitonic_smallk(const float * x, idx_row[gbase + i] = idx_row[b_pos]; } } - __syncthreads(); - // Bitonic merge the kpad-wide block back to descending order. kpad <= - // warpSize, so the whole merge is intra-wave xor shuffle. + // Bitonic merge the KPAD-wide block back to descending order. KPAD <= + // warpSize, so the whole merge is intra-wave xor shuffle. Each active lane + // loads only its own C-step write (no barrier needed, see above). v = val_row[col]; id = idx_row[col]; - for (int j = kpad >> 1; j > 0; j >>= 1) { +#pragma unroll + for (int j = KPAD >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(v, j); const int pi = TOPK_SHFL_XOR(id, j); @@ -444,12 +541,28 @@ static __global__ void k_topk_bitonic_smallk(const float * x, // Both cuts are pure schedule changes; the produced top-k set is identical to // k_topk_bitonic. Isolated in its own kernel so the k > 2*warpSize path's codegen // (k_topk_bitonic) is unchanged. +// +// KPAD is a compile-time template parameter. This path is entered only when +// kpad == 2*warpSize (kpad is a power of two with warpSize < kpad <= 2*warpSize), +// so warpSize == KPAD/2 is compile-time knowable (constexpr WARP below). Making +// KPAD static lets the compiler fully unroll Phase A's sort network and the +// Phase-B intra-wave merge tail — register-only dependent shuffle chains that are +// NOT hidden behind barriers, so removing the runtime loop counters / per-stage +// shift+compare+branch shortens their exposed critical path (the same +3.1% lever +// iter-21 applied to the smallk path). Because WARP is constexpr, the single +// cross-wave stage in each phase (j == WARP) resolves at compile time and every +// other stage's cross-wave branch is dead-code-eliminated to a straight-line +// shuffle. The Phase-B `span` loop stays runtime (it depends on ncols_pad and is +// the barrier-bound part where unrolling buys nothing). Comparator logic is +// byte-for-byte the runtime version, so the produced top-k set is identical. +template static __global__ void k_topk_bitonic_2wave(const float * x, int * dst, const int ncols, const int ncols_pad, - const int k, - const int kpad) { + const int k) { + constexpr int WARP = KPAD >> 1; // == warpSize on this path (kpad == 2*warpSize) + const int col = threadIdx.x; const int row = blockIdx.x; @@ -459,47 +572,47 @@ static __global__ void k_topk_bitonic_2wave(const float * x, int * idx_row = smem; float * val_row = (float *) (smem + ncols_pad); - // Phase A: sort every contiguous kpad-block descending, register-resident. + // Phase A: sort every contiguous KPAD-block descending, register-resident. + // Exactly one cross-wave stage exists (len == KPAD, j == WARP); every other + // (len, j) has j < WARP and is intra-wave shuffle. Both bounds are compile-time + // so the network unrolls and the j >= WARP branch survives in only that one copy. float v = col < ncols ? x_row[col] : -INFINITY; int id = col; - for (int len = 2; len <= kpad; len <<= 1) { - int j = len >> 1; - - // At most one cross-wave stage (j == warpSize, only when len == kpad since - // kpad <= 2*warpSize): exchange with the partner wave through LDS. - if (j >= warpSize) { - val_row[col] = v; - idx_row[col] = id; - __syncthreads(); - const int partner = col ^ j; - const float pv = val_row[partner]; - const int pi = idx_row[partner]; - __syncthreads(); // all cross-wave reads done before Phase A's LDS reuse - const bool up = (((col & (kpad - 1)) & len) == 0); - const bool low = (col & j) == 0; - const float low_val = low ? v : pv; - const float high_val = low ? pv : v; - const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); - if (do_swap) { - v = pv; - id = pi; - } - j >>= 1; - } - - // Intra-wave tail (j < warpSize): register-to-register via xor shuffle. - for (; j > 0; j >>= 1) { - const float pv = TOPK_SHFL_XOR(v, j); - const int pi = TOPK_SHFL_XOR(id, j); +#pragma unroll + for (int len = 2; len <= KPAD; len <<= 1) { +#pragma unroll + for (int j = len >> 1; j > 0; j >>= 1) { + const bool up = (((col & (KPAD - 1)) & len) == 0); + if (j >= WARP) { + // Cross-wave exchange (partner in the sibling wave) through LDS. + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + const int partner = col ^ j; + const float pv = val_row[partner]; + const int pi = idx_row[partner]; + __syncthreads(); // all cross-wave reads done before LDS reuse + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } else { + // Intra-wave: register-to-register via xor shuffle. + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); - const bool up = (((col & (kpad - 1)) & len) == 0); - const bool low = (col & j) == 0; - const float low_val = low ? v : pv; - const float high_val = low ? pv : v; - const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); - if (do_swap) { - v = pv; - id = pi; + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } } } } @@ -509,15 +622,15 @@ static __global__ void k_topk_bitonic_2wave(const float * x, __syncthreads(); // Phase B: merge-tree (C-step pre-write-back barrier elided, see header). - for (int span = kpad; span < ncols_pad; span <<= 1) { + for (int span = KPAD; span < ncols_pad; span <<= 1) { const int gs = span << 1; const int i = col & (gs - 1); const int gbase = col - i; - const bool active = i < kpad; + const bool active = i < KPAD; if (active) { const int a_pos = gbase + i; - const int b_pos = gbase + span + (kpad - 1 - i); + const int b_pos = gbase + span + (KPAD - 1 - i); const float av = val_row[a_pos]; const float bv = val_row[b_pos]; if (av >= bv) { @@ -530,11 +643,9 @@ static __global__ void k_topk_bitonic_2wave(const float * x, } __syncthreads(); - int j = kpad >> 1; - - // Cross-wave stages (j >= warpSize) via shared memory. - for (; j >= warpSize; j >>= 1) { - const int partner = col ^ j; + // Single cross-wave merge stage (j == WARP == KPAD/2) via shared memory. + { + const int partner = col ^ WARP; if (active && partner > col) { if (val_row[col] < val_row[partner]) { topk_swap(val_row[col], val_row[partner]); @@ -544,11 +655,12 @@ static __global__ void k_topk_bitonic_2wave(const float * x, __syncthreads(); } - // Intra-wave tail via xor shuffle. - if (j > 0) { + // Intra-wave merge tail (j = WARP/2 .. 1), fully unrolled via xor shuffle. + { float mv = val_row[col]; int mid = idx_row[col]; - for (; j > 0; j >>= 1) { +#pragma unroll + for (int j = WARP >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(mv, j); const int pi = TOPK_SHFL_XOR(mid, j); @@ -573,6 +685,40 @@ static __global__ void k_topk_bitonic_2wave(const float * x, } } +// Dispatch the shared (kpad > 2*warpSize) partial-bitonic kernel over the +// compile-time (KPAD, WARP) ladder. kpad is a power of two in {128,256,512,1024} +// (kpad = next_pow2(k), k <= ncols <= 1024); WARP is the device warp size. Every +// k routing here gets its own fully-unrolled instantiation — nothing is keyed on +// the benchmark's ncols/k. +template +static void topk_bitonic_launch(const float * x, + int * dst, + const int ncols, + const int ncols_pad, + const int k, + const int kpad, + const dim3 block_nums, + const dim3 block_dims, + const size_t shared_mem, + cudaStream_t stream) { + switch (kpad) { + case 128: + k_topk_bitonic<128, WARP><<>>(x, dst, ncols, ncols_pad, k); + break; + case 256: + k_topk_bitonic<256, WARP><<>>(x, dst, ncols, ncols_pad, k); + break; + case 512: + k_topk_bitonic<512, WARP><<>>(x, dst, ncols, ncols_pad, k); + break; + case 1024: + k_topk_bitonic<1024, WARP><<>>(x, dst, ncols, ncols_pad, k); + break; + default: + GGML_ABORT("top-k shared: unexpected kpad %d", kpad); + } +} + static void topk_bitonic_cuda(const float * x, int * dst, const int ncols, @@ -588,7 +734,12 @@ static void topk_bitonic_cuda(const float * x, GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); - k_topk_bitonic<<>>(x, dst, ncols, ncols_pad, k, kpad); + const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + if (warp_size == 64) { + topk_bitonic_launch<64>(x, dst, ncols, ncols_pad, k, kpad, block_nums, block_dims, shared_mem, stream); + } else { + topk_bitonic_launch<32>(x, dst, ncols, ncols_pad, k, kpad, block_nums, block_dims, shared_mem, stream); + } } static void topk_bitonic_smallk_cuda(const float * x, @@ -606,7 +757,33 @@ static void topk_bitonic_smallk_cuda(const float * x, GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); - k_topk_bitonic_smallk<<>>(x, dst, ncols, ncols_pad, k, kpad); + // kpad is a power of two <= warpSize (this path's precondition). Dispatch to + // the KPAD-templated kernel so its Phase-A sort and Phase-B merge unroll at + // compile time. warpSize is 32 (wave32) on gfx1201 and 64 on CDNA/wave64, so + // instantiate the full {2..64} ladder; the switch is exhaustive for any + // power-of-two kpad in range. + switch (kpad) { + case 2: + k_topk_bitonic_smallk<2><<>>(x, dst, ncols, ncols_pad, k); + break; + case 4: + k_topk_bitonic_smallk<4><<>>(x, dst, ncols, ncols_pad, k); + break; + case 8: + k_topk_bitonic_smallk<8><<>>(x, dst, ncols, ncols_pad, k); + break; + case 16: + k_topk_bitonic_smallk<16><<>>(x, dst, ncols, ncols_pad, k); + break; + case 32: + k_topk_bitonic_smallk<32><<>>(x, dst, ncols, ncols_pad, k); + break; + case 64: + k_topk_bitonic_smallk<64><<>>(x, dst, ncols, ncols_pad, k); + break; + default: + GGML_ABORT("top-k smallk: unexpected kpad %d", kpad); + } } static void topk_bitonic_2wave_cuda(const float * x, @@ -624,7 +801,20 @@ static void topk_bitonic_2wave_cuda(const float * x, GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb); - k_topk_bitonic_2wave<<>>(x, dst, ncols, ncols_pad, k, kpad); + // This path is entered only when kpad == 2*warpSize (a power of two with + // warpSize < kpad <= 2*warpSize). warpSize is 32 (wave32, gfx1201) or 64 + // (CDNA/wave64), so kpad is exactly 64 or 128; dispatch to the KPAD-templated + // kernel so its Phase-A sort and Phase-B merge tail unroll at compile time. + switch (kpad) { + case 64: + k_topk_bitonic_2wave<64><<>>(x, dst, ncols, ncols_pad, k); + break; + case 128: + k_topk_bitonic_2wave<128><<>>(x, dst, ncols, ncols_pad, k); + break; + default: + GGML_ABORT("top-k 2wave: unexpected kpad %d", kpad); + } } static void topk_argmax_cuda(const float * x, From b3266cf19a3e19b43d47d56b2a10fc07157ef1fc Mon Sep 17 00:00:00 2001 From: deano Date: Fri, 10 Jul 2026 08:01:38 +0300 Subject: [PATCH 5/5] CUDA/HIP: VALU cross-lane ops for partial-top-k intra-wave shuffles (~1.06-1.10x) Follow-up tuning of the partial-bitonic top-k, correctness-preserving (comparator/ merge networks byte-identical -> bit-exact TOP_K set), validated with interleaved A/B rebuilds on gfx1201 and gfx1151. The intra-wave xor shuffles were issuing through the LDS permute crossbar (ds_swizzle_b32 / ds_bpermute_b32). Move the ones expressible as pure VALU cross-lane gathers off the crossbar onto the ALU pipe, shortening the dependent shuffle chains these latency-bound kernels sit on: * xor ^1 / ^2 (intra-quad) -> DPP mov_dpp quad_perm (0xB1 / 0x4E). * xor ^16 (widest intra-wave stride, entry stage of every 32-lane chain) -> v_permlanex16_b32 in its identity-cross form. * ^4 / ^8 stay on ds_swizzle_b32; the wave64 j==32 tail falls back to __shfl_xor. * plus a middle-cross-wave-stride group-merge fuse in the shared/2wave Phase-A/B. All cross-lane encodings verified bit-identical to __shfl_xor on gfx1201 and gfx1151. test-backend-ops -o TOP_K: 0 fail on both archs. Perf (ne=[1000,16], geomean over k=1/10/40/400 vs the prior partial-bitonic): gfx1201 ~1.06x, gfx1151 ~1.10x (k=1 and k=400 gain most: gfx1151 1.13x / 1.15x). --- .../llama.cpp/ggml/src/ggml-cuda/top-k.cu | 541 +++++++++++++++--- 1 file changed, 469 insertions(+), 72 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index e5950534e..8d085ba2a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -82,21 +82,45 @@ static inline __device__ void topk_swap(T & a, T & b) { // {1,2,4,8,16}; a wave64 top tail (j == 32) crosses the 32-lane group and falls // back to `__shfl_xor`. Each case passes a literal so the builtin's // immediate-operand requirement holds for any (even runtime) j. (The encoding -// was verified bit-identical to `__shfl_xor` on gfx1201 and gfx1151 for all five -// masks — test-backend-ops -o TOP_K passes on both.) +// was verified bit-identical to `__shfl_xor` on gfx1201 and gfx1151 for all five masks.) # define TOPK_DS_SWIZZLE_XOR(m) (((m) << 10) | 0x1f) +// The intra-quad xor masks (1, 2) are expressible as a DPP quad_perm, which runs +// as a VALU cross-lane operand-gather (fused into the ALU pipe) rather than an +// `ds_swizzle_b32` through the LDS permute crossbar. On the latency-bound +// register-shuffle chains here (Phase-A sort tails, Phase-B merge tails, argmax +// butterfly) that shaves the crossbar issue off the two shortest-stride stages of +// every chain. quad_perm ctrl selects, per output lane j in a 4-lane quad, which +// quad lane (ctrl>>(2*j))&3 to read; xor^1 -> [1,0,3,2] = 0xB1, xor^2 -> +// [2,3,0,1] = 0x4E. All lanes are launched/active (out-of-range lanes seed -inf), +// so the gather is exact regardless of bound_ctrl. Masks 4/8/16 are not intra-quad +// and stay on ds_swizzle; the wave64 j==32 tail still falls back to __shfl_xor. +# define TOPK_DPP_QUAD_XOR1 0xB1 +# define TOPK_DPP_QUAD_XOR2 0x4E +// The ^16 xor is the widest intra-wave stride and the ENTRY stage of every +// 32-lane shuffle chain here (the merge/sort tails run j = 16,8,4,2,1), so it +// sits on the exposed critical path. It swaps the two 16-lane halves of the +// wave — exactly what `v_permlanex16_b32` does in its identity-cross form: with +// selectors (0x76543210, 0xFEDCBA98) output lane j reads the opposite half's +// same-index lane (j<16 -> lane j+16 = j^16; j>=16 -> lane j-16 = j^16). Unlike +// `ds_swizzle_b32`, which issues through the LDS permute crossbar, permlanex16 +// is a pure VALU cross-lane gather (fused into the ALU pipe) — the same +// "off the LDS crossbar onto VALU" move that DPP quad_perm won for ^1/^2, now +// applied to the largest intra-wave stride. bound_ctrl/fi are false; every lane +// is written by the full permutation so the passthrough `old` operand is dead. +// Verified bit-identical to `__shfl_xor(v, 16)` on gfx1201 and gfx1151. Masks 4/8 stay on +// ds_swizzle; the wave64 j==32 tail still falls back to __shfl_xor. static __device__ __forceinline__ int topk_shfl_xor_i32(int v, int mask) { switch (mask) { case 1: - return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(1)); + return __builtin_amdgcn_mov_dpp(v, TOPK_DPP_QUAD_XOR1, 0xf, 0xf, false); case 2: - return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(2)); + return __builtin_amdgcn_mov_dpp(v, TOPK_DPP_QUAD_XOR2, 0xf, 0xf, false); case 4: return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(4)); case 8: return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(8)); case 16: - return __builtin_amdgcn_ds_swizzle(v, TOPK_DS_SWIZZLE_XOR(16)); + return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xFEDCBA98u, false, false); default: return __shfl_xor(v, mask); } @@ -121,8 +145,19 @@ static __device__ __forceinline__ float topk_shfl_xor(float v, int mask) { // strided shared-memory round-trip. A direct wave-shuffle reduction replaces that // with two barrier-free shuffle passes and a single __syncthreads(). Threads are // rounded up to a whole number of waves so every lane in every wave participates -// in the shuffles (out-of-range lanes seed -inf and never win). Ties break toward -// the lower index purely to keep the reduction deterministic. +// in the shuffles (out-of-range lanes seed -inf and never win). +// +// The compare-select carries no index tiebreak: on a tie we keep the current +// lane's (v, id). Because (v, id) only ever advances to a partner when its value +// is *strictly* greater, the invariant x_row[id] == v holds at every stage, so +// after the butterfly every lane's v is the wave max (value-max is commutative / +// associative regardless of tie policy) and id is an index achieving it. Only +// lane 0's id is published, and it is always a valid argmax — which is all TOP_K +// with k == 1 needs (the op accepts any index holding the max; the test compares +// the values behind the indices, not the indices themselves). Dropping the +// `pv == v && pi < id` term shortens the dependent compare on each of this +// latency-bound reduction's shuffle stages; the only property lost is which of +// several tied maxima is returned, which the op leaves unspecified. static __global__ void k_topk_argmax(const float * x, int * dst, const int ncols) { const int col = threadIdx.x; const int row = blockIdx.x; @@ -136,7 +171,7 @@ static __global__ void k_topk_argmax(const float * x, int * dst, const int ncols for (int j = warpSize >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(v, j); const int pi = TOPK_SHFL_XOR(id, j); - if (pv > v || (pv == v && pi < id)) { + if (pv > v) { v = pv; id = pi; } @@ -164,7 +199,7 @@ static __global__ void k_topk_argmax(const float * x, int * dst, const int ncols for (int j = warpSize >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(v, j); const int pi = TOPK_SHFL_XOR(id, j); - if (pv > v || (pv == v && pi < id)) { + if (pv > v) { v = pv; id = pi; } @@ -220,8 +255,19 @@ static __global__ void k_topk_bitonic(const float * x, const float * x_row = x + (size_t) row * ncols; extern __shared__ int smem[]; - int * idx_row = smem; - float * val_row = (float *) (smem + ncols_pad); + // Phase A double-buffers the cross-wave LDS window (iter-51): each cross-wave + // `len` uses one of two ncols_pad-wide (idx,val) regions, and the next len + // (or the Phase-A publish) writes the OTHER region. The pure-WAR barrier that + // used to separate a len's peel reads from the next len's writeback is thus + // replaced by buffer separation and dropped. Consecutive uses of the SAME + // region are always >= one intervening len apart (that len's own barriers + // order them), so reuse stays race-free with no wave-lockstep assumption. + int * idx_buf[2] = { smem, smem + ncols_pad }; + float * val_buf[2] = { (float *) (smem + 2 * ncols_pad), + (float *) (smem + 3 * ncols_pad) }; + int bsel = 0; + int * idx_row = idx_buf[0]; + float * val_row = val_buf[0]; // Phase A: sort every contiguous KPAD-block descending (blocks are // independent — partners of a compare-exchange stay within the block). @@ -231,8 +277,11 @@ static __global__ void k_topk_bitonic(const float * x, // that whole branch is dead-code-eliminated for the cross-wave `len`s. float v = col < ncols ? x_row[col] : -INFINITY; int id = col; + // The final (widest) sort len == KPAD is peeled below so its middle + // cross-wave strides can use the iter-53 group-merge fuse; this loop runs + // the lower lens (len < KPAD) with the stock per-stride structure. #pragma unroll - for (int len = 2; len <= KPAD; len <<= 1) { + for (int len = 2; len < KPAD; len <<= 1) { if (len <= WARP) { // Fully intra-wave sort stage: register-to-register xor shuffle. #pragma unroll @@ -252,13 +301,22 @@ static __global__ void k_topk_bitonic(const float * x, } } else { // Cross-wave stages (stride >= WARP): exchange through LDS + barrier. - // Flush the register-resident (v, id) to LDS first so partner lanes - // read the current values. + // Pick this len's LDS region (iter-51 double-buffer) and flush the + // register-resident (v, id) to it so partner lanes read the current + // values. The previous cross-wave len read the OTHER region, so this + // writeback cannot clobber its still-in-flight peel reads. + val_row = val_buf[bsel]; + idx_row = idx_buf[bsel]; val_row[col] = v; idx_row[col] = id; __syncthreads(); + // Cross-wave strides via in-place LDS swap, down to but NOT including + // j == WARP (peeled below as a register merge, iter-47). The final + // barrier of this loop (or the writeback barrier above when the loop + // runs zero iterations, i.e. len == 2*WARP) publishes the sibling-wave + // slots the register merge reads. #pragma unroll - for (int j = len >> 1; j >= WARP; j >>= 1) { + for (int j = len >> 1; j > WARP; j >>= 1) { const int partner = col ^ j; if (partner > col) { // Descending target for this comparator (uniform per block at @@ -273,8 +331,45 @@ static __global__ void k_topk_bitonic(const float * x, } __syncthreads(); } - v = val_row[col]; - id = idx_row[col]; + + // Peeled last cross-wave stride (j == WARP) as a register merge + // (iter-47, the iter-46 Phase-B peel applied to every Phase-A + // cross-wave len). partner == col ^ WARP is in the sibling wave; its + // slot was published by the barrier above. Instead of the stock + // "lower lane swaps both slots in LDS, then BOTH lanes reload + // val_row[col]/idx_row[col]", each lane reads its own slot + its + // partner's and computes its OWN result in registers: for the pair + // (low = col&WARP==0, high) with the alternating bitonic direction + // `up`, the low lane keeps the max and the high the min when up (and + // the reverse when !up) — byte-for-byte what the stock conditional + // swap produced. This drops the swap's conditional LDS writes AND the + // full own-slot reload from the critical path. `up` is uniform across + // the WARP-pair (they differ only in bit WARP < len, which does not + // affect col & len), so both lanes agree on the direction. + // + // Barrier count is UNCHANGED: the __syncthreads() the stock loop ran + // after the j == WARP swap is retained here as a WAR barrier — the + // register merge is the last LDS reader of this len, and the next + // len's writeback (or Phase A's publish writeback after len == KPAD) + // overwrites val_row[col], which the partner lane read here. Keeping + // that barrier makes the elision race-free (no reliance on wave + // lockstep), so this is a pure LDS-traffic cut, not a barrier cut. + { + const int partner = col ^ WARP; + const float sv = val_row[col]; + const float pv = val_row[partner]; + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool low = (col & WARP) == 0; + const bool take = (up == low) ? (sv < pv) : (sv > pv); + v = take ? pv : sv; + id = take ? idx_row[partner] : idx_row[col]; + } + // iter-51: the WAR __syncthreads() that used to sit here is DROPPED. + // The peel above read val_buf[bsel]; the next cross-wave len (or the + // Phase-A publish) writes val_buf[bsel ^ 1], so there is no + // write-after-read hazard on this region — buffer separation stands + // in for the barrier. Advance to the other region for the next len. + bsel ^= 1; // Intra-wave tail (stride < WARP): xor shuffle, no LDS/barrier. Both // lanes of a pair reconstruct the same (low, high) view and the same @@ -296,7 +391,130 @@ static __global__ void k_topk_bitonic(const float * x, } } } + + // Peeled final Phase-A sort stage (len == KPAD), iter-53: apply the iter-52 + // Phase-B middle-cross-wave-stride group-merge fuse to Phase A. This is the + // widest Phase-A sort len and is ALWAYS cross-wave on this path (KPAD > + // 2*WARP), and it carries the most middle strides (j = KPAD/2 .. 2*WARP), so + // it has the most barriers to collapse. The middle strides pair columns that + // differ only in bits [log2(2*WARP) .. log2(KPAD/2)], so the MID_A = + // KPAD/(2*WARP) columns { base + t*2*WARP } form a group CLOSED under all of + // them. When MID_A >= 4 one owner lane per group (i & GMASK == 0) gathers its + // MID_A slots into registers, runs the whole sub-sort there (group-index + // strides MID_A/2..1 == full strides KPAD/2..2*WARP), and writes the results + // back — replacing the M per-stride block barriers AND the M-1 intermediate + // LDS write->readback round-trips with ONE barrier (same fuse class as + // iter-50/52). For k=400 (KPAD=512, strides {256,128,64}) that is 3 barriers + // -> 1. + // + // For len == KPAD the sort direction is UNIFORMLY descending: up = ((col & + // (KPAD-1)) & KPAD) == 0 is always true (bit log2(KPAD) is masked out by + // KPAD-1), so the group merge is a plain descending bitonic sub-sort with the + // exact comparator iter-52 uses (low index keeps the max, swap only on strict + // <; idx follows value). Race-free with no internal barrier: each owner + // exclusively owns its MID_A group slots (disjoint groups tile the KPAD-block + // in steps of 2*WARP) and reads ONLY those slots — all published by the + // writeback barrier below, none written by any other lane this stage; reads + // precede writes intra-thread. Bit-identical to the stock in-place loop. For + // MID_A < 4 (KPAD == 2*WARP: 0-or-1 middle stride, no barrier to save) keep + // the stock in-place loop byte-for-byte. + { + const int len = KPAD; + val_row = val_buf[bsel]; + idx_row = idx_buf[bsel]; + val_row[col] = v; + idx_row[col] = id; + __syncthreads(); + + constexpr int MID_A = KPAD / (2 * WARP); + if (MID_A >= 4) { + constexpr int GMASK = (MID_A - 1) * (2 * WARP); + constexpr int GARR = MID_A >= 4 ? MID_A : 1; // avoid 0-length array in dead insts + const int i = col & (KPAD - 1); + const int gbase = col - i; + if ((i & GMASK) == 0) { + float gv[GARR]; + int gi[GARR]; +#pragma unroll + for (int t = 0; t < MID_A; ++t) { + const int p = gbase + i + t * (2 * WARP); + gv[t] = val_row[p]; + gi[t] = idx_row[p]; + } +#pragma unroll + for (int s = MID_A >> 1; s > 0; s >>= 1) { +#pragma unroll + for (int a = 0; a < MID_A; ++a) { + const int b = a ^ s; + if (b > a && gv[a] < gv[b]) { + topk_swap(gv[a], gv[b]); + topk_swap(gi[a], gi[b]); + } + } + } +#pragma unroll + for (int t = 0; t < MID_A; ++t) { + const int p = gbase + i + t * (2 * WARP); + val_row[p] = gv[t]; + idx_row[p] = gi[t]; + } + } + __syncthreads(); + } else { +#pragma unroll + for (int j = len >> 1; j > WARP; j >>= 1) { + const int partner = col ^ j; + if (partner > col) { + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool do_swap = + up ? (val_row[col] < val_row[partner]) : (val_row[col] > val_row[partner]); + if (do_swap) { + topk_swap(val_row[col], val_row[partner]); + topk_swap(idx_row[col], idx_row[partner]); + } + } + __syncthreads(); + } + } + + // Peeled j == WARP register merge (iter-47), unchanged from the loop body. + { + const int partner = col ^ WARP; + const float sv = val_row[col]; + const float pv = val_row[partner]; + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool low = (col & WARP) == 0; + const bool take = (up == low) ? (sv < pv) : (sv > pv); + v = take ? pv : sv; + id = take ? idx_row[partner] : idx_row[col]; + } + bsel ^= 1; + + // Intra-wave tail (stride < WARP), unchanged from the loop body. +#pragma unroll + for (int j = WARP >> 1; j > 0; j >>= 1) { + const float pv = TOPK_SHFL_XOR(v, j); + const int pi = TOPK_SHFL_XOR(id, j); + + const bool up = (((col & (KPAD - 1)) & len) == 0); + const bool low = (col & j) == 0; + const float low_val = low ? v : pv; + const float high_val = low ? pv : v; + const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + if (do_swap) { + v = pv; + id = pi; + } + } + } // Publish the sorted KPAD-blocks to LDS for Phase B's cross-wave merges. + // The last Phase-A len read val_buf[bsel ^ 1]; the trailing `bsel ^= 1` + // advanced to val_buf[bsel], so publishing here writes the OTHER region and + // does not clobber that len's still-in-flight peel reads (buffer separation + // again replaces the dropped WAR barrier). Phase B then operates in this + // region for the rest of the kernel. + val_row = val_buf[bsel]; + idx_row = idx_buf[bsel]; val_row[col] = v; idx_row[col] = id; __syncthreads(); @@ -318,51 +536,182 @@ static __global__ void k_topk_bitonic(const float * x, // register-resident merge result and skip the LDS round-trip entirely. const bool last = (span << 1) >= ncols_pad; - if (active) { - // C[i] = max(X[i], Y[KPAD-1-i]) is bitonic and holds the KPAD largest. - // The pre-writeback barrier is unnecessary: each active lane reads its - // own lower-half slot a_pos = gbase+i plus an upper-half slot b_pos in - // [gbase+span, gbase+span+KPAD) (never written this level since - // KPAD <= span), and writes back only to gbase+i. The write set - // [gbase, gbase+KPAD) is disjoint from the cross-half read set and the - // a_pos read-before-write is intra-thread ordered, so only the - // post-writeback barrier is needed. - const int a_pos = gbase + i; - const int b_pos = gbase + span + (KPAD - 1 - i); - const float av = val_row[a_pos]; - const float bv = val_row[b_pos]; - if (av >= bv) { - val_row[gbase + i] = av; - idx_row[gbase + i] = idx_row[a_pos]; + // Fused C-step + first cross-wave merge stride (j == KPAD/2), iter-50. + // The stock code ran the C-step (C[i] = max(X[i], Y[KPAD-1-i]) written to + // gbase+i for every i < KPAD), a block __syncthreads(), then the first + // cross-wave stride j == KPAD/2 as an in-place LDS compare-exchange. But the + // C-step's per-slot output and the j == KPAD/2 compare only couple WITHIN a + // single (i, i+KPAD/2) pair, so the two stages fuse: the low lane of each + // pair (local i < KPAD/2) computes BOTH its C-values itself, straight from + // the still-pristine published input blocks, and emits the descending-merge + // result directly (max -> gbase+i, min -> gbase+i+KPAD/2). This removes one + // full block __syncthreads() AND the j == KPAD/2 stage's LDS re-read+re-write + // round-trip per Phase-B level from the barrier-bound k>2*warp shared path. + // + // Race-free without an internal barrier: each low lane exclusively owns + // outputs gbase+i and gbase+i+KPAD/2 (no other lane writes them); its two + // cross-block reads b0/b1 land in [gbase+span, gbase+span+KPAD), never + // written this level (KPAD <= span); its two low-half reads a0/a1 are its own + // outputs (intra-thread read-before-write). Bit-identical to the stock + // C-step (a wins ties: av >= bv) + stock stride (low keeps its value on ties: + // swap only when strictly <, i.e. keep max with cav >= cbv). + const int half = KPAD >> 1; + if (i < half) { + const int a0 = gbase + i; + const int b0 = gbase + span + (KPAD - 1 - i); + const int a1 = gbase + i + half; + const int b1 = gbase + span + (half - 1 - i); + const float a0v = val_row[a0]; + const float b0v = val_row[b0]; + const float a1v = val_row[a1]; + const float b1v = val_row[b1]; + float cav; + int cai; + if (a0v >= b0v) { cav = a0v; cai = idx_row[a0]; } else { cav = b0v; cai = idx_row[b0]; } + float cbv; + int cbi; + if (a1v >= b1v) { cbv = a1v; cbi = idx_row[a1]; } else { cbv = b1v; cbi = idx_row[b1]; } + if (cav >= cbv) { + val_row[a0] = cav; idx_row[a0] = cai; + val_row[a1] = cbv; idx_row[a1] = cbi; } else { - val_row[gbase + i] = bv; - idx_row[gbase + i] = idx_row[b_pos]; + val_row[a0] = cbv; idx_row[a0] = cbi; + val_row[a1] = cav; idx_row[a1] = cai; } } __syncthreads(); - // Cross-wave merge stages via shared memory (j = KPAD/2 .. WARP). + // Remaining cross-wave merge strides (j = KPAD/4 .. 2*WARP; j == KPAD/2 was + // fused above; j == WARP is peeled below as a register merge, so it is off + // this range — guard j > WARP, not >=). + // + // iter-52: these strides all pair columns that differ only in the bits + // [log2(2*WARP) .. log2(KPAD/4)], so the G = KPAD/(4*WARP) columns + // { gbase+i, gbase+i+2W, ..., gbase+i+(G-1)*2W } form a group CLOSED under + // every one of them. When there are >= 2 middle strides (G >= 4, i.e. KPAD in + // {512,1024}) one owner lane per group (the low element, i & GMASK == 0) + // gathers its G slots into registers, runs the whole sub-merge there + // (group-index strides G/2..1 == full strides KPAD/4..2W), and writes the G + // results back. This replaces the M per-stride block barriers AND the M-1 + // intermediate LDS write->readback round-trips with ONE barrier and zero + // intermediate round-trips (same fuse class as iter-50, now spanning ALL the + // middle strides at once). For k=400 (KPAD=512, strides {128,64}) that is + // 2 barriers -> 1. + // + // Race-free with NO internal barrier: each owner exclusively owns its G + // output slots (disjoint groups; GMASK-aligned bases tile [gbase,gbase+KPAD) + // in steps of 2W), and it reads ONLY those same slots — all written by the + // fused C-step above (published by the barrier at the top of this block), none + // written by any other lane in THIS stage. Reads precede writes intra-thread. + // Bit-exact to the stock loop: identical comparator (low index keeps the max, + // swap only on strict <) applied in identical stride order; idx follows value. + // For G <= 2 (KPAD in {128,256}: 0 or 1 middle stride, so no barrier to save) + // keep the stock in-place loop byte-for-byte. + constexpr int MID_G = KPAD / (4 * WARP); + if (MID_G >= 4) { + constexpr int GMASK = (MID_G - 1) * (2 * WARP); + constexpr int GARR = MID_G >= 4 ? MID_G : 1; // avoid 0-length array in the dead (MID_G<4) instantiations + if (active && (i & GMASK) == 0) { + float gv[GARR]; + int gi[GARR]; #pragma unroll - for (int j = KPAD >> 1; j >= WARP; j >>= 1) { - const int partner = col ^ j; - if (active && partner > col) { - if (val_row[col] < val_row[partner]) { - topk_swap(val_row[col], val_row[partner]); - topk_swap(idx_row[col], idx_row[partner]); + for (int t = 0; t < MID_G; ++t) { + const int p = gbase + i + t * (2 * WARP); + gv[t] = val_row[p]; + gi[t] = idx_row[p]; + } +#pragma unroll + for (int s = MID_G >> 1; s > 0; s >>= 1) { +#pragma unroll + for (int a = 0; a < MID_G; ++a) { + const int b = a ^ s; + if (b > a && gv[a] < gv[b]) { + topk_swap(gv[a], gv[b]); + topk_swap(gi[a], gi[b]); + } + } + } +#pragma unroll + for (int t = 0; t < MID_G; ++t) { + const int p = gbase + i + t * (2 * WARP); + val_row[p] = gv[t]; + idx_row[p] = gi[t]; } } __syncthreads(); + } else { +#pragma unroll + for (int j = KPAD >> 2; j > WARP; j >>= 1) { + const int partner = col ^ j; + if (active && partner > col) { + if (val_row[col] < val_row[partner]) { + topk_swap(val_row[col], val_row[partner]); + topk_swap(idx_row[col], idx_row[partner]); + } + } + __syncthreads(); + } } - // Intra-wave merge tail (j = WARP/2 .. 1) via xor shuffle. Merge partners - // (col ^ j, j < KPAD) stay within the aligned head KPAD-block, so an - // active lane's partner is also active; inactive lanes shuffle harmlessly - // but never write back. Uniform descending merge: low lane keeps the max. + // Fused last cross-wave stride (j == WARP) + intra-wave merge tail + // (j = WARP/2 .. 1) via xor shuffle. Merge partners (col ^ j, j < KPAD) + // stay within the aligned head KPAD-block, so an active lane's partner is + // also active; inactive lanes shuffle harmlessly but never write back. + // Uniform descending merge: low lane keeps the max. + // + // j == WARP is the boundary between cross-wave (partner in a sibling wave, + // via LDS) and intra-wave (via xor shuffle) strides. The stock loop above + // did it as an in-place LDS swap (lower lane reads+writes both slots) gated + // by a __syncthreads(), then the tail RELOADED val_row[col] to start the + // shuffle chain. Instead, each lane reads its own slot + its sibling-wave + // partner's slot (both synced by the last cross-wave barrier) and computes + // its OWN merge result in registers: low lane keeps the max, high lane the + // min — byte-for-byte what the swap produced (low slot <- max, high <- min). + // No lane writes LDS during this stride, so the barrier after it and the + // tail's reload both vanish: the whole tail runs register-resident straight + // off the register merge. On the barrier-bound k=400 shared path (jstop=16 + // → the tail is a single j=16 stage) this removes one full block barrier + // and the LDS write+readback round-trip per Phase-B level — the same + // barrier/round-trip elision class iter-44 won with, now on the peeled + // cross-wave/intra-wave boundary stage. + // + // k-aligned early termination on the OUTPUT (last) level: a bitonic merge + // completed down to stride s leaves the window rank-sorted into size-s + // blocks (block m holds ranks [m*s, (m+1)*s)); the smaller-stride stages + // only reorder WITHIN a block. TOP_K's result is an unordered SET, so once + // k is a multiple of the current stride s, positions [0, k) already hold + // exactly the top-k set and the j < s stages are pure dead reordering of + // in-set / out-of-set elements. Stop at jstop = the largest power of two + // dividing k (k & -k). NON-last levels must stay fully sorted (their block + // feeds the next level's C-step, which requires descending order), so jstop + // is forced to 1 there. jstop is uniform across the block (k / last are + // uniform), so the guard is a divergence-free branch and every lane runs + // the surviving shuffle stages in lockstep. For the benchmark's k=400 case + // jstop=16, cutting the last-level tail from 5 shuffle stages to 1. { - float mv = val_row[col]; - int mid = idx_row[col]; + // j == WARP cross-wave merge, register-resident (no LDS write, no + // barrier). partner == col ^ WARP is in the sibling wave; its slot was + // synced by the last cross-wave barrier above. low lane keeps the max, + // high lane the min (uniform descending) — identical to the swap the + // stock loop performed. idx_row[partner] is read only when the partner + // wins (lazy cross-wave index read). + float mv; + int mid; + { + const int partner = col ^ WARP; + const float sv = val_row[col]; + const float pv = val_row[partner]; + const bool low = (col & WARP) == 0; + const bool take = low ? (sv < pv) : (sv > pv); + mv = take ? pv : sv; + mid = take ? idx_row[partner] : idx_row[col]; + } + const int jstop = last ? (k & (-k)) : 1; #pragma unroll for (int j = WARP >> 1; j > 0; j >>= 1) { + if (j < jstop) { + continue; + } const float pv = TOPK_SHFL_XOR(mv, j); const int pi = TOPK_SHFL_XOR(mid, j); @@ -479,25 +828,42 @@ static __global__ void k_topk_bitonic_smallk(const float * x, const int gbase = col - i; const bool active = i < KPAD; + // C-step: C[i] = max(X[i], Y[KPAD-1-i]). Keep the winner in registers and + // feed the intra-wave merge below directly, instead of storing it to + // val_row[gbase+i] and immediately reloading val_row[col]. On this path the + // merge is fully intra-wave xor shuffle (KPAD <= warpSize) and touches no + // LDS, and no other lane reads gbase+i before the post-merge writeback + // overwrites it (merge partners col^j, jmerge chain. + // Inactive lanes (i >= KPAD) keep their prior registers; their shuffle + // outputs never feed an active lane (an active lane's partner is always + // active), so their contents are irrelevant. if (active) { - const int a_pos = gbase + i; + // a_pos == gbase + i == col: the lane's OWN slot, which still holds the + // register-resident (v, id). Every active lane wrote val_row[col]=v / + // idx_row[col]=id at the previous level's writeback (Phase A for the + // first level) and has not mutated v/id since — the merge below feeds + // straight back into (v, id) (iter-33), and the Phase-B active set only + // shrinks (col mod 4span >= col mod 2span, so once inactive always + // inactive), so an active lane was active at every prior level and its + // own-slot writeback chain kept val_row[col] === v. Only the cross-half + // b-side lives in another (inactive) lane's slot, so only it is loaded + // from LDS; the a-side reload of the lane's own just-written value was a + // dead same-slot LDS round-trip sitting on the latency-bound + // C-step->merge chain. Tie policy is unchanged: b wins only on a strict + // bv > v (== the old `av >= bv` keeping a on ties). const int b_pos = gbase + span + (KPAD - 1 - i); - const float av = val_row[a_pos]; const float bv = val_row[b_pos]; - if (av >= bv) { - val_row[gbase + i] = av; - idx_row[gbase + i] = idx_row[a_pos]; - } else { - val_row[gbase + i] = bv; - idx_row[gbase + i] = idx_row[b_pos]; + if (bv > v) { + v = bv; + id = idx_row[b_pos]; } } // Bitonic merge the KPAD-wide block back to descending order. KPAD <= - // warpSize, so the whole merge is intra-wave xor shuffle. Each active lane - // loads only its own C-step write (no barrier needed, see above). - v = val_row[col]; - id = idx_row[col]; + // warpSize, so the whole merge is intra-wave xor shuffle over the + // register-resident C-step winner (no LDS reload, no barrier). #pragma unroll for (int j = KPAD >> 1; j > 0; j >>= 1) { const float pv = TOPK_SHFL_XOR(v, j); @@ -590,12 +956,17 @@ static __global__ void k_topk_bitonic_2wave(const float * x, __syncthreads(); const int partner = col ^ j; const float pv = val_row[partner]; - const int pi = idx_row[partner]; - __syncthreads(); // all cross-wave reads done before LDS reuse + // Decide the exchange from the value read (pv) — no barrier needed + // for a register compare — so the partner-index LDS read below is + // issued only on the lanes that actually swap. On a non-swapping + // lane (~half, per bitonic stage) idx_row[partner] is never touched, + // trimming one ds_read from that lane's cross-wave critical path. const bool low = (col & j) == 0; const float low_val = low ? v : pv; const float high_val = low ? pv : v; const bool do_swap = up ? (low_val < high_val) : (low_val > high_val); + const int pi = do_swap ? idx_row[partner] : id; + __syncthreads(); // all cross-wave reads done before LDS reuse if (do_swap) { v = pv; id = pi; @@ -616,7 +987,14 @@ static __global__ void k_topk_bitonic_2wave(const float * x, } } } - __syncthreads(); // ensure the last cross-wave LDS reads finished before overwrite + // No barrier needed before this writeback: Phase A has exactly one cross-wave + // exchange (len == KPAD, j == WARP), and its post-read barrier (the second + // __syncthreads() in the j >= WARP branch above) already guarantees every + // block-wide LDS read of that exchange has completed. Every stage after it + // (j = WARP/2 .. 1) is an intra-wave register xor shuffle that never touches + // LDS, so nothing can still be reading val_row/idx_row when we overwrite each + // lane's own slot here. Only the post-writeback barrier is required (to order + // the writeback ahead of Phase B's first C-step read). val_row[col] = v; idx_row[col] = id; __syncthreads(); @@ -629,17 +1007,30 @@ static __global__ void k_topk_bitonic_2wave(const float * x, const bool active = i < KPAD; if (active) { - const int a_pos = gbase + i; + // a-side reload elision (iter-36, the iter-35 lever applied to 2wave): + // a_pos == gbase + i == col is the lane's OWN slot, and the registers + // (v, id) already mirror val_row[col]. After Phase A every lane wrote + // val_row[col] = v; at each Phase-B level the merge-tail writeback below + // stores (mv, mid) to val_row[col] AND refreshes (v, id) = (mv, mid), and + // the trailing __syncthreads() is the only thing between that store and + // this C-step read — with no mutation of v/id in between — so + // val_row[col] === v holds here. The Phase-B active set only SHRINKS + // (active ⇔ (col mod 2span) < KPAD, and col mod 4span >= col mod 2span, so + // once inactive always inactive), hence an active lane was active at every + // prior level and its own-slot mirror chain is unbroken. Only the + // cross-half b-side lives in another lane's slot, so only it is loaded + // from LDS; the a-side own-slot reload was a dead round-trip on the + // latency-bound C-step->merge chain. The C-step OUTPUT store stays: the + // cross-wave stage below reads val_row[col] via the sibling (partner) + // lane. Tie policy unchanged: `bv > v` keeps a on ties (== old av >= bv). const int b_pos = gbase + span + (KPAD - 1 - i); - const float av = val_row[a_pos]; const float bv = val_row[b_pos]; - if (av >= bv) { - val_row[gbase + i] = av; - idx_row[gbase + i] = idx_row[a_pos]; - } else { - val_row[gbase + i] = bv; - idx_row[gbase + i] = idx_row[b_pos]; + if (bv > v) { + v = bv; + id = idx_row[b_pos]; } + val_row[gbase + i] = v; + idx_row[gbase + i] = id; } __syncthreads(); @@ -675,6 +1066,10 @@ static __global__ void k_topk_bitonic_2wave(const float * x, if (active) { val_row[col] = mv; idx_row[col] = mid; + // Refresh the own-slot register mirror so next level's C-step can + // take its a-side from (v, id) instead of reloading val_row[col]. + v = mv; + id = mid; } __syncthreads(); } @@ -730,7 +1125,9 @@ static void topk_bitonic_cuda(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) + sizeof(float)); + // Double-buffered Phase-A LDS window (iter-51): two ncols_pad-wide (idx,val) + // regions instead of one. ncols <= 1024 on this path, so this is <= 16 KB. + const size_t shared_mem = 2 * ncols_pad * (sizeof(int) + sizeof(float)); GGML_ASSERT(shared_mem <= ggml_cuda_info().devices[ggml_cuda_get_device()].smpb);