onnxruntime

e8b9b24a - [WebGPU] Fuse Q/K RMSNorm into GroupQueryAttention for Qwen3-style models (#28484)

[WebGPU] Fuse Q/K RMSNorm into GroupQueryAttention for Qwen3-style models (#28484) ### Description Adds a per-head Q/K RMS normalization prologue to the WebGPU `GroupQueryAttention` contrib op so that Qwen3-style models can fold the standalone `SimplifiedLayerNormalization` dispatches on Q and K into the attention kernel. The MS-domain `GroupQueryAttention` schema gains two optional inputs and one attribute: * `q_norm_weight` (input 14, optional, shape `(head_size,)`) * `k_norm_weight` (input 15, optional, shape `(head_size,)`) * `qk_norm_epsilon` (FLOAT attribute, default `1e-6`) A new Level 2 optimizer pass `GroupQueryAttentionPreNormFusion` rewrites the `Reshape -> SimplifiedLayerNormalization -> Reshape -> GroupQueryAttention` pattern produced by the Qwen3 ONNX export into a single `GroupQueryAttention` node that carries the norm weights directly. The pass is scoped to the WebGPU EP only. ### WebGPU runtime paths | sequence_length | Behavior | |---|---| | 1 (decode) | RMSNorm is folded into `FusedQKRotaryEmbedding`. | | > 1 (prefill) | Falls back to two standalone `SimplifiedLayerNorm` dispatches into scratch `qNorm`/`kNorm` tensors, then runs the unfused `FusedQKRotaryEmbedding`. Matches the pre-fusion graph timing exactly so prefill cannot regress. | ### Gating The optimizer only rewrites nodes that the WebGPU kernel can handle: * `do_rotary == 1` (the fused decode path runs inside the rotary kernel; the prefill fallback also runs only on the `do_rotary=1` branch) * Non-packed QKV (K at slot 1, V at slot 2 wired) * Already-fused nodes (q/k_norm_weight present) are skipped * The pattern must use a single norm weight per side with matching epsilon and `axis=-1`, and the surrounding reshapes must collapse to `(... head_size)` / expand back to `(... hidden_size)`. ### Defense-in-depth on other EPs The CPU, CUDA, and JSEP `GroupQueryAttention` kernels reject hand-authored models that wire inputs 14/15 with EP-specific error messages, since none of those runtimes implement the prologue: * CPU: `GroupQueryAttention (CPU): q_norm_weight / k_norm_weight inputs are not supported...` * CUDA: `GroupQueryAttention (CUDA): q_norm_weight / k_norm_weight inputs are not supported...` * JSEP: `GroupQueryAttention (JSEP): q_norm_weight / k_norm_weight inputs are not supported...` ### Motivation and Context Perf improvement on **D3D12 backend** for **Qwen3-1.7B** model on a **Windows** machine with **RTX 5060Ti** card <img width="426" height="205" alt="image" src="https://github.com/user-attachments/assets/83e3a0f1-4c2f-4c0f-845d-d5342091d004" />