From f5299f72cd20d258eec96cb9b81277226a2afbcd Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Mon, 7 Jul 2025 06:06:01 -0700
Subject: [PATCH] Fix layernorm race condition (#2340)

---
 mlx/backend/metal/kernels/layer_norm.metal | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlx/backend/metal/kernels/layer_norm.metal b/mlx/backend/metal/kernels/layer_norm.metal
index 06b8be55f..ea77b53dc 100644
--- a/mlx/backend/metal/kernels/layer_norm.metal
+++ b/mlx/backend/metal/kernels/layer_norm.metal
@@ -31,6 +31,7 @@ inline void threadgroup_sum(
   for (int i = 0; i < N; i++) {
     x[i] = simd_sum(x[i]);
   }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
   if (simd_lane_id == 0) {
     for (int i = 0; i < N; i++) {
       xs[N * simd_group_id + i] = x[i];