diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index ae4b63ef293e1..0c6628eae0a5e 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/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. // diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc index bf3bb53341418..3a397dc75788a 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -59,19 +59,21 @@ static std::string MaxVector(const std::string& name, int components) { 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 row_max_shared : x_value_t;\n" << "var row_sum_shared : x_value_t;\n" - << "var thread_shared : array;\n"; + << "var thread_max_scalar_shared : array;\n" + << "var thread_sum_scalar_shared : array;\n"; // Define helper functions to get and set values shader.AdditionalImplementation() @@ -93,49 +95,37 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << " 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" diff --git a/onnxruntime/core/providers/webgpu/math/softmax.h b/onnxruntime/core/providers/webgpu/math/softmax.h index 532a56ff0be41..69f1ab0989fe7 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.h +++ b/onnxruntime/core/providers/webgpu/math/softmax.h @@ -38,8 +38,8 @@ class Softmax final : public WebGpuKernel { class SoftmaxProgram final : public Program { 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 { private: uint32_t wg_; bool is_fp32_; + SoftmaxAlgorithm algorithm_; }; } // namespace webgpu diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index ed9396db344f4..4fbf189ae0d63 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/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( [this]() -> const webgpu::BufferManager& { return context_.InitializerBufferManager(); }, false)} { if (enable_graph_capture_ && config.session_buffer_pool_generations > 0) { diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h index cf0c19e2e44f8..e27bc71c7953e 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.h +++ b/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 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 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 graph_id_to_run_count_; // Required regular runs before graph capture for any necessary allocations. const int min_num_runs_before_graph_capture_ = 0; diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc b/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc index e9a2343be1b8d..5d4cd52c46bc8 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc +++ b/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; } diff --git a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h b/onnxruntime/core/providers/webgpu/webgpu_provider_options.h index 390d325d66ec5..904444c6c9a16 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_provider_options.h +++ b/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"; diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 962a055b5fcbe..9d0006b0b7d1a 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -6,6 +6,10 @@ #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" @@ -13,6 +17,32 @@ namespace onnxruntime { namespace test { +#ifdef USE_WEBGPU +static void RunWebGpuSoftmaxOnlineTest(const std::vector& x_vals, + const std::vector& expected_vals, + const std::vector& 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("X", dimensions, x_vals); + test.AddOutput("Y", dimensions, expected_vals); + test.ConfigEp(std::move(provider)).RunWithConfig(); +} +#endif + static void RunTest(const std::vector& x_vals, const std::vector& expected_vals, const std::vector& dimensions, @@ -64,6 +94,43 @@ TEST(SoftmaxOperator, webgpu_nan) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCpuExecutionProvider, kCoreMLExecutionProvider, kDmlExecutionProvider}); } + +TEST(SoftmaxOperator, WebGpuOnlineSimple) { + std::vector x_vals = {-1.0f, 0.0f, 1.0f}; + std::vector expected_vals = {0.09003058f, 0.24472848f, 0.66524094f}; + std::vector dimensions = {1, 3}; + + RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions); +} + +TEST(SoftmaxOperator, WebGpuOnlineLargeNumber) { + std::vector x_vals = {0.0f, 1.0f, 2.0f, 3.0f, 10000.0f, 10001.0f, 10002.0f, 10003.0f}; + std::vector expected_vals = {0.0320586f, 0.08714432f, 0.23688284f, 0.64391428f, + 0.0320586f, 0.08714432f, 0.23688284f, 0.64391428f}; + std::vector dimensions = {2, 4}; + + RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions); +} + +TEST(SoftmaxOperator, WebGpuOnlineAllNegativeInfinity) { + std::vector x_vals = {-INFINITY, -INFINITY, -INFINITY}; + std::vector expected_vals = {0.0f, 0.0f, 0.0f}; + std::vector 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 dimensions = {rows, cols}; + + const size_t element_count = static_cast(rows * cols); + std::vector x_vals(element_count, 0.0f); + std::vector expected_vals(element_count, 1.0f / static_cast(cols)); + + RunWebGpuSoftmaxOnlineTest(x_vals, expected_vals, dimensions); +} #endif #if defined(USE_CUDA) || defined(USE_XNNPACK)