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] Fix reductions (#2314)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-06-27 12:59:20 -07:00
committed by GitHub
parent 2c11d10f8d
commit 772f471ff2
16 changed files with 862 additions and 419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -51,7 +51,7 @@ __global__ void softmax(const T* in, T* out, int axis_size) {
make_cast_iterator<AccT>(in),
vals,
axis_size,
Limits<AccT>::finite_min());
Limits<AccT>::min());
prevmax = maxval;
maxval = max_op(maxval, cub::ThreadReduce(vals, max_op));
// Online normalizer calculation for softmax:
@@ -79,7 +79,7 @@ __global__ void softmax(const T* in, T* out, int axis_size) {
block.sync();
maxval = warp.thread_rank() < warp.meta_group_size()
? local_max[warp.thread_rank()]
: Limits<AccT>::finite_min();
: Limits<AccT>::min();
maxval = cg::reduce(warp, maxval, max_op);
normalizer = normalizer * softmax_exp(prevmax - maxval);
if (warp.thread_rank() == 0) {