onnxruntime
e1f27d15 - [WebGPU] Support models with opset 24 ops and KV-shared decoder layers (Gemma 4) (#28501)

Commit
79 days ago
[WebGPU] Support models with opset 24 ops and KV-shared decoder layers (Gemma 4) (#28501) ### Description This PR adds WebGPU EP changes to support models with opset 24 ops and KV-shared decoder layers (like Gemma 4): **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) kv_empty path for KV-shared layers:** - Support `kv_sequence_length==0` where layers reuse another layer's KV cache via `past_key`/`past_value` - Apply Q-only rotary embedding by reusing `RotaryEmbeddingProgram` through a shared `RunRotaryEmbedding` helper. When invoked with `use_seqlens_for_position=true`, the shader derives `position_id` per batch directly from the `seqlens` tensor (`past_seqlen = seqlens[batch] + 1 - global_shape[1]`, then `position_id = past_seqlen + sequence_index`) — no host-side offset uniform and no scratch position_ids tensor are needed - `RunRotaryEmbedding` is now a single shared helper used by both the contrib GQA kernel (kv_empty Q-only rotary) and the standalone RotaryEmbedding kernel (`ComputeInternal`); the boolean `use_seqlens_for_position` toggles between the seqlens-derived path and the legacy `position_ids` path - Skip wasteful `present_key`/`present_value` GPU allocation for kv_empty path (aliased to past instead) - Make `present_key`/`present_value` outputs optional (nullptr when model doesn't request them) - Fix `kv_sequence_length_` initialization in `WebgpuAttentionParameters` (was incorrectly using Q sequence length) - Fix `present_kv_heads` to use `kv_num_heads_` instead of `num_heads_` for GQA internal buffer shapes **Tests:** - `WebGPU_SharedKV_Decode` — kv_empty decode path (S=1) - `WebGPU_SharedKV_Prefill` — kv_empty prefill path (S>1, tiled attention) - `WebGPU_SharedKV_Rotary` — kv_empty with Q-only rotary embedding - `WebGPU_SharedKV_Rotary_Prefill` — Q-only rotary with multi-token prefill (q_seq=4) - `WebGPU_SharedKV_Rotary_MultiBatch` — Q-only rotary with batch_size=2 - `WebGPU_SharedKV_SlidingWindow` — kv_empty with sliding window attention All WebGPU tests cross-check against CPU for numerical correctness. ### Motivation and Context Models with KV-shared decoder layers (e.g. Gemma 4 E2B: 20 of 35 layers share KV cache) pass `kv_sequence_length=0` and empty K/V inputs to GQA nodes that reuse another layer's cached KV. The WebGPU GQA kernel previously rejected this pattern. These changes enable correct and efficient execution by: 1. Applying rotary to Q only via the shared `RunRotaryEmbedding` helper with `use_seqlens_for_position=true`, which lets the shader derive each token's position from `seqlens[batch_idx]` (K is already rotated in the shared cache) 2. Aliasing past as present (avoiding unnecessary allocation and copy) 3. Reusing the existing `RotaryEmbeddingProgram` — no new shader classes needed; the only addition is a `use_seqlens_for_position_` branch inside `GenerateShaderCode` that selects between reading `position_ids` and computing `position_id` from `seqlens` Opset 24 Cast/Shape registrations prevent fallback to CPU for models exported at opset 24. --------- Co-authored-by: Claude Opus 4 <noreply@anthropic.com> Co-authored-by: Jiajia Qin <jiajiaqin@microsoft.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Author
Parents
Loading