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

format

Browse Source
This commit is contained in:
Awni Hannun
2025-07-29 17:17:47 -07:00
parent 5694f764fc
commit 968868fa26
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -54,7 +54,8 @@ __global__ void logsumexp(const T* in, T* out, int axis_size) {
// https://github.com/NVIDIA/online-softmax // https://github.com/NVIDIA/online-softmax
normalizer = normalizer * softmax_exp(prevmax - maxval); normalizer = normalizer * softmax_exp(prevmax - maxval);
for (int i = 0; i < N_READS; i++) { for (int i = 0; i < N_READS; i++) {
normalizer = normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval); normalizer =
normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
} }
} }