From 4400d477a797e013235bba97367abdb58bf8d4bb Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Thu, 9 Jul 2026 15:22:18 -0700 Subject: [PATCH 1/6] impl --- .../core/providers/webgpu/compute_context.h | 4 + .../core/providers/webgpu/math/softmax.cc | 140 +++++++++++------- .../core/providers/webgpu/math/softmax.h | 5 +- .../webgpu/webgpu_execution_provider.cc | 1 + .../webgpu/webgpu_execution_provider.h | 8 + .../webgpu/webgpu_provider_factory.cc | 23 +++ .../webgpu/webgpu_provider_options.h | 4 + 7 files changed, 127 insertions(+), 58 deletions(-) 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..282ff9245f529 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -73,6 +73,11 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << "var row_sum_shared : x_value_t;\n" << "var thread_shared : array;\n"; + if (algorithm_ == SoftmaxAlgorithm::Online) { + shader.AdditionalImplementation() + << "var thread_sum_shared : array;\n"; + } + // Define helper functions to get and set values shader.AdditionalImplementation() << "fn getValue(row: i32, col: i32, row_stride: i32) -> x_value_t {\n" @@ -84,63 +89,85 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << " result[index] = value;\n" << "}\n"; - // Main function body - shader.MainFunctionBody() - << " let gindex = i32(global_idx);\n" + auto& body = shader.MainFunctionBody(); + body << " let gindex = i32(global_idx);\n" << " let lindex = i32(local_idx);\n" << " const wg = " << wg_ << ";\n" << " let row = gindex / wg;\n" << " let cols = uniforms.packedCols;\n" - << " let row_stride : i32 = uniforms.packedCols;\n" - - // Find the row's max value - << thread_max_decl - << " for (var col = lindex; col < cols; col += wg) {\n" - << " let value = getValue(row, col, row_stride);\n" - << " thread_max = max(thread_max, value);\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" - << " workgroupBarrier();\n" - - // Reduce to find the sum of exponentials - << " 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" - << " }\n" - << " workgroupBarrier();\n" - << " }\n" - << " if (lindex == 0) {\n" - << " row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n" - << " }\n" - << " workgroupBarrier();\n" - - // Calculate the final value for each element in the row - << " for (var col = lindex; col < cols; col += wg) {\n" + << " let row_stride : i32 = uniforms.packedCols;\n"; + + if (algorithm_ == SoftmaxAlgorithm::Online) { + body << thread_max_decl + << " var thread_sum = x_value_t(0.0);\n" + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = getValue(row, col, row_stride);\n" + << " let new_max = max(thread_max, value);\n" + << " thread_sum = thread_sum * exp(thread_max - new_max) + exp(value - new_max);\n" + << " thread_max = new_max;\n" + << " }\n" + << " thread_shared[lindex] = thread_max;\n" + << " thread_sum_shared[lindex] = thread_sum;\n" + << " workgroupBarrier();\n" + << " for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n" + << " if (lindex < curr_size) {\n" + << " let rhs_max = thread_shared[lindex + curr_size];\n" + << " let rhs_sum = thread_sum_shared[lindex + curr_size];\n" + << " let lhs_max = thread_shared[lindex];\n" + << " let lhs_sum = thread_sum_shared[lindex];\n" + << " let merged_max = max(lhs_max, rhs_max);\n" + << " thread_sum_shared[lindex] = lhs_sum * exp(lhs_max - merged_max) + rhs_sum * exp(rhs_max - merged_max);\n" + << " thread_shared[lindex] = merged_max;\n" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " row_max_shared = thread_shared[0];\n" + << " row_sum_shared = thread_sum_shared[0];\n" + << " }\n" + << " workgroupBarrier();\n"; + } else { + body << thread_max_decl + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = getValue(row, col, row_stride);\n" + << " thread_max = max(thread_max, value);\n" + << " }\n" + << " if (lindex < cols) {\n" + << " thread_shared[lindex] = thread_max;\n" + << " }\n" + << " workgroupBarrier();\n" + << " 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" + << " 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" + << " workgroupBarrier();\n" + << " 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" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n" + << " }\n" + << " workgroupBarrier();\n"; + } + + body << " for (var col = lindex; col < cols; col += wg) {\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" @@ -187,13 +214,14 @@ Status Softmax::ComputeInternal(ComputeContext& context) const { // one part is treated as batch size, and the other part is performed by Softmax. const int64_t cols = is_transpose_required ? transposed_input_shape[input_rank - 1] : (opset_ >= 13 ? input_shape[input_rank - 1] : input_shape.SizeFromDimension(axis)); const int64_t rows = input_shape.Size() / cols; - const int64_t components = GetMaxComponents(cols); + const auto algorithm = context.GetSoftmaxAlgorithm(); + const int64_t components = algorithm == SoftmaxAlgorithm::Online ? 1 : GetMaxComponents(cols); const auto packed_cols = cols / components; uint32_t workgroup_size = rows == 1 ? 256 : 64; // 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}; + SoftmaxProgram program{workgroup_size, is_fp32, algorithm}; if (is_transpose_required) { program .AddInputs({{&transposed_input_tensor, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components)}}) @@ -205,7 +233,7 @@ Status Softmax::ComputeInternal(ComputeContext& context) const { } program - .CacheHint(std::to_string(components), std::to_string(workgroup_size)) + .CacheHint(std::to_string(components), std::to_string(workgroup_size), algorithm == SoftmaxAlgorithm::Online ? "online" : "naive") .SetWorkgroupSize(workgroup_size) .SetDispatchGroupSize(static_cast(rows)) .AddUniformVariables({{static_cast(packed_cols)}}); diff --git a/onnxruntime/core/providers/webgpu/math/softmax.h b/onnxruntime/core/providers/webgpu/math/softmax.h index 532a56ff0be41..4ee784eeddc53 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..c317125b23b20 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"; From 3161c8411a437d9b1796d2c8713b83d952b4905e Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Fri, 10 Jul 2026 12:04:44 -0700 Subject: [PATCH 2/6] lint --- .../core/providers/webgpu/math/softmax.cc | 166 +++++++++--------- .../core/providers/webgpu/math/softmax.h | 2 +- .../webgpu/webgpu_execution_provider.cc | 2 +- 3 files changed, 85 insertions(+), 85 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc index 282ff9245f529..653d2bad1e76e 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -89,90 +89,90 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << " result[index] = value;\n" << "}\n"; - auto& body = shader.MainFunctionBody(); - body << " let gindex = i32(global_idx);\n" - << " let lindex = i32(local_idx);\n" - << " const wg = " << wg_ << ";\n" - << " let row = gindex / wg;\n" - << " let cols = uniforms.packedCols;\n" - << " let row_stride : i32 = uniforms.packedCols;\n"; - - if (algorithm_ == SoftmaxAlgorithm::Online) { - body << thread_max_decl - << " var thread_sum = x_value_t(0.0);\n" - << " for (var col = lindex; col < cols; col += wg) {\n" - << " let value = getValue(row, col, row_stride);\n" - << " let new_max = max(thread_max, value);\n" - << " thread_sum = thread_sum * exp(thread_max - new_max) + exp(value - new_max);\n" - << " thread_max = new_max;\n" - << " }\n" - << " thread_shared[lindex] = thread_max;\n" - << " thread_sum_shared[lindex] = thread_sum;\n" - << " workgroupBarrier();\n" - << " for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n" - << " if (lindex < curr_size) {\n" - << " let rhs_max = thread_shared[lindex + curr_size];\n" - << " let rhs_sum = thread_sum_shared[lindex + curr_size];\n" - << " let lhs_max = thread_shared[lindex];\n" - << " let lhs_sum = thread_sum_shared[lindex];\n" - << " let merged_max = max(lhs_max, rhs_max);\n" - << " thread_sum_shared[lindex] = lhs_sum * exp(lhs_max - merged_max) + rhs_sum * exp(rhs_max - merged_max);\n" - << " thread_shared[lindex] = merged_max;\n" - << " }\n" - << " workgroupBarrier();\n" - << " }\n" - << " if (lindex == 0) {\n" - << " row_max_shared = thread_shared[0];\n" - << " row_sum_shared = thread_sum_shared[0];\n" - << " }\n" - << " workgroupBarrier();\n"; - } else { - body << thread_max_decl - << " for (var col = lindex; col < cols; col += wg) {\n" - << " let value = getValue(row, col, row_stride);\n" - << " thread_max = max(thread_max, value);\n" - << " }\n" - << " if (lindex < cols) {\n" - << " thread_shared[lindex] = thread_max;\n" - << " }\n" - << " workgroupBarrier();\n" - << " 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" - << " 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" - << " workgroupBarrier();\n" - << " 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" - << " }\n" - << " workgroupBarrier();\n" - << " }\n" - << " if (lindex == 0) {\n" - << " row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n" - << " }\n" - << " workgroupBarrier();\n"; - } + auto& body = shader.MainFunctionBody(); + body << " let gindex = i32(global_idx);\n" + << " let lindex = i32(local_idx);\n" + << " const wg = " << wg_ << ";\n" + << " let row = gindex / wg;\n" + << " let cols = uniforms.packedCols;\n" + << " let row_stride : i32 = uniforms.packedCols;\n"; + + if (algorithm_ == SoftmaxAlgorithm::Online) { + body << thread_max_decl + << " var thread_sum = x_value_t(0.0);\n" + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = getValue(row, col, row_stride);\n" + << " let new_max = max(thread_max, value);\n" + << " thread_sum = thread_sum * exp(thread_max - new_max) + exp(value - new_max);\n" + << " thread_max = new_max;\n" + << " }\n" + << " thread_shared[lindex] = thread_max;\n" + << " thread_sum_shared[lindex] = thread_sum;\n" + << " workgroupBarrier();\n" + << " for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n" + << " if (lindex < curr_size) {\n" + << " let rhs_max = thread_shared[lindex + curr_size];\n" + << " let rhs_sum = thread_sum_shared[lindex + curr_size];\n" + << " let lhs_max = thread_shared[lindex];\n" + << " let lhs_sum = thread_sum_shared[lindex];\n" + << " let merged_max = max(lhs_max, rhs_max);\n" + << " thread_sum_shared[lindex] = lhs_sum * exp(lhs_max - merged_max) + rhs_sum * exp(rhs_max - merged_max);\n" + << " thread_shared[lindex] = merged_max;\n" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " row_max_shared = thread_shared[0];\n" + << " row_sum_shared = thread_sum_shared[0];\n" + << " }\n" + << " workgroupBarrier();\n"; + } else { + body << thread_max_decl + << " for (var col = lindex; col < cols; col += wg) {\n" + << " let value = getValue(row, col, row_stride);\n" + << " thread_max = max(thread_max, value);\n" + << " }\n" + << " if (lindex < cols) {\n" + << " thread_shared[lindex] = thread_max;\n" + << " }\n" + << " workgroupBarrier();\n" + << " 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" + << " 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" + << " workgroupBarrier();\n" + << " 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" + << " }\n" + << " workgroupBarrier();\n" + << " }\n" + << " if (lindex == 0) {\n" + << " row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n" + << " }\n" + << " workgroupBarrier();\n"; + } - body << " for (var col = lindex; col < cols; col += wg) {\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"; + body << " for (var col = lindex; col < cols; col += wg) {\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"; return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/math/softmax.h b/onnxruntime/core/providers/webgpu/math/softmax.h index 4ee784eeddc53..69f1ab0989fe7 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.h +++ b/onnxruntime/core/providers/webgpu/math/softmax.h @@ -39,7 +39,7 @@ class Softmax final : public WebGpuKernel { class SoftmaxProgram final : public Program { public: SoftmaxProgram(uint32_t wg, bool is_fp32, SoftmaxAlgorithm algorithm) - : Program{"Softmax"}, wg_{wg}, is_fp32_{is_fp32}, algorithm_{algorithm} { + : Program{"Softmax"}, wg_{wg}, is_fp32_{is_fp32}, algorithm_{algorithm} { } Status GenerateShaderCode(ShaderHelper& sh) const override; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index c317125b23b20..4fbf189ae0d63 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -598,7 +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}, + 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) { From a17ff98787aade2ad4ec4a910bb3836dc1d8d8f2 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Fri, 10 Jul 2026 14:44:28 -0700 Subject: [PATCH 3/6] unit tests --- .../test/providers/cpu/math/softmax_test.cc | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 962a055b5fcbe..1d41dd341bab1 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,37 @@ namespace onnxruntime { namespace test { +#ifdef USE_WEBGPU +static std::unique_ptr CreateWebGpuProvider(const char* softmax_algorithm) { + ConfigOptions config_options{}; + if (softmax_algorithm != nullptr) { + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kSoftmaxAlgorithm, softmax_algorithm)); + } + + return WebGpuExecutionProviderWithOptions(config_options); +} + +static void RunWebGpuSoftmaxOnlineTest(const std::vector& x_vals, + const std::vector& expected_vals, + const std::vector& dimensions, + int opset = 13, + int64_t axis = -1) { + auto provider = CreateWebGpuProvider(webgpu::options::kSoftmaxAlgorithm_Online); + 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 +99,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) From ae33bf20b9cbf053d40998ba4632cf4684aa28f7 Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Fri, 10 Jul 2026 14:58:15 -0700 Subject: [PATCH 4/6] copilot Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- onnxruntime/test/providers/cpu/math/softmax_test.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 1d41dd341bab1..119f204e5bed1 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -21,7 +21,7 @@ namespace test { static std::unique_ptr CreateWebGpuProvider(const char* softmax_algorithm) { ConfigOptions config_options{}; if (softmax_algorithm != nullptr) { - ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kSoftmaxAlgorithm, softmax_algorithm)); + ASSERT_STATUS_OK(config_options.AddConfigEntry(webgpu::options::kSoftmaxAlgorithm, softmax_algorithm)); } return WebGpuExecutionProviderWithOptions(config_options); From bc688950244577fc2ca1b6221f1761cf669abaff Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 13 Jul 2026 10:58:45 -0700 Subject: [PATCH 5/6] reviewer feedback --- onnxruntime/core/providers/webgpu/math/softmax.cc | 6 +++--- .../test/providers/cpu/math/softmax_test.cc | 15 +++++---------- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc index 653d2bad1e76e..00cadd3c5cc1e 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -122,8 +122,8 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << " workgroupBarrier();\n" << " }\n" << " if (lindex == 0) {\n" - << " row_max_shared = thread_shared[0];\n" - << " row_sum_shared = thread_sum_shared[0];\n" + << " row_max_shared = x_value_t(" << MaxVector("thread_shared[0]", components) << ");\n" + << " row_sum_shared = x_value_t(" << SumVector("thread_sum_shared[0]", components) << ");\n" << " }\n" << " workgroupBarrier();\n"; } else { @@ -215,7 +215,7 @@ Status Softmax::ComputeInternal(ComputeContext& context) const { const int64_t cols = is_transpose_required ? transposed_input_shape[input_rank - 1] : (opset_ >= 13 ? input_shape[input_rank - 1] : input_shape.SizeFromDimension(axis)); const int64_t rows = input_shape.Size() / cols; const auto algorithm = context.GetSoftmaxAlgorithm(); - const int64_t components = algorithm == SoftmaxAlgorithm::Online ? 1 : GetMaxComponents(cols); + const int64_t components = GetMaxComponents(cols); const auto packed_cols = cols / components; uint32_t workgroup_size = rows == 1 ? 256 : 64; // check input tensor element type is float diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 1d41dd341bab1..9d0006b0b7d1a 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -18,21 +18,16 @@ namespace onnxruntime { namespace test { #ifdef USE_WEBGPU -static std::unique_ptr CreateWebGpuProvider(const char* softmax_algorithm) { - ConfigOptions config_options{}; - if (softmax_algorithm != nullptr) { - ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kSoftmaxAlgorithm, softmax_algorithm)); - } - - return WebGpuExecutionProviderWithOptions(config_options); -} - static void RunWebGpuSoftmaxOnlineTest(const std::vector& x_vals, const std::vector& expected_vals, const std::vector& dimensions, int opset = 13, int64_t axis = -1) { - auto provider = CreateWebGpuProvider(webgpu::options::kSoftmaxAlgorithm_Online); + 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."; } From ea8b9ea0582f96eee41ede2f3282820b87a15e4b Mon Sep 17 00:00:00 2001 From: Prathik Rao Date: Mon, 13 Jul 2026 12:31:10 -0700 Subject: [PATCH 6/6] webgpu softmax: replace naive shader with online implementation --- .../core/providers/webgpu/math/softmax.cc | 156 +++++++----------- 1 file changed, 59 insertions(+), 97 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/softmax.cc b/onnxruntime/core/providers/webgpu/math/softmax.cc index 00cadd3c5cc1e..3a397dc75788a 100644 --- a/onnxruntime/core/providers/webgpu/math/softmax.cc +++ b/onnxruntime/core/providers/webgpu/math/softmax.cc @@ -59,24 +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"; - - if (algorithm_ == SoftmaxAlgorithm::Online) { - shader.AdditionalImplementation() - << "var thread_sum_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() @@ -89,90 +86,56 @@ Status SoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const { << " result[index] = value;\n" << "}\n"; - auto& body = shader.MainFunctionBody(); - body << " let gindex = i32(global_idx);\n" - << " let lindex = i32(local_idx);\n" - << " const wg = " << wg_ << ";\n" - << " let row = gindex / wg;\n" - << " let cols = uniforms.packedCols;\n" - << " let row_stride : i32 = uniforms.packedCols;\n"; - - if (algorithm_ == SoftmaxAlgorithm::Online) { - body << thread_max_decl - << " var thread_sum = x_value_t(0.0);\n" - << " for (var col = lindex; col < cols; col += wg) {\n" - << " let value = getValue(row, col, row_stride);\n" - << " let new_max = max(thread_max, value);\n" - << " thread_sum = thread_sum * exp(thread_max - new_max) + exp(value - new_max);\n" - << " thread_max = new_max;\n" - << " }\n" - << " thread_shared[lindex] = thread_max;\n" - << " thread_sum_shared[lindex] = thread_sum;\n" - << " workgroupBarrier();\n" - << " for (var curr_size = wg >> 1; curr_size > 0; curr_size = curr_size >> 1) {\n" - << " if (lindex < curr_size) {\n" - << " let rhs_max = thread_shared[lindex + curr_size];\n" - << " let rhs_sum = thread_sum_shared[lindex + curr_size];\n" - << " let lhs_max = thread_shared[lindex];\n" - << " let lhs_sum = thread_sum_shared[lindex];\n" - << " let merged_max = max(lhs_max, rhs_max);\n" - << " thread_sum_shared[lindex] = lhs_sum * exp(lhs_max - merged_max) + rhs_sum * exp(rhs_max - merged_max);\n" - << " thread_shared[lindex] = merged_max;\n" - << " }\n" - << " workgroupBarrier();\n" - << " }\n" - << " if (lindex == 0) {\n" - << " row_max_shared = x_value_t(" << MaxVector("thread_shared[0]", components) << ");\n" - << " row_sum_shared = x_value_t(" << SumVector("thread_sum_shared[0]", components) << ");\n" - << " }\n" - << " workgroupBarrier();\n"; - } else { - body << thread_max_decl - << " for (var col = lindex; col < cols; col += wg) {\n" - << " let value = getValue(row, col, row_stride);\n" - << " thread_max = max(thread_max, value);\n" - << " }\n" - << " if (lindex < cols) {\n" - << " thread_shared[lindex] = thread_max;\n" - << " }\n" - << " workgroupBarrier();\n" - << " 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" - << " 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" - << " workgroupBarrier();\n" - << " 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" - << " }\n" - << " workgroupBarrier();\n" - << " }\n" - << " if (lindex == 0) {\n" - << " row_sum_shared = x_value_t(" << SumVector("thread_shared[0]", components) << ");\n" - << " }\n" - << " workgroupBarrier();\n"; - } - - body << " for (var col = lindex; col < cols; col += wg) {\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"; + // Main function body + shader.MainFunctionBody() + << " let gindex = i32(global_idx);\n" + << " let lindex = i32(local_idx);\n" + << " const wg = " << wg_ << ";\n" + << " let row = gindex / wg;\n" + << " let cols = uniforms.packedCols;\n" + << " let row_stride : i32 = uniforms.packedCols;\n" + + // 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" + << " 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" + << " thread_max_scalar_shared[lindex] = thread_max;\n" + << " thread_sum_scalar_shared[lindex] = thread_sum;\n" + << " workgroupBarrier();\n" + + // 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" + << " 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_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" + + // Calculate the final value for each element in the row + << " for (var col = lindex; col < cols; col += wg) {\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"; return Status::OK(); } @@ -214,14 +177,13 @@ Status Softmax::ComputeInternal(ComputeContext& context) const { // one part is treated as batch size, and the other part is performed by Softmax. const int64_t cols = is_transpose_required ? transposed_input_shape[input_rank - 1] : (opset_ >= 13 ? input_shape[input_rank - 1] : input_shape.SizeFromDimension(axis)); const int64_t rows = input_shape.Size() / cols; - const auto algorithm = context.GetSoftmaxAlgorithm(); const int64_t components = GetMaxComponents(cols); const auto packed_cols = cols / components; uint32_t workgroup_size = rows == 1 ? 256 : 64; // 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, algorithm}; + SoftmaxProgram program{workgroup_size, is_fp32}; if (is_transpose_required) { program .AddInputs({{&transposed_input_tensor, ProgramTensorMetadataDependency::TypeAndRank, static_cast(components)}}) @@ -233,7 +195,7 @@ Status Softmax::ComputeInternal(ComputeContext& context) const { } program - .CacheHint(std::to_string(components), std::to_string(workgroup_size), algorithm == SoftmaxAlgorithm::Online ? "online" : "naive") + .CacheHint(std::to_string(components), std::to_string(workgroup_size)) .SetWorkgroupSize(workgroup_size) .SetDispatchGroupSize(static_cast(rows)) .AddUniformVariables({{static_cast(packed_cols)}});