[WebGPU] implements online softmax for webgpu ep

onnxruntime/core/providers/webgpu/compute_context.h

@@ -117,6 +117,10 @@ class ComputeContextBase {
     return ep_.KvCacheQuantizationEnabled();
   }
 
+  inline webgpu::SoftmaxAlgorithm GetSoftmaxAlgorithm() const {
+    return ep_.GetSoftmaxAlgorithm();
+  }
+
   //
   // Get the logger.
   //

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

@@ -59,19 +59,21 @@
 
 Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
   // Add input and output variables
-  const auto& input = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias);
+  const auto& input = shader.AddInput("x", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias |
+                                               ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias);
   shader.AddOutput("result", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias);
   int components = input.NumComponents();
 
-  const std::string thread_max_decl = is_fp32_
-                                          ? "var thread_max = x_value_t(-3.4028234663852886e+38f);\n"
-                                          : "var thread_max = x_value_t(-65504.0h);\n";
+  const std::string thread_max_scalar_decl = is_fp32_
+                                                 ? "var thread_max = x_element_t(-3.4028234663852886e+38f);\n"
+                                                 : "var thread_max = x_element_t(-65504.0h);\n";
 
   // Define shared memory for row max and row sum
   shader.AdditionalImplementation()
       << "var<workgroup> row_max_shared : x_value_t;\n"
       << "var<workgroup> row_sum_shared : x_value_t;\n"
-      << "var<workgroup> thread_shared : array<x_value_t, " << wg_ << ">;\n";
+      << "var<workgroup> thread_max_scalar_shared : array<x_element_t, " << wg_ << ">;\n"
+      << "var<workgroup> thread_sum_scalar_shared : array<x_element_t, " << wg_ << ">;\n";
 
   // Define helper functions to get and set values
   shader.AdditionalImplementation()
@@ -93,49 +95,37 @@
       << "  let cols = uniforms.packedCols;\n"
       << "  let row_stride : i32 = uniforms.packedCols;\n"
 
-      // Find the row's max value
-      << thread_max_decl
+      // Online accumulation: track scalar running max and scalar running sum.
+      << thread_max_scalar_decl
+      << "  var thread_sum = x_element_t(0.0);\n"
       << "  for (var col = lindex; col < cols; col += wg) {\n"
       << "    let value = getValue(row, col, row_stride);\n"
-      << "    thread_max = max(thread_max, value);\n"
+      << "    let value_max = x_element_t(" << MaxVector("value", components) << ");\n"
+      << "    let new_max = max(thread_max, value_max);\n"
+      << "    let shifted_exp = exp(value - x_value_t(new_max));\n"
+      << "    thread_sum = thread_sum * exp(thread_max - new_max) + x_element_t(" << SumVector("shifted_exp", components) << ");\n"
+      << "    thread_max = new_max;\n"
       << "  }\n"
-      << "  if (lindex < cols) {\n"
-      << "    thread_shared[lindex] = thread_max;\n"
-      << "  }\n"
-      << "  workgroupBarrier();\n"
-
-      // Reduce to find the max value
-      << "  var reduce_size = min(cols, wg);\n"
-      << "  for (var curr_size = reduce_size >> 1; curr_size > 0; curr_size = reduce_size >> 1) {\n"
-      << "    reduce_size = curr_size + (reduce_size & 1);\n"
-      << "    if (lindex < curr_size) {\n"
-      << "      thread_shared[lindex] = max(thread_shared[lindex], thread_shared[lindex + reduce_size]);\n"
-      << "    }\n"
-      << "    workgroupBarrier();\n"
-      << "  }\n"
-      << "  if (lindex == 0) {\n"
-      << "    row_max_shared = x_value_t(" << MaxVector("thread_shared[0]", components) << ");\n"
-      << "  }\n"
-      << "  workgroupBarrier();\n"
-
-      // Find the row's sum of exponentials
-      << "  var thread_sum = x_value_t(0.0);\n"
-      << "  for (var col = lindex; col < cols; col += wg) {\n"
-      << "    let sub_exp = exp(getValue(row, col, row_stride) - row_max_shared);\n"
-      << "    thread_sum += sub_exp;\n"
-      << "  }\n"
-      << "  thread_shared[lindex] = thread_sum;\n"
+      << "  thread_max_scalar_shared[lindex] = thread_max;\n"
+      << "  thread_sum_scalar_shared[lindex] = thread_sum;\n"
       << "  workgroupBarrier();\n"
 
