[WebGPU] Add per-head Q/K normalization to PagedAttention (#32808)
**Summary** Adds WebGPU PagedAttention support for q_norm_weight , k_norm_weight , and qk_norm_epsilon , removing the unsupported-input error encountered by the Qwen3.8 / DFlash2 path. The implementation reuses the existing WebGPU normalization pass to apply RMSNorm independently to each query head and each newly appended key head. Reduction and gain multiplication use FP32, with normalized Q/K materialized as FP16 before rotary embedding, matching CUDA/reference arithmetic ordering. **Changes** - Wire Q/K normalization into PagedAttention preprocessing, supporting separate and packed QKV inputs, with or without RoPE. - Add an opt-in FP32 arithmetic mode to the shared normalization shader. Existing callers retain their current behavior. - Preserve previously cached K and all V values; only newly appended K is normalized. - Validate positive, finite epsilon when normalization is enabled. - Update WebGPU PagedAttention documentation. No changes to attention kernels, masking, block-table layout, or direct-paged dispatch restrictions. Behavior without norm weights is unchanged. **Validation** Built and installed a local native WebGPU wheel and ran on NVIDIA A100-SXM4-80GB: - Full test_paged_attention.py : 161 passed, 281 CUDA-only tests skipped, 0 failed. - Added 24 cases covering GQA, prefill/decode, ragged batches, packed QKV, rotary, non-causal local windows, and epsilon/FP16 edge cases. - Added deterministic CPU-reference regressions distinguishing per-head normalization from full-hidden normalization and verifying cached K remains unchanged across consecutive calls. - Verified both normalization shader branches and ran 12 additional checks of existing normalization behavior with CPU fallback disabled. Normalization adds two preprocessing dispatches and Q/K scratch tensors. Fusing normalization with rotary is outside this PR’s scope. --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Prathik Rao <prathikrao@microsoft.com> Co-authored-by: Prathik Rao <prathik.rao@gmail.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
A
aciddelgado committed
c971ca264e0fd7b505d17e4e45ffaed5b795ff67
Parent: 2cdff3c
Committed by GitHub <noreply@github.com>
on 10/2/2026, 5:39:16 AM