[WebGPU] Add opset 24 support and Gemma 4 GQA enhancements

[feich-ms]

Models exported at opset 24 (e.g. Gemma 4) require these registrations to avoid falling back to CPU for basic tensor operations.

Description

This PR adds two categories of WebGPU EP changes to support Gemma 4 E2B model inference:

Opset 24 kernel registrations:

• Cast op: version the opset 23 registration to 23-23 and add opset 24 registration
• Shape op: version the opset 23 registration to 23-23 and add opset 24 registration
• Updated webgpu_execution_provider.cc registration table accordingly

GQA (GroupQueryAttention) enhancements for Gemma 4:

• Support variable head_dim (Gemma 4 uses head_dim=256 vs the typical 128)
• Make present_key/present_value outputs optional to support KV-shared layers (Gemma 4 has 15 KV layers shared across 35 GQA nodes, where 20 nodes reuse cached KV from earlier layers)
• Add unit tests for variable head_dim and optional present KV outputs

Motivation and Context

Gemma 4 E2B is exported at ONNX opset 24. Without Cast and Shape registrations for opset 24, these ops fall back to CPU, causing unnecessary data transfers and potential failures.

Gemma 4 also uses an unconventional attention architecture with head_dim=256 and KV-shared layers where only 15 of 35 GQA nodes produce new KV cache outputs. The GQA kernel changes enable the WebGPU EP to handle both of these patterns correctly.

🤖 Generated with Claude Code

[Copilot]

Pull request overview

Adds WebGPU EP support for Gemma 4 E2B by registering opset 24 versions of Cast and Shape, and by extending the GroupQueryAttention kernel to handle the model's variable head_dim (e.g. 256) and KV-shared layers where kv_sequence_length==0 and present_key/present_value are optional.

Changes:

• Cast and Shape get explicit opset 23–23 versioned registrations plus new opset 24 registrations (kernels themselves unchanged).
• WebGPU GroupQueryAttention now accepts a missing present_key/present_value and a zero-length new KV input, routing the call through flash attention using past_key/past_value as the KV context; rotary on Q is supported via a dummy K buffer.
• WebgpuAttentionParameters(GroupQueryAttentionParameters) now initializes kv_sequence_length_ from parameters.kv_sequence_length (previously the Q sequence length), and ApplyFlashAttention's internal present K/V buffer shape uses kv_num_heads_ instead of num_heads_.
• New WebGPU tests cross-check shared-KV decode/prompt/GQA-ratio/rotary against CPU.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc	Splits Shape and Cast opset 23 registrations into 23–23 versioned and adds opset 24 registrations.
onnxruntime/core/providers/webgpu/tensor/shape_op.cc	Adds Shape kernel definition for opset 23–23 and bumps the open-ended registration to opset 24.
onnxruntime/core/providers/webgpu/tensor/cast.cc	Adds explicit CreateCastKernelInfo<23,23> and CreateCastKernelInfo<24> template instantiations.
onnxruntime/contrib_ops/webgpu/bert/attention_common.h	Fixes GQA kv_sequence_length_ to read from parameters.kv_sequence_length rather than Q's sequence length.
onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc	Uses kv_num_heads_ for internal present K/V shape; adds a kv_empty branch that aliases past_key/past_value as present_key/present_value via const_cast.
onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc	Drops the hard requirement on present K/V outputs, adds a kv_empty fast path that skips K/V processing (with dummy K for rotary), and errors out if the non-flash path is reached with shared KV.
onnxruntime/test/contrib_ops/group_query_attention_op_test.cc	Adds a use_webgpu parameter to two shared-KV test helpers and four new WebGPU shared-KV tests cross-checked against CPU.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[feich-ms]

Fixed in 9cf7ea3. Sized kDummy/kRotary with parameters.sequence_length_ instead of hardcoded 1, so the kernel's K write offsets stay in bounds during prefill. Also added a new test SharedKV_EmptyKV_WithPast_Rotary_Prompt_WebGPU with q_seq_len=6 to cover this case.

[feich-ms]

Fixed in 9cf7ea3. Added an ORT_ENFORCE(!parameters.past_present_share_buffer_) after the aliasing to guard against future refactors that might invoke CopyKVCache (which writes to present buffers). This makes the contract explicit: the kv_empty path must never reach a write path through present_key/present_value.