-      // Reduce to find the sum of exponentials
+      // Merge per-thread online states.
       << "  for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n"
       << "    if (lindex < curr_size) {\n"
-      << "      thread_shared[lindex] = thread_shared[lindex] + thread_shared[lindex + curr_size];\n"
+      << "      let rhs_max = thread_max_scalar_shared[lindex + curr_size];\n"
+      << "      let rhs_sum = thread_sum_scalar_shared[lindex + curr_size];\n"
+      << "      let lhs_max = thread_max_scalar_shared[lindex];\n"
+      << "      let lhs_sum = thread_sum_scalar_shared[lindex];\n"
+      << "      let merged_max = max(lhs_max, rhs_max);\n"
+      << "      thread_sum_scalar_shared[lindex] = lhs_sum * exp(lhs_max - merged_max) + rhs_sum * exp(rhs_max - merged_max);\n"
+      << "      thread_max_scalar_shared[lindex] = merged_max;\n"
       << "    }\n"
       << "    workgroupBarrier();\n"
       << "  }\n"
       << "  if (lindex == 0) {\n"
-      << "    row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n"
+      << "    row_max_shared = x_value_t(thread_max_scalar_shared[0]);\n"
+      << "    row_sum_shared = x_value_t(thread_sum_scalar_shared[0]);\n"
       << "  }\n"
       << "  workgroupBarrier();\n"
 
