Commit 481c65f09 for llama.cpp
commit 481c65f091f74c5e7089dd0a3a1cc6b50cced31e
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Fri Sep 11 00:44:13 2026 -0500
vulkan: fix data race and OOB access in argsort(large) (#28705)
argsort had a data race in the inner loop, which VVL caught. But I don't think
this was causing failures in practice.
argsort_large has OOB accesses which might explain the failures in CI, but I
couldn't reproduce it locally and I don't think it's a convincing explanation
of the failures.
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp
index 0fc2b9b72..4ba63f7ae 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp
@@ -33,7 +33,11 @@ void argsort(bool needs_bounds_check, const uint row) {
const uint row_offset = row * p.ncols;
// initialize indices
- dst_row[col] = ivec2(col, floatBitsToInt(data_a[row_offset + col]));
+ ivec2 value = ivec2(col, 0);
+ if (!needs_bounds_check || col < p.ncols) {
+ value.y = floatBitsToInt(data_a[row_offset + col]);
+ }
+ dst_row[col] = value;
barrier();
uint num_outer_loop_iters = NCOLS_PADDED_LOG2;
@@ -42,18 +46,20 @@ void argsort(bool needs_bounds_check, const uint row) {
[[unroll]] for (uint j = k / 2, inner_idx = 0; inner_idx < num_inner_loop_iters; j /= 2, inner_idx++) {
const int ixj = int(col ^ j);
- int idx_0 = (col & k) == 0 ? col : ixj;
- int idx_1 = (col & k) == 0 ? ixj : col;
+ if (ixj > col) {
+ int idx_0 = (col & k) == 0 ? col : ixj;
+ int idx_1 = (col & k) == 0 ? ixj : col;
- ivec2 sh_idx_0 = dst_row[idx_0];
- ivec2 sh_idx_1 = dst_row[idx_1];
- bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
- bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
+ ivec2 sh_idx_0 = dst_row[idx_0];
+ ivec2 sh_idx_1 = dst_row[idx_1];
+ bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
+ bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
- if ((idx_0_oob ||
- (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) && (ixj > col)) {
- dst_row[idx_0] = sh_idx_1;
- dst_row[idx_1] = sh_idx_0;
+ if (idx_0_oob ||
+ (!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) {
+ dst_row[idx_0] = sh_idx_1;
+ dst_row[idx_1] = sh_idx_0;
+ }
}
barrier();
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp b/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp
index 920bac6bb..b2df44137 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp
@@ -42,7 +42,10 @@ void argsort(bool needs_bounds_check, const uint row) {
[[unroll]] for (int u = 0; u < WG_UNROLL_FACTOR; ++u) {
uint c = u*BLOCK_SIZE + col;
if (c < p.ncols_padded) {
- ivec2 v = ivec2(c, floatBitsToInt(data_a[row_offset + c]));
+ ivec2 v = ivec2(c, 0);
+ if (!needs_bounds_check || c < p.ncols) {
+ v.y = floatBitsToInt(data_a[row_offset + c]);
+ }
tmp_idx[idx_offset + c] = v;
}
}