[WebGPU] implements online softmax for webgpu ep

[prathikr]

Feature request from #29393

Performance/Accurary Metrics:

Intel Core Ultra 7 266v (Lunar Lake)

input shape=[1024, 35840] size_mb=140.00 seed=0
warmup_runs=3 timed_runs=500
naive mean=0.000028 min=0.000000 max=0.006227 time_ms avg=43.382 median=40.710 min=34.660
online mean=0.000028 min=0.000000 max=0.006227 time_ms avg=39.789 median=39.018 min=35.731
max_abs_diff=0.00000000
speedup naive/online=1.090x

Snapdragon(R) X Elite - X1E80100 - Qualcomm(R) Hexagon(TM) NPU

input shape=[1024, 35840] size_mb=140.00 seed=0
warmup_runs=3 timed_runs=500
naive mean=0.000028 min=0.000000 max=0.006227 time_ms avg=29.570 median=29.455 min=25.665
online mean=0.000028 min=0.000000 max=0.006227 time_ms avg=28.621 median=28.724 min=24.834
max_abs_diff=0.00000000
speedup naive/online=1.033x

Benchmarking Script/Model

softmax_benchmark.zip

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR adds an optional online safe softmax implementation to the WebGPU Execution Provider, selectable via a new EP configuration option, to improve performance on certain hardware while retaining the existing naive softmax path as the default.

Changes:

• Added a new WebGPU EP option (ep.webgpuexecutionprovider.softmaxAlgorithm) with supported values "naive" and "online", including parsing and verbose logging.
• Implemented an online softmax shader path in the WebGPU Softmax kernel and ensured program cache separation between algorithms.
• Added WebGPU-focused unit tests to validate correctness of the online softmax path (including large values and all -inf inputs).

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
onnxruntime/test/providers/cpu/math/softmax_test.cc	Adds WebGPU online-softmax coverage tests and helper to configure the EP option.
onnxruntime/core/providers/webgpu/webgpu_provider_options.h	Introduces the new softmaxAlgorithm option key and allowed string values.
onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc	Parses softmaxAlgorithm from config options and logs the selected algorithm.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.h	Adds SoftmaxAlgorithm enum and stores the selected algorithm in EP config/state.
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Wires the configured softmax algorithm into the EP instance.
onnxruntime/core/providers/webgpu/math/softmax.h	Extends SoftmaxProgram to carry the chosen softmax algorithm.
onnxruntime/core/providers/webgpu/math/softmax.cc	Implements the online softmax shader path and selects scalar packing for online mode.
onnxruntime/core/providers/webgpu/compute_context.h	Exposes the selected softmax algorithm to kernels via ComputeContextBase.

[qjia7]

Please also support components > 1 in online. The per-thread loop and the cross-thread reduction already work component-wise on vec types — the only blocker is the final lane combine. Once vectorization in place, online is numerically identical to naive and does one fewer pass over the row, so I don't see a reason to keep both implementations. Please re-benchmark across the shapes naive currently covers. If online ≥ naive everywhere, we can just keep the online version and drop the user-facing option.

[prathikr]

@qjia7 can you clarify what you mean by "re-benchmark across the shapes naive currently covers"? I attached my benchmarking script in the PR description, are there any specific input shapes/contents you think would be good to test?

[qjia7]

I suppose you can get better performance with components > 1 support in online. I'd like to know the new perf data comparison after this change.

[qjia7]

fyi: The non-flash attention path looks like a good candidate for online softmax. InPlaceSoftmaxProgram currently makes three passes over probs (max, then sum, then normalize) with two barriers; folding the max and sum into a single running-max pass would cut that to two passes over probs and one barrier. Might be worth prototyping to measure the impact — likely a separate PR (The flash path already does online softmax by construction, so this only helps the non-flash fallback.)

[prathikr]

continued in #29694