zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 01:17:26 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Make sure softmax doesn't change the actual max

Browse Source
This commit is contained in:
Angelos Katharopoulos 2025-06-22 23:34:32 -07:00
parent 818e8e663e
commit fd1d0821d2
2 changed files with 2 additions and 3 deletions
Show Stats 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) {

1
python/tests/cuda_skip.py
View File

@ -13,7 +13,6 @@ cuda_skip = {
"TestLayers.test_upsample",
"TestOps.test_complex_ops",
"TestOps.test_dynamic_slicing",
"TestOps.test_softmax",
"TestReduce.test_axis_permutation_sums",
"TestReduce.test_dtypes",
"TestReduce.test_expand_sums",