Fix numerical stability in Triton masked softmax (#8782)
Fixes #8780. A mask that suppresses the row max made every `exp()`
underflow to 0, giving 0/0 = NaN. Now the row max is subtracted again
after the mask is added.
The first shift stays to keep FP32 mask precision.
Adds one regression test.
---------
Signed-off-by: Srijan Keshri <212402043+arcusbuilds@users.noreply.github.com>
Co-authored-by: pengdurice <pengduhit@gmail.com>