@@ -193,7 +183,7 @@
   // check input tensor element type is float
   const bool is_fp32 = input_tensor->GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT;
 
   SoftmaxProgram program{workgroup_size, is_fp32};
   if (is_transpose_required) {
     program
         .AddInputs({{&transposed_input_tensor, ProgramTensorMetadataDependency::TypeAndRank, static_cast<int>(components)}})

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

@@ -38,8 +38,8 @@ class Softmax final : public WebGpuKernel {
 
 class SoftmaxProgram final : public Program<SoftmaxProgram> {
  public:
-  SoftmaxProgram(uint32_t wg, bool is_fp32)
-      : Program{"Softmax"}, wg_{wg}, is_fp32_{is_fp32} {
+  SoftmaxProgram(uint32_t wg, bool is_fp32, SoftmaxAlgorithm algorithm)
+      : Program{"Softmax"}, wg_{wg}, is_fp32_{is_fp32}, algorithm_{algorithm} {
   }
 
   Status GenerateShaderCode(ShaderHelper& sh) const override;
@@ -49,6 +49,7 @@ class SoftmaxProgram final : public Program<SoftmaxProgram> {
  private:
   uint32_t wg_;
   bool is_fp32_;
+  SoftmaxAlgorithm algorithm_;
 };
 
 }  // namespace webgpu

onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc

@@ -598,6 +598,7 @@ WebGpuExecutionProvider::WebGpuExecutionProvider(int context_id,
       enable_int64_{config.enable_graph_capture || config.enable_int64},
       multi_rotary_cache_concat_offset_{config.multi_rotary_cache_concat_offset},
       kv_cache_quantization_bits_{config.kv_cache_quantization_bits},
+      softmax_algorithm_{config.softmax_algorithm},
       prepack_allocator_{std::make_shared<webgpu::GpuBufferAllocator>(
           [this]() -> const webgpu::BufferManager& { return context_.InitializerBufferManager(); }, false)} {
   if (enable_graph_capture_ && config.session_buffer_pool_generations > 0) {

onnxruntime/core/providers/webgpu/webgpu_execution_provider.h

@@ -26,6 +26,11 @@ struct pthreadpool;
 namespace onnxruntime {
 namespace webgpu {
 
+enum class SoftmaxAlgorithm {
+  Naive,
+  Online,
+};
+
 // forward declaration for this EP's namespace.
 template <typename T>
 KernelCreateInfo BuildKernelCreateInfo();
@@ -52,6 +57,7 @@ struct WebGpuExecutionProviderConfig {
   // generator's worth of intermediate buffers.
   size_t session_buffer_pool_generations{1};
   uint32_t kv_cache_quantization_bits{0};  // KV cache quantization bits (0 = off, 4 = 4-bit)
+  webgpu::SoftmaxAlgorithm softmax_algorithm{webgpu::SoftmaxAlgorithm::Naive};
   std::vector<std::string> force_cpu_node_names{};
 };
 
@@ -115,6 +121,7 @@ class WebGpuExecutionProvider : public IExecutionProvider {
   uint32_t MultiRotaryCacheConcatOffset() const { return multi_rotary_cache_concat_offset_; }
   uint32_t KvCacheQuantizationBits() const { return kv_cache_quantization_bits_; }
   bool KvCacheQuantizationEnabled() const { return kv_cache_quantization_bits_ != 0; }
+  webgpu::SoftmaxAlgorithm GetSoftmaxAlgorithm() const { return softmax_algorithm_; }
 
 #if defined(ORT_USE_EP_API_ADAPTERS)
   inline onnxruntime::ep::adapter::Logger& GetEpLogger() const {
@@ -139,6 +146,7 @@ class WebGpuExecutionProvider : public IExecutionProvider {
   bool enable_int64_ = false;
   uint32_t multi_rotary_cache_concat_offset_ = 0;
   uint32_t kv_cache_quantization_bits_ = 0;
+  webgpu::SoftmaxAlgorithm softmax_algorithm_{webgpu::SoftmaxAlgorithm::Naive};
   std::unordered_map<int, int> graph_id_to_run_count_;
   // Required regular runs before graph capture for any necessary allocations.
   const int min_num_runs_before_graph_capture_ = 0;

onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc

@@ -36,6 +36,17 @@ struct WebGpuProviderFactory : IExecutionProviderFactory {
 
 namespace {
 
+const char* ToString(SoftmaxAlgorithm algorithm) {
+  switch (algorithm) {
+    case SoftmaxAlgorithm::Naive:
+      return kSoftmaxAlgorithm_Naive;
+    case SoftmaxAlgorithm::Online:
+      return kSoftmaxAlgorithm_Online;
+    default:
+      return "unknown";
+  }
+}
+
 WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options) {
   WebGpuExecutionProviderConfig webgpu_ep_config{};
 
@@ -109,6 +120,17 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options)
     }
   }
 
+  if (std::string softmax_algorithm_str;
+      config_options.TryGetConfigEntry(kSoftmaxAlgorithm, softmax_algorithm_str)) {
+    if (softmax_algorithm_str == kSoftmaxAlgorithm_Naive) {
+      webgpu_ep_config.softmax_algorithm = SoftmaxAlgorithm::Naive;
+    } else if (softmax_algorithm_str == kSoftmaxAlgorithm_Online) {
+      webgpu_ep_config.softmax_algorithm = SoftmaxAlgorithm::Online;
+    } else {
+      ORT_THROW("Invalid softmaxAlgorithm value: ", softmax_algorithm_str, ". Must be \"naive\" or \"online\".");
+    }
+  }
+
   // parse force CPU node names
   // The force CPU node names are separated by EOL (\n or \r\n) in the config entry.
   // each line is a node name that will be forced to run on CPU.
@@ -147,6 +169,7 @@ WebGpuExecutionProviderConfig ParseEpConfig(const ConfigOptions& config_options)
   LOGS_DEFAULT(VERBOSE) << "WebGPU EP enable int64: " << webgpu_ep_config.enable_int64;
   LOGS_DEFAULT(VERBOSE) << "WebGPU EP multi rotary cache concat offset: " << webgpu_ep_config.multi_rotary_cache_concat_offset;
   LOGS_DEFAULT(VERBOSE) << "WebGPU EP session buffer pool generations: " << webgpu_ep_config.session_buffer_pool_generations;
+  LOGS_DEFAULT(VERBOSE) << "WebGPU EP softmax algorithm: " << ToString(webgpu_ep_config.softmax_algorithm);
 
   return webgpu_ep_config;
 }

onnxruntime/core/providers/webgpu/webgpu_provider_options.h

@@ -18,6 +18,7 @@ constexpr const char* kSessionBufferPoolGenerations = "ep.webgpuexecutionprovide
 constexpr const char* kEnableInt64 = "ep.webgpuexecutionprovider.enableInt64";
 constexpr const char* kMultiRotaryCacheConcatOffset = "ep.webgpuexecutionprovider.multiRotaryCacheConcatOffset";
 constexpr const char* kKvCacheQuantizationBits = "ep.webgpuexecutionprovider.kvCacheQuantizationBits";
+constexpr const char* kSoftmaxAlgorithm = "ep.webgpuexecutionprovider.softmaxAlgorithm";
 
 constexpr const char* kDawnProcTable = "ep.webgpuexecutionprovider.dawnProcTable";
 
@@ -74,6 +75,9 @@ constexpr const char* kPreserveDevice_OFF = "0";
 constexpr const char* kKvCacheQuantizationBits_OFF = "0";
 constexpr const char* kKvCacheQuantizationBits_4Bit = "4";
 
+constexpr const char* kSoftmaxAlgorithm_Naive = "naive";
+constexpr const char* kSoftmaxAlgorithm_Online = "online";
+
 constexpr const char* kBufferCacheMode_Disabled = "disabled";
 constexpr const char* kBufferCacheMode_LazyRelease = "lazyRelease";
 constexpr const char* kBufferCacheMode_Simple = "simple";

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

@@ -6,13 +6,43 @@
 #include "gtest/gtest.h"
 
 #include "core/session/environment.h"
+#ifdef USE_WEBGPU
+#include "core/providers/webgpu/webgpu_provider_options.h"
+#include "test/util/include/default_providers.h"
+#endif
 #include "test/providers/provider_test_utils.h"
 #include "test/common/cuda_op_test_utils.h"
 #include "test/common/dnnl_op_test_utils.h"
 
 namespace onnxruntime {
 namespace test {
 
+#ifdef USE_WEBGPU
+static void RunWebGpuSoftmaxOnlineTest(const std::vector<float>& x_vals,
+                                       const std::vector<float>& expected_vals,
+                                       const std::vector<int64_t>& dimensions,
+                                       int opset = 13,
+                                       int64_t axis = -1) {
+  ConfigOptions config_options{};
+  ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kSoftmaxAlgorithm,
+                                                 webgpu::options::kSoftmaxAlgorithm_Online));
+
+  auto provider = WebGpuExecutionProviderWithOptions(config_options);
+  if (!provider) {
+    GTEST_SKIP() << "WebGPU execution provider is not available.";
+  }
+
+  OpTester test("Softmax", opset);
+  if (axis != -1) {
+    test.AddAttribute("axis", axis);
+  }
+
+  test.AddInput<float>("X", dimensions, x_vals);
+  test.AddOutput<float>("Y", dimensions, expected_vals);
+  test.ConfigEp(std::move(provider)).RunWithConfig();
+}
+#endif
+
 static void RunTest(const std::vector<float>& x_vals,
                     const std::vector<float>& expected_vals,
                     const std::vector<int64_t>& dimensions,
@@ -64,6 +94,43 @@ TEST(SoftmaxOperator, webgpu_nan) {
   test.Run(OpTester::ExpectResult::kExpectSuccess, "",
            {kCpuExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider});
 }
+
+TEST(SoftmaxOperator, WebGpuOnlineSimple) {
+  std::vector<float> x_vals = {-1.0f, 0.0f, 1.0f};
+  std::vector<float> expected_vals = {0.09003058f, 0.24472848f, 0.66524094f};
+  std::vector<int64_t> dimensions = {1, 3};
+
+  RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions);
+}
+
+TEST(SoftmaxOperator, WebGpuOnlineLargeNumber) {
+  std::vector<float> x_vals = {0.0f, 1.0f, 2.0f, 3.0f, 10000.0f, 10001.0f, 10002.0f, 10003.0f};
+  std::vector<float> expected_vals = {0.0320586f, 0.08714432f, 0.23688284f, 0.64391428f,
+                                      0.0320586f, 0.08714432f, 0.23688284f, 0.64391428f};
+  std::vector<int64_t> dimensions = {2, 4};
+
+  RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions);
+}
+
+TEST(SoftmaxOperator, WebGpuOnlineAllNegativeInfinity) {
+  std::vector<float> x_vals = {-INFINITY, -INFINITY, -INFINITY};
+  std::vector<float> expected_vals = {0.0f, 0.0f, 0.0f};
+  std::vector<int64_t> dimensions = {1, 3};
+
+  RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions);
+}
+
+TEST(SoftmaxOperator, DISABLED_WebGpuOnlineLargeInputSize) {
+  constexpr int64_t rows = 1024;
+  constexpr int64_t cols = 256;
+  const std::vector<int64_t> dimensions = {rows, cols};
+
+  const size_t element_count = static_cast<size_t>(rows * cols);
+  std::vector<float> x_vals(element_count, 0.0f);
+  std::vector<float> expected_vals(element_count, 1.0f / static_cast<float>(cols));
+
+  RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions);
+}
 #endif
 
 #if defined(USE_CUDA) || defined(USE_XNNPACK)