Repository navigation
[WebGPU] Add per-head Q/K normalization to PagedAttention - #32808
Merged
Merged
Conversation
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
…oft/onnxruntime into aciddelgado/qwen38_webgpu
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Added enforcement for is_causal attribute values and updated its assignment. Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
…upport-stateupdatecapacity
…elgado/qwen38_webgpu Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
aciddelgado
changed the title
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Active state-update tests assert unspecified tail data, and substantial non-causal and GatedDeltaNet changes are missing from the stated scope.
Get a fresh assessment by requesting another Copilot review.
Review effort: Balanced
Findings: 3
Open (3)
What changed in this PR
Adds WebGPU PagedAttention Q/K normalization while also expanding non-causal attention and GatedDeltaNet compact-update support.
Changes:
- Adds FP32 Q/K RMSNorm preprocessing before RoPE.
- Enables non-causal and local-window attention paths.
- Adds WebGPU GatedDeltaNet compact state updates and tests.
| File | Description |
|---|---|
onnxruntime/​test/​python/​transformers/​test_paged_attention.py |
Adds Q/K norm and masking tests. |
onnxruntime/​test/​contrib_ops/​paged_attention_op_test.cc |
Adds non-causal reference coverage. |
onnxruntime/​test/​contrib_ops/​gated_delta_net_op_test.cc |
Adds compact-update coverage. |
onnxruntime/​core/​providers/​webgpu/​nn/​layer_norm.h |
Exposes FP32 normalization mode. |
onnxruntime/​core/​providers/​webgpu/​nn/​layer_norm.cc |
Implements FP32 affine arithmetic. |
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc |
Updates GatedDeltaNet schema documentation. |
onnxruntime/​contrib_ops/​webgpu/​bert/​paged_attention.h |
Stores causal mode. |
onnxruntime/​contrib_ops/​webgpu/​bert/​paged_attention.cc |
Adds normalization and non-causal dispatch. |
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.wgsl.template |
Emits compact state updates. |
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.h |
Adds capture and clear programs. |
onnxruntime/​contrib_ops/​webgpu/​bert/​gated_delta_net.cc |
Enables compact-update execution. |
onnxruntime/​contrib_ops/​webgpu/​bert/​flash_attention.wgsl.template |
Revises local-window masking. |
onnxruntime/​contrib_ops/​webgpu/​bert/​flash_attention_decode_qkv.wgsl.template |
Adds generalized visible-key bounds. |
onnxruntime/​contrib_ops/​webgpu/​bert/​attention_common.h |
Propagates causal configuration. |
docs/​design/​webgpu_paged_attention.md |
Documents normalization and masking. |
docs/​ContribOperators.md |
Regenerates GatedDeltaNet documentation. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Preserve Q/K normalization coverage alongside main's updated WebGPU regression tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
|
Copilot resolve the merge conflicts in this pull request |
…orm_webgpu # Conflicts: # onnxruntime/core/providers/webgpu/nn/layer_norm.cc Co-authored-by: aciddelgado <139922440+aciddelgado@users.noreply.github.com>
Contributor
Merged the latest |
Baiju Meswani (baijumeswani)
left a comment
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I found two focused follow-ups for the WebGPU Q/K-normalization change and the CUDA routing repair. The suggestions are inline.
aciddelgado
added a commit
that referenced
this pull request
Sep 30, 2026
### Description
<!-- Describe your changes. -->
Fix the CUDA QMoE packed-INT GEMV path's remaining reference to the
removed
`routing_record` API and add a focused expert-counting regression.
- Capture selected expert IDs immediately after routing and consume the
snapshot
before the packed-INT path returns, using the same lifecycle as the
neighboring
paths. Preserve the existing non-plugin/non-minimal guards.
- Add `MoeExpertCountingTest.CudaQMoEPackedIntGemv`: FP16, INT2, fused
SwiGLU,
three rows, four experts, top-k 1, 512-wide hidden/intermediate
dimensions,
and block size 64.
- Verify expert-dependent numerical outputs and per-invocation selection
counters:
duplicate routes `[0, 0, 2]`, a repeated invocation, then changed routes
`[1, 1, 1]`. Expected counters are `[0.1, 0, 0.1, 0]`,
`[0.19, 0, 0.19, 0]`, and `[0.171, 0.1, 0.171, 0]`.
- Disable CPU fallback and cap dense dequantization scratch at one byte,
so a
successful run cannot silently use the dense fallback.
Only `moe_quantization.cc` and `moe_expert_counting_test.cc` change.
Routing
selection and GEMV arithmetic are unchanged; no WebGPU/PagedAttention
changes
are included.
**Validation**
Built and executed locally on NVIDIA A100-SXM4-80GB (SM80), CUDA
12.8.93,
GCC 11.4, cuDNN 9.10.2.21, Release, warnings-as-errors, QUICK_BUILD=OFF.
The existing CUDA build tree was reconfigured from SM86 to SM80,
recompiling
the CUDA targets. This is not a claim of CUDA 13/GCC 14 CI validation.
After configuring with the same options plus `--update`, the following
build
command passed:
```bash
source /workspace/venv/bin/activate
python tools/ci_build/build.py \
--build_dir /workspace/onnxruntime/build/cuda-ci-20260929 \
--config Release --build --parallel 16 --cmake_generator Ninja \
--build_shared_lib --build_wheel --enable_onnx_tests --use_cuda \
--use_binskim_compliant_compile_flags --nvcc_threads 4 --flash_nvcc_threads 4 \
--cuda_version=12.8 --cuda_home=/usr/local/cuda-12.8 \
--cudnn_home=/workspace/venv/lib/python3.11/site-packages/nvidia/cudnn \
--enable_cuda_profiling --skip_tests --allow_running_as_root \
--targets onnxruntime_test_all onnxruntime_provider_test \
--cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=80 \
onnxruntime_BUILD_UNIT_TESTS=ON \
onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS=ON \
onnxruntime_BUILD_CUDA_QUANT_PREPROCESS=ON
```
Tests were run from the build output directory:
```bash
cd /workspace/onnxruntime/build/cuda-ci-20260929/Release
export LD_LIBRARY_PATH=/workspace/venv/lib/python3.11/site-packages/nvidia/cudnn/lib:/usr/local/cuda-12.8/lib64:${LD_LIBRARY_PATH:-}
./onnxruntime_test_all --gtest_filter='MoeExpertCountingTest.CudaQMoE*'
# PASS: 4 tests, no skips.
./onnxruntime_provider_test --gtest_filter='MoETest.QMoETest_Int2CudaPackedDecodeMultiRowExpertParity'
# PASS: 1 test, no skips.
ORT_ENABLE_QMOE_KERNEL_DEBUG_INFO=1 ./onnxruntime_test_all \
--gtest_filter='MoeExpertCountingTest.CudaQMoEPackedIntGemv'
# PASS: 1 test; all three invocations log Route=packed_int_gemv Rows=3.
```
**Negative control:** temporarily omitted both packed-path `Capture` and
`Consume`, rebuilt with
`cmake --build build/cuda-ci-20260929/Release --target
onnxruntime_providers_cuda --parallel 16`,
and ran
`./onnxruntime_test_all
--gtest_filter='MoeExpertCountingTest.CudaQMoEPackedIntGemv'`.
It failed as intended: outputs still matched, but seven counter
assertions
failed because all counters remained zero (first failure: expert 0
expected
0.1, actual 0). Restored the repair, rebuilt the provider, and reran
both
focused suites: all five tests passed. No negative-control edits are
committed.
### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
The expert-counting migration in #32738 left a stale
`routing_record->CaptureTile`
call in the packed-INT GEMV early-return path. This prevents ordinary
built-in
CUDA compilation on current main; simply deleting that call would
compile but
silently omit expert counts for this path.
This separates the repair originally included in #32808, following
[the scope/coverage review
comment](#32808 (comment)).
The production patch is carried unchanged from commit `1f57e92abe`; the
new
regression verifies the runtime counting behavior as well as
compilation.
After this lands, merging main into #32808 will remove the unrelated MoE
file
from that PR's diff.
---------
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
aciddelgado
enabled auto-merge (squash)
September 30, 2026 01:27
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Baiju Meswani (baijumeswani)
left a comment
Baiju Meswani (baijumeswani)
left a comment
Contributor
There was a problem hiding this comment.
I did not find a confirmed correctness bug in the Q/K-normalization implementation. These two follow-ups would help qualify the Qwen3.8 WebGPU path.
Baiju Meswani (baijumeswani)
approved these changes
Oct 2, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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
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:
Normalization adds two preprocessing dispatches and Q/K scratch tensors. Fusing normalization with rotary is outside this PR’s scope.