zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-15 04:51:13 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix layernorm race condition (#2340)

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-07-07 06:06:01 -07:00 committed by GitHub
parent 0e0d9ac522
commit f5299f72cd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlx/backend/metal/kernels/layer_norm.metal
View File

@ -31,6 +31,7 @@ inline void threadgroup_sum(
for (int i = 0; i < N; i++) { for (int i = 0; i < N; i++) {
x[i] = simd_sum(x[i]); x[i] = simd_sum(x[i]);
} }
threadgroup_barrier(mem_flags::mem_threadgroup);
if (simd_lane_id == 0) { if (simd_lane_id == 0) {
for (int i = 0; i < N; i++) { for (int i = 0; i < N; i++) {
xs[N * simd_group_id + i] = x[i]; xs[N * simd_group_id + i] = x[i];