zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-30 13:41:14 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

nit

Browse Source
This commit is contained in:
Awni Hannun 2025-06-12 16:28:02 -07:00
parent bd2ea38397
commit 50dcaa6a7c
2 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/cuda/layer_norm.cu
View File

@ -244,8 +244,7 @@ void LayerNorm::eval_gpu(
} }
}; };
array o = set_output(inputs[0]); const array x = set_output(inputs[0]);
const array& x = o.data_shared_ptr() ? o : out;
const array& w = inputs[1]; const array& w = inputs[1];
const array& b = inputs[2]; const array& b = inputs[2];

4
mlx/backend/cuda/rms_norm.cu
View File

@ -156,7 +156,6 @@ __global__ void rms_norm_vjp(
cub::LoadDirectBlocked(index, g, gn, axis_size); cub::LoadDirectBlocked(index, g, gn, axis_size);
cub::LoadDirectBlocked(index, strided_iterator(w, w_stride), wn, axis_size); cub::LoadDirectBlocked(index, strided_iterator(w, w_stride), wn, axis_size);
for (int i = 0; i < N_READS; i++) { for (int i = 0; i < N_READS; i++) {
// thread_g * thread_w * norm - thread_x * meangwx * norm3
float xi = xn[i]; float xi = xn[i];
float wi = wn[i]; float wi = wn[i];
float gi = gn[i]; float gi = gn[i];
@ -214,8 +213,7 @@ void RMSNorm::eval_gpu(
} }
}; };
array o = set_output(inputs[0]); const array x = set_output(inputs[0]);
const array& x = o.data_shared_ptr() ? o : out;
const array& w = inputs[1]; const array& w = inputs[1];
int32_t axis_size = x.shape().back(); int32_t axis_size = x.shape().back();