SIGN IN SIGN UP

[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