[webgpu] support smooth softmax for non-FA GQA implementation

[fs-eire]

Description

support smooth softmax for non-FA GQA implementation

This change depends on:

• [CPU] GQA supports head_sink input for smooth softmax #25269

Work items:

• support smooth softmax
• support bias
• support head sink (per-head smooth softmax)

The following will not be included in this PR:

• support for FlashAttention
• support sliding window

[Copilot]

Pull Request Overview

Adds support for smooth softmax, attention bias, and a per-head “sink” value in the non-FlashAttention path of GroupQueryAttention.

• Introduce three new optional inputs (position_ids, attention_bias, head_sink) in both WebGPU and CPU GQA kernels
• Wire attention_bias and head_sink through ComputeInternal, ApplyAttention, and the softmax shader
• Update program constructors and shader generation to respect use_smooth_softmax, has_seqlen_k, and has_head_sink flags

Reviewed Changes

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

Show a summary per file

File	Description
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc	Added new inputs and custom‐input checks; passed attention_bias & head_sink into ComputeInternal and ApplyAttention
onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc	Noted TODO for smooth softmax/head_sink support in FlashAttention
onnxruntime/contrib_ops/webgpu/bert/attention_common.h	Extended ApplyAttention signature to include head_sink
onnxruntime/contrib_ops/webgpu/bert/attention.h	Modified program constructors to take use_smooth_softmax, has_seqlen_k, has_head_sink flags
onnxruntime/contrib_ops/webgpu/bert/attention.cc	Enhanced InPlaceSoftmaxProgram generation to handle the new flags in the shader
onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h	Added validation logic for the new head_sink tensor
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Loaded head_sink input and invoked validation in the CPU kernel

Comments suppressed due to low confidence (3)

onnxruntime/contrib_ops/webgpu/bert/attention.cc:226

• The new smooth_softmax and head_sink branches in the shader code introduce distinct behavior; please add unit or integration tests to cover both code paths in WebGPU and CPU implementations.

Status InPlaceSoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc:155

• The operator schema and user‐facing documentation need updating to expose the new optional inputs (position_ids, attention_bias, and head_sink) and explain their semantics.

  const Tensor* position_ids = context.Input<Tensor>(9);  // TODO: support sliding window

onnxruntime/contrib_ops/webgpu/bert/attention.h:72

• [nitpick] The term head_sink may be unfamiliar to new readers; consider renaming it to something more descriptive or adding a comment in the header to clarify its intended effect on the softmax.

  InPlaceSoftmaxProgram(int work_group_size, int components, bool use_smooth_softmax, bool has_seqlen_k, bool has_head_sink)