zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[CUDA] RMSNorm and VJP (#2280)

Browse Source 
* rms norm start

* nit
This commit is contained in:
Awni Hannun
2025-06-12 17:09:49 -07:00
committed by GitHub
parent a4fc671d3e
commit 918761a25a
4 changed files with 345 additions and 4 deletions
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 = o.data_shared_ptr() ? o : out;
const array x = set_output(inputs[0]);
const array& w = inputs[1];
const array& b = inputs[2];