From 86e7fc4ffc310c2432453f571745155e6474082f Mon Sep 17 00:00:00 2001
From: Aman Gupta <amangupta052@gmail.com>
Date: Sun, 6 Sep 2026 14:01:34 +0800
Subject: [PATCH] cuda: fixes races in mmid and mmf

---
 ggml/src/ggml-cuda/mmf.cuh | 19 +++++++++++++++++++
 ggml/src/ggml-cuda/mmid.cu |  1 +
 2 files changed, 20 insertions(+)

diff --git a/ggml/src/ggml-cuda/mmf.cuh b/ggml/src/ggml-cuda/mmf.cuh
index d55cc1ec7b..60bf5ce6e6 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;
             }
+            __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) {
+            __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];
             }
+            __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) {
+            __syncwarp();
             if constexpr (std::is_same_v<T, float>) {
 #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<T, void>, "unsupported type");
             }
+            __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 {
+        __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 {
+        __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) {
+            __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];
             }
+            __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) {
+                __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]);
                 }
 
+                __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) {
+                __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]);
                 }
 
+                __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 {
+        __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 {
+        __syncwarp();
     }
 
 #pragma unroll
diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu
index ed0851dcf8..6c53613ef0 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<warp_size>(nex_prev);
+    __syncwarp();
 
     for (int itc = threadIdx.x; itc < it_compact; itc += warp_size) {
         const mm_ids_helper_store store_it = store[itc];