[WebGPU EP] Fix NaN bug in softmax operator (microsoft#24855)

Handle NaN in softmax operator for WebGPU EP and JSEP.

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: prathikr

js/web/lib/wasm/jsep/webgpu/ops/softmax.ts

@@ -152,7 +152,9 @@ const createSoftmaxProgramInfo = (context: ComputeContext, attributes: SoftmaxAt
 
         // calculate final value for each element in the row
         for (var col = lindex; col < cols; col += wg) {
-          let value = exp(getValue(row, col, row_stride) - rowMaxShared) / rowSumShared;
+          var value = exp(getValue(row, col, row_stride) - rowMaxShared) / rowSumShared;
+          // max operation protects against NaN since all values should be >=0
+          value = max(value, ${valueType}(0.0));
           setValue(row, col, row_stride, value);
         }
       }`;

onnxruntime/core/providers/webgpu/math/softmax.cc

@@ -141,7 +141,9 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
 
       // Calculate the final value for each element in the row
       << "  for (var col = lindex; col < cols; col += wg) {\n"
-      << "    let value = exp(getValue(row, col, row_stride) - row_max_shared) / row_sum_shared;\n"
+      << "    var value = exp(getValue(row, col, row_stride) - row_max_shared) / row_sum_shared;\n"
+      << "    // max operation protects against NaN since all values should be >=0\n"
+      << "    value = max(value, x_value_t(0.0));\n"
       << "    setValue(row, col, row_stride, value);\n"
       << "  }\n";
 

onnxruntime/test/providers/cpu/math/softmax_test.cc

@@ -49,6 +49,22 @@ TEST(SoftmaxOperator, Simple) {
   RunTest(x_vals, expected_vals, dimensions);
 }
 
+#ifdef USE_WEBGPU
+TEST(SoftmaxOperator, webgpu_nan) {
+  OpTester test("Softmax", 13);  // axis default is -1
+
+  std::vector<float> x_vals = {-INFINITY, -INFINITY, -INFINITY};
+  std::vector<float> expected_result = {0.0f, 0.0f, 0.0f};
+  std::vector<int64_t> dimensions = {1, 3};
+
+  test.AddInput<float>("X", dimensions, x_vals);
+  test.AddOutput<float>("Y", dimensions, expected_result);
+
+  // explicitly disable CPU EP for this test since CPU implementation does not handle NaN
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider});
+}
+#endif
+
 #if defined(USE_CUDA) || defined(USE_ROCM) || defined(USE_XNNPACK)
 TEST(SoftmaxOperator, Simple_fp16) {
 #ifdef USE_CUDA