Commit 57d04820d for llama.cpp
commit 57d04820dfef7a3519aac276d018a6fd71251895
Author: Alexey Kopytko <alexey@kopytko.com>
Date: Sun Nov 8 23:42:57 2026 +0900
ggml : update CUDA/Metal radix Top-K to make tied selections deterministic (#31084)
* ggml : fix for the radix TOP_K thread timing bug
* make the key an ordered 64-bit pair
* a rewrite without shift-skip logic
* fix casts
* warp-contiguous layout
* reduce the scope removing vulkan+metal changes
* turn it off for Metal/ROCm, note about cub::DeviceTopK
* fix the hip build
* skip Vulkan as well
* popc lambda from fattn
* do block_cols math in int64_t
* make it allowlist bc SYCL looks non deterministic too
* ggml : make Metal radix Top-K tied selections deterministic
The radix top-k kernel gathered the selected indices with atomic
counters, so the order of the tied candidates at the threshold (and,
with more ties than remaining slots, the subset itself) depended on
the thread scheduling.
Replace the two atomic compaction loops with a deterministic scheme:
count the elements above and tied with the threshold in contiguous
per-thread chunks, rank them with a simdgroup scan and a serial scan
over the (now unused) radix histogram, and write each kept element
exactly once in column order. The output is now byte-identical
across runs.
This is also faster than the atomic version: the strided atomic
loops become contiguous passes, reducing top-k times by ~20-45% on
M5 Ultra, with tie-heavy shapes within a few percent.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* tests : enable deterministic top-k check on Metal, add radix-path tie cases
Metal now produces deterministic tied selections for the radix top-k
path, so the determinism self-check runs for the MTL backend as well.
Add eval cases covering the radix path (ncols > 2048) with more ties
at the threshold than remaining slots, and tie-heavy perf cases to
monitor the cost of the deterministic tie handling.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* ggml : use coalesced accesses in the Metal radix Top-K compaction
The deterministic compaction walked contiguous per-thread chunks with
scalar sequential loads, which the memory system of the M2 Ultra can
not overlap well: the radix Top-K times regressed by ~1.5-2.4x on the
large-ncols shapes.
Give each warp a contiguous chunk of the row instead and let the lanes
interleave within it with stride 32, so every access is fully
coalesced. The lanes write their kept elements at lane-major ranks,
which is deterministic, and only a tie overflow falls back to a
per-step cross-lane ranking to keep the lowest tied indices in column
order.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* tests : remove the extra top-k tie and perf cases
The radix-path eval and tie-heavy perf cases were added to validate and
benchmark the Metal deterministic tie handling. Keep only the enabling
of the determinism self-check for the MTL backend.
Assisted-by: dsh:llama.cpp/ggml-org/GLM-5.3-Flash-GGUF (Mac Studio, M5 Ultra, 256GB)
* metal: scan the radix bins and warp totals with a simdgroup
The radix top-k selected the threshold bucket with a serial scan of
the 256 bins by thread 0 in each of the 4 passes, and ranked the warps
of the compaction with two serial scans by thread 0. The first
simdgroup now does both: each lane sums 8 bins, an exclusive prefix
sum and a ballot find the lane that reaches the k-th value, and the
warp totals get their exclusive bases from simd_prefix_exclusive_sum.
* tests: enable the deterministic top-k check on ROCm
HIP has no cub::DeviceTopK, so every row goes through the radix or
bitonic path, which break ties in the same way in each run. The check
covers all tie cases on ROCm, including the 1 and 2 row ones that CUDA
skips.
* cuda: remove the redundant radix top-k counter reset
Address review from ORippler. top_k_radix_init already zeroes
greater_count for each chunk of rows and top_k_radix_select keeps it,
so the reset kernel before the gather is one launch with no effect.
---------
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
Co-authored-by: Pascal <admin@serveurperso.com>
diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu
index b0f609990..19cb45d1f 100644
--- a/ggml/src/ggml-cuda/top-k.cu
+++ b/ggml/src/ggml-cuda/top-k.cu
@@ -88,18 +88,18 @@ static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) {
return bits ^ mask;
}
-struct top_k_radix_state {
+// one state per cache line, so the gather atomics of different rows do not compete
+struct alignas(128) top_k_radix_state {
uint32_t prefix;
uint32_t prefix_mask;
int rank;
int greater_count;
- int equal_count;
};
static __global__ void top_k_radix_init(top_k_radix_state * states, int nrows, int k) {
const int row = blockIdx.x * blockDim.x + threadIdx.x;
if (row < nrows) {
- states[row] = {0, 0, k, 0, 0};
+ states[row] = {0, 0, k, 0};
}
}
@@ -109,6 +109,7 @@ static __global__ void top_k_radix_histogram(
const top_k_radix_state * __restrict__ states,
int * __restrict__ block_histograms,
int ncols,
+ int block_cols,
int blocks_per_row,
int shift) {
constexpr int NBINS = 1 << RADIX_BITS;
@@ -117,15 +118,15 @@ static __global__ void top_k_radix_histogram(
const int row_block = blockIdx.x % blocks_per_row;
const int tid = threadIdx.x;
const float * row_src = src + (size_t) row * ncols;
+ const int64_t col_begin = (int64_t) row_block * block_cols;
+ const int64_t col_end = min(col_begin + block_cols, (int64_t) ncols);
__shared__ int histogram[NBINS];
histogram[tid] = 0;
__syncthreads();
const top_k_radix_state state = states[row];
- for (int64_t col = row_block * BLOCK_SIZE + tid;
- col < ncols;
- col += blocks_per_row * BLOCK_SIZE) {
+ for (int64_t col = col_begin + tid; col < col_end; col += BLOCK_SIZE) {
const uint32_t key = top_k_float_to_ordered(row_src[col]);
if ((key & state.prefix_mask) == state.prefix) {
atomicAdd(&histogram[(key >> shift) & (NBINS - 1)], 1);
@@ -170,42 +171,90 @@ static __global__ void top_k_radix_select(
}
}
-static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, int nrows) {
- const int row = blockIdx.x * blockDim.x + threadIdx.x;
- if (row < nrows) {
- states[row].greater_count = 0;
- states[row].equal_count = 0;
- }
-}
-
-template<int BLOCK_SIZE>
+template<int BLOCK_SIZE, int RADIX_BITS>
static __global__ void top_k_radix_gather(
const float * __restrict__ src,
int * __restrict__ dst,
+ const int * __restrict__ block_histograms,
top_k_radix_state * __restrict__ states,
int ncols,
+ int block_cols,
int k,
int blocks_per_row) {
+ constexpr int NBINS = 1 << RADIX_BITS;
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size();
+ constexpr int nwarps = BLOCK_SIZE / warp_size;
+ using lane_mask_t = std::conditional_t<warp_size == 64, uint64_t, uint32_t>;
+ // HIP requires a 64-bit mask also for a 32-wide warp
+ constexpr uint64_t ballot_mask = ~(lane_mask_t) 0;
+ const auto popc = [](const lane_mask_t mask) {
+ if constexpr (warp_size == 64) {
+ return __popcll(mask);
+ } else {
+ return __popc(mask);
+ }
+ };
+ __shared__ int warp_ties[nwarps];
+
const int row = blockIdx.x / blocks_per_row;
const int row_block = blockIdx.x % blocks_per_row;
- const int tid = threadIdx.x;
+ const int lane = threadIdx.x % warp_size;
+ const int warp = threadIdx.x / warp_size;
const float * row_src = src + (size_t) row * ncols;
int * row_dst = dst + (size_t) row * k;
top_k_radix_state * state = &states[row];
+ const uint32_t threshold = state->prefix;
+ const int rank = state->rank;
+
+ // each warp reads a contiguous part of the contiguous range of the block
+ const int warp_cols = block_cols / nwarps;
+ const int64_t warp_begin = (int64_t) row_block * block_cols + warp * warp_cols;
+ const int64_t warp_end = min(warp_begin + warp_cols, (int64_t) ncols);
- for (int64_t col = row_block * BLOCK_SIZE + tid;
- col < ncols;
- col += blocks_per_row * BLOCK_SIZE) {
+ int nties = 0;
+ for (int64_t col = warp_begin + lane; col < warp_end; col += warp_size) {
const uint32_t key = top_k_float_to_ordered(row_src[col]);
- if (key > state->prefix) {
+ if (key > threshold) {
const int pos = atomicAdd(&state->greater_count, 1);
row_dst[pos] = col;
- } else if (key == state->prefix) {
- const int pos = atomicAdd(&state->equal_count, 1);
- if (pos < state->rank) {
- row_dst[k - state->rank + pos] = col;
- }
+ } else if (key == threshold) {
+ nties++;
+ }
+ }
+
+ // the ties fill the last rank slots in column order, the lowest columns win.
+ // the histograms of the last pass count the ties of each block in the bin of the threshold.
+ const int * block_ties = block_histograms + (size_t) row * blocks_per_row * NBINS + (threshold & (NBINS - 1));
+ int pos = 0;
+ for (int b = lane; b < row_block; b += warp_size) {
+ pos += block_ties[b * NBINS];
+ }
+ pos = warp_reduce_sum<warp_size>(pos);
+ if (pos >= rank || block_ties[row_block * NBINS] == 0) {
+ return; // the same result in all threads of the block
+ }
+
+ nties = warp_reduce_sum<warp_size>(nties);
+ if (lane == 0) {
+ warp_ties[warp] = nties;
+ }
+ __syncthreads();
+ for (int w = 0; w < warp; ++w) {
+ pos += warp_ties[w];
+ }
+
+ // pos is the same in all lanes of the warp, so the lanes leave the loop together
+ const int pos_end = min(pos + nties, rank);
+ for (int64_t tile = warp_begin; tile < warp_end && pos < pos_end; tile += warp_size) {
+ const int64_t col = tile + lane;
+ const bool is_tie = col < warp_end && top_k_float_to_ordered(row_src[col]) == threshold;
+
+ const lane_mask_t ties = (lane_mask_t) __ballot_sync(ballot_mask, is_tie);
+ const int tie_pos = pos + popc(ties & (((lane_mask_t) 1 << lane) - 1));
+ if (is_tie && tie_pos < pos_end) {
+ row_dst[k - rank + tie_pos] = col;
}
+ pos += popc(ties);
}
}
@@ -216,6 +265,8 @@ static void top_k_radix_cuda(
constexpr int RADIX_BITS = 8;
constexpr int NBINS = 1 << RADIX_BITS;
const int blocks_per_row = (int) std::min<int64_t>(((int64_t) ncols + 1023) / 1024, 64);
+ // each block of a row reads a contiguous range of whole tiles
+ const int block_cols = (int) (((int64_t) ncols + blocks_per_row * BLOCK_SIZE - 1) / (blocks_per_row * BLOCK_SIZE) * BLOCK_SIZE);
// chunk the rows to bound the histogram memory to 64 MB
const int64_t chunk_nrows = ggml_cuda_chunk_nrows((size_t) blocks_per_row * NBINS * sizeof(int), nrows);
@@ -234,16 +285,14 @@ static void top_k_radix_cuda(
for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) {
top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS>
<<<row_grid, BLOCK_SIZE, 0, stream>>>(
- src, states, histograms, ncols, blocks_per_row, shift);
+ src, states, histograms, ncols, block_cols, blocks_per_row, shift);
top_k_radix_select<BLOCK_SIZE, RADIX_BITS>
<<<iter_nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift);
}
- top_k_radix_reset_counters
- <<<(iter_nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, iter_nrows);
- top_k_radix_gather<BLOCK_SIZE>
+ top_k_radix_gather<BLOCK_SIZE, RADIX_BITS>
<<<row_grid, BLOCK_SIZE, 0, stream>>>(
- src, dst, states, ncols, k, blocks_per_row);
+ src, dst, histograms, states, ncols, block_cols, k, blocks_per_row);
src += (size_t) ncols * iter_nrows;
dst += (size_t) k * iter_nrows;
diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp
index 3d317a0a6..5c93bc4f0 100644
--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp
+++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp
@@ -6008,11 +6008,10 @@ static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) {
/*.top_k =*/ (int32_t) op->ne[0],
};
- // shared memory: 256-entry histogram + bucket/above scalars + output counter
+ // shared memory: 256-entry histogram + bucket/above scalars
const size_t smem_histo = GGML_PAD(256*sizeof(uint32_t), 16);
const size_t smem_bucket = GGML_PAD( sizeof(uint32_t), 16);
const size_t smem_above = GGML_PAD( sizeof(uint32_t), 16);
- const size_t smem_out = GGML_PAD( sizeof(uint32_t), 16);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
@@ -6022,7 +6021,6 @@ static void ggml_metal_op_top_k_radix(ggml_metal_op_t ctx, int idx) {
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_histo, 0);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_bucket, 1);
ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_above, 2);
- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem_out, 3);
ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1);
}
diff --git a/ggml/src/ggml-metal/kernels/argsort.metal b/ggml/src/ggml-metal/kernels/argsort.metal
index 6328528a9..6c6fcb534 100644
--- a/ggml/src/ggml-metal/kernels/argsort.metal
+++ b/ggml/src/ggml-metal/kernels/argsort.metal
@@ -256,7 +256,6 @@ kernel void kernel_top_k_f32_i32(
threadgroup atomic_uint * histo [[threadgroup(0)]],
threadgroup uint * sh_bucket [[threadgroup(1)]],
threadgroup uint * sh_above [[threadgroup(2)]],
- threadgroup atomic_uint * out_count [[threadgroup(3)]],
uint3 tgpig[[threadgroup_position_in_grid]],
ushort3 tpitg[[thread_position_in_threadgroup]],
ushort3 ntg[[threads_per_threadgroup]]) {
@@ -294,20 +293,32 @@ kernel void kernel_top_k_f32_i32(
}
threadgroup_barrier(mem_flags::mem_threadgroup);
- // top-down scan for the bucket holding the k-th value
- if (tid == 0) {
- uint acc = 0;
- uint b = 0;
- for (int bb = 255; bb >= 0; --bb) {
- const uint c = atomic_load_explicit(&histo[bb], memory_order_relaxed);
- if (acc + c >= desired) {
- b = uint(bb);
- break;
+ // top-down scan for the bucket holding the k-th value: each lane of the first
+ // simdgroup sums 8 bins, the lane whose range reaches desired scans its own bins
+ if (tid < 32) {
+ const uint top = 255 - 8*tid;
+
+ uint c[8];
+ uint sum = 0;
+ for (uint j = 0; j < 8; ++j) {
+ c[j] = atomic_load_explicit(&histo[top - j], memory_order_relaxed);
+ sum += c[j];
+ }
+ uint acc = simd_prefix_exclusive_sum(sum);
+
+ const simd_vote::vote_t hit = simd_vote::vote_t(simd_ballot(acc + sum >= desired));
+ if (tid == (uint) ctz(hit)) {
+ uint b = top;
+ for (uint j = 0; j < 8; ++j) {
+ if (acc + c[j] >= desired) {
+ b = top - j;
+ break;
+ }
+ acc += c[j];
}
- acc += c;
+ *sh_bucket = b;
+ *sh_above = acc;
}
- *sh_bucket = b;
- *sh_above = acc;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
@@ -318,28 +329,101 @@ kernel void kernel_top_k_f32_i32(
threadgroup_barrier(mem_flags::mem_threadgroup);
}
- if (tid == 0) {
- atomic_store_explicit(out_count, 0u, memory_order_relaxed);
+ // deterministic compaction: rank the kept elements by column order, so the
+ // result does not depend on the thread scheduling
+ const uint threshold = prefix;
+
+ // each warp scans a contiguous chunk of the row, and the lanes of the warp
+ // interleave within the chunk with stride 32 to keep the accesses coalesced
+ const uint lane = tid & 31u;
+ const uint wg = tid >> 5u;
+ const uint nw = ntg_x >> 5u;
+ const uint wlo = wg*((ncols + nw - 1) / nw);
+ const uint whi = MIN(wlo + (ncols + nw - 1) / nw, ncols);
+ const uint wlen = wlo < whi ? whi - wlo : 0;
+ const uint nstep = (wlen + 31) >> 5;
+
+ uint ca = 0; // elements above the threshold in this lane's substream
+ uint ct = 0; // elements tied with the threshold in this lane's substream
+ for (uint i = wlo + lane; i < whi; i += 32u) {
+ const uint key = ggml_top_k_f2ui(src0_row[i]);
+ ca += (key > threshold) ? 1u : 0u;
+ ct += (key == threshold) ? 1u : 0u;
+ }
+
+ // per-warp totals into the now unused histogram, so that the first warp can
+ // compute the exclusive base of each warp with a scan across its lanes
+ const uint wca = simd_sum(ca);
+ const uint wct = simd_sum(ct);
+ if (lane == 0) {
+ atomic_store_explicit(&histo[wg], wca, memory_order_relaxed);
+ atomic_store_explicit(&histo[32 + wg], wct, memory_order_relaxed);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
- // emit everything above the threshold, then fill the rest from ties
- const uint threshold = prefix;
-
- for (uint i = tid; i < ncols; i += ntg_x) {
- if (ggml_top_k_f2ui(src0_row[i]) > threshold) {
- const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
- dst_row[pos] = (int32_t) i;
+ if (tid < 32) {
+ const bool valid = tid < nw;
+ const uint sa = valid ? atomic_load_explicit(&histo[tid], memory_order_relaxed) : 0u;
+ const uint st = valid ? atomic_load_explicit(&histo[32 + tid], memory_order_relaxed) : 0u;
+ const uint ea = simd_prefix_exclusive_sum(sa);
+ const uint et = simd_prefix_exclusive_sum(st);
+ if (valid) {
+ atomic_store_explicit(&histo[tid], ea, memory_order_relaxed);
+ atomic_store_explicit(&histo[32 + tid], et, memory_order_relaxed);
+ }
+ const uint ta = simd_sum(sa);
+ const uint tt = simd_sum(st);
+ if (tid == 0) {
+ *sh_bucket = ta; // total above the threshold
+ *sh_above = tt; // total tied with the threshold
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
- for (uint i = tid; i < ncols; i += ntg_x) {
- if (ggml_top_k_f2ui(src0_row[i]) == threshold) {
- const uint pos = atomic_fetch_add_explicit(out_count, 1u, memory_order_relaxed);
- if (pos < top_k) {
+ const uint above_total = *sh_bucket;
+ const uint tie_total = *sh_above;
+ const uint base_wa = atomic_load_explicit(&histo[wg], memory_order_relaxed);
+ const uint base_wt = above_total + atomic_load_explicit(&histo[32 + wg], memory_order_relaxed);
+
+ // the lanes write their kept elements at lane-major ranks, which does not
+ // depend on the thread scheduling. with more tied candidates than remaining
+ // slots, the ties are instead ranked in column order with a per-step
+ // cross-lane scan, so that the lowest tied indices are the ones kept
+ const uint excl_a = simd_prefix_exclusive_sum(ca);
+ const uint excl_t = simd_prefix_exclusive_sum(ct);
+
+ if (tie_total <= top_k - above_total) {
+ uint na = 0;
+ uint nt = 0;
+ for (uint i = wlo + lane; i < whi; i += 32u) {
+ const uint key = ggml_top_k_f2ui(src0_row[i]);
+ if (key > threshold) {
+ dst_row[base_wa + excl_a + na] = (int32_t) i;
+ ++na;
+ } else if (key == threshold) {
+ dst_row[base_wt + excl_t + nt] = (int32_t) i;
+ ++nt;
+ }
+ }
+ } else {
+ uint na = 0;
+ for (uint i = wlo + lane; i < whi; i += 32u) {
+ if (ggml_top_k_f2ui(src0_row[i]) > threshold) {
+ dst_row[base_wa + excl_a + na] = (int32_t) i;
+ ++na;
+ }
+ }
+
+ uint tt = 0; // tied elements in the earlier steps of this warp
+ for (uint s = 0; s < nstep; ++s) {
+ const uint i = wlo + 32u*s + lane;
+ const bool ok = i < whi;
+ const bool t = ok && ggml_top_k_f2ui(src0_row[i]) == threshold;
+ const uint pos = base_wt + tt + simd_prefix_exclusive_sum((uint) t);
+ if (t && pos < top_k) {
dst_row[pos] = (int32_t) i;
}
+ tt += simd_sum((uint) t);
}
}
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index bfbd5322c..016fca3c3 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -1341,6 +1341,18 @@ struct test_case {
}
virtual bool run_whole_graph() { return false; }
+ // the output must be the same in each run on the same backend
+ virtual bool deterministic(ggml_backend_t backend) {
+ GGML_UNUSED(backend);
+ return false;
+ }
+ virtual bool same_output(ggml_tensor * t1, ggml_tensor * t2) {
+ std::vector<uint8_t> t1_data(ggml_nbytes(t1));
+ std::vector<uint8_t> t2_data(ggml_nbytes(t2));
+ ggml_backend_tensor_get(t1, t1_data.data(), 0, ggml_nbytes(t1));
+ ggml_backend_tensor_get(t2, t2_data.data(), 0, ggml_nbytes(t2));
+ return t1_data == t2_data;
+ }
virtual std::vector<ggml_tensor *> fusion_test_nodes() { return {}; }
virtual bool use_weight_context() { return false; }
@@ -1616,6 +1628,24 @@ struct test_case {
run_whole_graph() ? fused_nodes_to_verify.data() : nullptr,
fused_nodes_to_verify.size());
+ // compare backend1 with itself: a second copy of the graph must give the same bytes
+ if (deterministic(backend1)) {
+ auto callback_same = [](int index, ggml_tensor * t1, ggml_tensor * t2, void * user_data) -> bool {
+ callback_userdata * ud = (callback_userdata *) user_data;
+ if (!ud->tc->same_output(t1, t2)) {
+ printf("[%s] not deterministic ", ggml_op_desc(t1));
+ ud->ok = false;
+ }
+ return true;
+
+ GGML_UNUSED(index);
+ };
+
+ for (int i = 0; i < 4 && ud.ok; i++) {
+ ggml_backend_compare_graph_backend(backend1, backend1, gf, callback_same, &ud, &out, 1);
+ }
+ }
+
// Create test result
bool test_passed = ud.ok && cmp_ok;
std::string error_msg = test_passed ? "" : (!cmp_ok ? "compare failed" : "test failed");
@@ -7346,6 +7376,34 @@ struct test_top_k : public test_case {
// The logic in err can't handle the sentinel tensors.
bool run_whole_graph() override { return ties; }
+ // the backends may break ties in different ways, but each backend must break them in the same way in each run
+ bool deterministic(ggml_backend_t backend) override {
+ // TODO: Vulkan is not deterministic yet, SYCL gathers the ties with an atomic counter as well
+ const char * name = ggml_backend_reg_name(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)));
+ // ROCm has no cub::DeviceTopK, every row goes through the radix or bitonic path
+ if (strcmp(name, "ROCm") == 0) {
+ return ties;
+ }
+ if (strcmp(name, "CUDA") != 0 && strcmp(name, "MTL") != 0) {
+ return false;
+ }
+ // CUDA uses cub::DeviceTopK for 1 or 2 rows, which does not guarantee deterministic ties (CCCL <= 3.5)
+ return ties && ne[1]*ne[2]*ne[3] > 2;
+ }
+
+ // the order in a row is not specified, so compare the sets
+ bool same_output(ggml_tensor * t1, ggml_tensor * t2) override {
+ std::vector<int32_t> i1(ggml_nelements(t1));
+ std::vector<int32_t> i2(ggml_nelements(t2));
+ ggml_backend_tensor_get(t1, i1.data(), 0, ggml_nbytes(t1));
+ ggml_backend_tensor_get(t2, i2.data(), 0, ggml_nbytes(t2));
+ for (size_t r = 0; r < i1.size(); r += k) {
+ std::sort(i1.begin() + r, i1.begin() + r + k);
+ std::sort(i2.begin() + r, i2.begin() + r + k);
+ }
+ return i1 == i2;
+ }
+
double err(const float * a, const float * b, size_t n) override {
// When there are no ties, we expect the exact same set of indices,
// but possibly in a different order. When there are ties, the indices
@@ -12312,6 +12370,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_top_k_inf({4096, 2, 1, 1}, k));
test_cases.emplace_back(new test_top_k_inf({248320, 1, 1, 1}, k));
}
+ // a selection of the indexer of sparse attention, with more ties at the threshold than slots
+ test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {1088, 2048, 1, 1}, 512, true));
for (int i = 0; i < 20; ++i) {
for (int k : {1, 2, 3, 7, 15, 100, 500, 1023, 9999}) {