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

faster softmax and logsumexp

Browse Source
This commit is contained in:
Awni Hannun
2025-07-29 14:34:00 -07:00
parent 483221631a
commit 5694f764fc
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -54,7 +54,7 @@ __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(vals[i] - maxval); normalizer = normalizer + softmax_exp(static_cast<AccT>(vals[i]) - maxval);
} }
} }

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

@@ -57,7 +57,8 @@ __global__ void softmax(const T* in, T* out, int axis_size) {
normalizer = normalizer * softmax_exp(prevmax - maxval); normalizer = normalizer * softmax_exp(prevmax - maxval);
#pragma unroll #pragma unroll
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);
} }
} }