[SreecharanGundaboluAMD]

Motivation

When beta=0 is passed to the MIOpen softmax kernel, the output buffer (y forward, dx backward) may contain uninitialized memory or NaN. The kernel was unconditionally reading and scaling these buffers by beta, which produces NaN under IEEE 754 semantics (NaN * 0.0 == NaN), poisoning the result even though beta=0 means the buffer should be ignored entirely. This fixes ROCM-26026 and adds a regression gtest to catch any future reintroduction.

Technical Details

Kernel fix (projects/miopen/src/kernels/MIOpenSoftmax.cpp)

Four accumulation sites shared the same pattern:

value = computed * alpha + CVT_FLOAT2ACCUM(output_buf[idx]) * CVT_FP32_2ACCUM(beta);

Each is replaced with a runtime guard:

value = computed * alpha +
        (beta != 0.0f ? CVT_FLOAT2ACCUM(output_buf[idx]) * CVT_FP32_2ACCUM(beta)
                      : FLOAT_ACCUM{0});

Sites fixed:

• softmaxfwd scalar path (line ~261)
• softmaxfwd vectorized path (line ~390)
• softmaxbwd scalar path (line ~460)
• softmaxbwd vectorized path (line ~529)

This matches the standard BLAS contract: when beta=0, the output buffer must not be read regardless of its contents. The branch is cheap relative to the memory load it avoids.

Regression test (projects/miopen/test/gtest/soft_max.cpp)

Added GPU_Softmax_BetaZeroNaN_FP32, a parameterized fixture with two test methods:

• ForwardTest: pre-poisons the y output buffer with quiet_NaN, calls SoftmaxForward with alpha=1, beta=0, asserts all output elements are finite.

• BackwardTest: pre-poisons the dx output buffer with quiet_NaN, calls SoftmaxBackward with alpha=1, beta=0, asserts all output elements are finite.

• Smoke: all 3 algos (FAST, ACCURATE, LOG) × 2 modes (INSTANCE, CHANNEL) = 6 forward + 6 backward cases

Risk Assessment

Risk Level: 🟢 Low

Level	Indicators
🟢 Low	Single kernel file changed; fix is additive (guard added, no logic removed); only affects the beta=0 code path which was previously broken; no architectural changes

Impacted Components

• projects/miopen/src/kernels/MIOpenSoftmax.cpp — forward and backward softmax GPU kernels (scalar and vectorized paths)
• projects/miopen/test/gtest/soft_max.cpp — softmax gtest suite

Potential Side Effects

• No behavioral change when beta != 0; all four sites are untouched for that case.
• When beta == 0, the output buffer is now skipped — this is the correct BLAS contract. Previously broken inputs now produce correct results.

Mitigation Steps

• Change is isolated to one kernel source file and four symmetric sites.
• Regression test directly reproduces the failure mode and will catch any future reintroduction.
• No new dependencies or API surface.

Test Plan

Added GPU_Softmax_BetaZeroNaN_FP32 gtest fixture that pre-poisons the output buffer with NaN before each kernel call. Without the fix, both ForwardTest and BackwardTest fail on every element. With the fix, both pass.

Run smoke tests:

./test_soft_max --gtest_filter="Smoke/GPU_Softmax_BetaZeroNaN_FP32*"

Test Result

• Both ForwardTest and BackwardTest pass with the kernel fix applied across all algo/mode combinations. Confirmed that without the fix the tests fail on every output element with NaN. No regressions in existing SoftmaxCommon (FP16, BF16) tests.
• Passes CI.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[JonathanLichtnerAMD]

LGTM with a few additions:

• can we add a changelog item for this fix?
• let's add a Risk Assessment category to this PR
• can we add a test or tests to cover this scenario?

Thanks!

[JonathanLichtnerAMD]

LGTM, thanks!