Optimize q4_matmul: add missing assignment back

0cc4m · Sep 5, 2023 · 215a715 · 215a715
1 parent caf4de8
commit 215a715
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/exllama_ext/cuda_func/q4_matmul.cu b/exllama_ext/cuda_func/q4_matmul.cu
@@ -96,14 +96,15 @@ __global__ void q4_matmul_kernel
             {
                 half2 w_scale = w_scales_.item_half2half2(group, w_column);
                 uint32_t w_zero = w_zeros_.item(group, w_column) + 1;
-                dot_product_8(acc, x_cache, w_, k, w_column, w_scale, w_zero, groupsize / 8);
+                acc = dot_product_8(acc, x_cache, w_, k, w_column, w_scale, w_zero, groupsize / 8);
             }
             else
             {
                 half w_scale = w_scales_.item(group, w_column);
                 uint32_t w_zero = w_zeros_.item(group, w_column) + 1;
-                dot_product_8_h(acc_h, x_cache_h, w_, k, w_column, w_scale, w_zero, groupsize / 8);
+                acc_h = dot_product_8_h(acc_h, x_cache_h, w_, k, w_column, w_scale, w_zero, groupsize / 8);
             }
+            __syncthreads();
         }
     }
     else
@@ -124,15 +125,16 @@ __global__ void q4_matmul_kernel
                 int group = k / groupsize;
                 half2 w_scale = w_scales_.item_half2half2(group, w_column);
                 uint32_t w_zero = w_zeros_.item(group, w_column) + 1;
-                dot_product_8(acc, x_cache, w_, k, w_column, w_scale, w_zero, 1);
+                acc = dot_product_8(acc, x_cache, w_, k, w_column, w_scale, w_zero, 1);
             }
             else
             {
                 int group = k / groupsize;
                 half w_scale = w_scales_.item(group, w_column);
                 uint32_t w_zero = w_zeros_.item(group, w_column) + 1;
-                dot_product_8_h(acc_h, x_cache_h, w_, k, w_column, w_scale, w_zero, 1);
+                acc_h = dot_product_8_h(acc_h, x_cache_h, w_, k, w_column, w_scale, w_zero, 1);
             }
+            __syncthreads();
         }
     }