From 73a43d1f69345aee8bb186ef4b3172cef892f2e5 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Sun, 6 Sep 2026 19:45:01 +0800 Subject: [PATCH] cuda: fixes races in mmid and mmf (#28475) --- ggml/src/ggml-cuda/common.cuh | 6 ++++++ ggml/src/ggml-cuda/fattn-vec.cuh | 4 +--- ggml/src/ggml-cuda/mmf.cuh | 19 +++++++++++++++++++ ggml/src/ggml-cuda/mmid.cu | 1 + 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9918c03947..ed0ea60bdd 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -121,6 +121,12 @@ # define GGML_CUDA_USE_PDL #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && (CUDART_VERSION >= 12030 || (!(defined(_MSC_VER) && !defined(__clang__)) && CUDART_VERSION >= 11080)) +static __device__ __forceinline__ void ggml_cuda_syncwarp() { +#ifndef GGML_USE_HIP + __syncwarp(); +#endif // GGML_USE_HIP +} + static __device__ __forceinline__ void ggml_cuda_pdl_sync() { #if defined(GGML_CUDA_USE_PDL) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER cudaGridDependencySynchronize(); diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh index 519b36b9ff..57a2855659 100644 --- a/ggml/src/ggml-cuda/fattn-vec.cuh +++ b/ggml/src/ggml-cuda/fattn-vec.cuh @@ -317,9 +317,7 @@ static __global__ void flash_attn_ext_vec( #endif // V_DOT2_F32_F16_AVAILABLE } -#ifndef GGML_USE_HIP - __syncwarp(); -#endif // GGML_USE_HIP + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < WARP_SIZE; k0 += V_cols_per_iter) { diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh index d55cc1ec7b..879a865275 100644 --- a/ggml/src/ggml-cuda/mmf.cuh +++ b/ggml/src/ggml-cuda/mmf.cuh @@ -143,6 +143,7 @@ static __global__ void mul_mat_f( if (threadIdx.x == 0) { slot_map[j] = -1; } + ggml_cuda_syncwarp(); if (col_base + j >= ncols_dst_total) { continue; @@ -171,10 +172,12 @@ static __global__ void mul_mat_f( tile_A A[ntA][warp_size / tile_A::J]; #pragma unroll for (int itA = 0; itA < ntA; ++itA) { + ggml_cuda_syncwarp(); #pragma unroll for (int i = 0; i < tile_A::I; ++i) { tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col]; } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) { load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded); @@ -183,6 +186,7 @@ static __global__ void mul_mat_f( #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); if constexpr (std::is_same_v) { #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { @@ -212,6 +216,7 @@ static __global__ void mul_mat_f( } else { static_assert(std::is_same_v, "unsupported type"); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -229,6 +234,8 @@ static __global__ void mul_mat_f( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll for (int itB = 0; itB < ntB; ++itB) { @@ -245,6 +252,8 @@ static __global__ void mul_mat_f( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll @@ -382,10 +391,12 @@ static __global__ void mul_mat_f_ids( tile_A A[ntA][warp_size / tile_A::J]; #pragma unroll for (int itA = 0; itA < ntA; ++itA) { + ggml_cuda_syncwarp(); #pragma unroll for (int i = 0; i < tile_A::I; ++i) { tile_xy[i*tile_k_padded + threadIdx.x] = x[(itA*tile_A::I + i)*stride_row + col]; } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_A::J) { load_ldmatrix(A[itA][k0/tile_A::J], tile_xy + k0, tile_k_padded); @@ -419,6 +430,7 @@ static __global__ void mul_mat_f_ids( int next_buf = 1; #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { tile_xy[j0*tile_k_padded + threadIdx.x] = vals_buf[curr_buf][j0]; @@ -428,6 +440,7 @@ static __global__ void mul_mat_f_ids( gather_tile(itB + 1, vals_buf[next_buf]); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -472,6 +485,7 @@ static __global__ void mul_mat_f_ids( int next_buf = 1; #pragma unroll for (int itB = 0; itB < ntB; ++itB) { + ggml_cuda_syncwarp(); #pragma unroll for (int j0 = 0; j0 < tile_B::I; ++j0) { const float2 tmp = vals_buf[curr_buf][j0]; @@ -482,6 +496,7 @@ static __global__ void mul_mat_f_ids( gather_tile(itB + 1, vals_buf[next_buf]); } + ggml_cuda_syncwarp(); #pragma unroll for (int k0 = 0; k0 < warp_size; k0 += tile_B::J) { tile_B B; @@ -507,6 +522,8 @@ static __global__ void mul_mat_f_ids( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll for (int itB = 0; itB < ntB; ++itB) { @@ -523,6 +540,8 @@ static __global__ void mul_mat_f_ids( if (nwarps > 1) { __syncthreads(); + } else { + ggml_cuda_syncwarp(); } #pragma unroll diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu index ed0851dcf8..0b222e63ac 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -101,6 +101,7 @@ static __global__ void mm_ids_helper( } } nex_prev = warp_reduce_sum(nex_prev); + ggml_cuda_syncwarp(); for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) { const mm_ids_helper_store store_it = store[itc];