Repository navigation
Add WebGPU INT2 and mixed-width QMoE support - #32878
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Homogeneous INT2, FC3-specific widths, and multi-token mixed-width execution lack direct WebGPU coverage.
Review effort: Balanced
Findings: 3
Open (3)
What changed in this PR
Adds WebGPU INT2 and mixed-width QMoE execution support.
Changes:
- Accepts 2-bit expert weights.
- Dispatches FC1/FC2/FC3 with independent bit widths.
- Adds single-token mixed-width numerical validation.
| File | Description |
|---|---|
qmoe.h |
Permits INT2 weights. |
qmoe.cc |
Applies projection-specific bit widths. |
moe_test.cc |
Tests INT2 FC1 with INT4 FC2. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Reject partial u32-packed FC1, FC2, and FC3 rows and blocks before MatMulNBits dispatch. Test inherited INT2, routed mixed-width decode and multi-token experts, and independent FC3 widths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Use the non-fused SwiGLU branch when testing independent INT2 FC3 packing and its alignment rejection. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
kunal-vaishnavi
approved these changes
Oct 5, 2026
Contributor
There was a problem hiding this comment.
Warning
Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.
Copilot review overview
Review effort: Lite
Findings: 1
Open (3)
Resolved since last review (5)
This was referenced Oct 5, 2026
kunal-vaishnavi
pushed a commit
that referenced
this pull request
Oct 8, 2026
## Summary Adds direct WebGPU W8A16 support for QMoE expert weights stored as FP8 E4M3 with 128x128 FP32 block scales, matching the Qwen 3.8 Flash Next MTP expert checkpoint format without requantization or dense weight materialization. ## Changes - Extends the existing `quant_type="fp8"` QMoE mode: `block_size=128` consumes `fc*_scales` shaped `[E, ceil(N/128), ceil(K/128)]`; omitted block size retains the existing CUDA global-scale FP8 contract. - Adds a selected-expert WebGPU shader that decodes raw E4M3 bytes and applies block scales during FP16/FP32 accumulation. - Supports direct expert dispatch for prefill and indirect top-k expert dispatch for single-token decode. - Keeps integer, FP4, NVFP4, and WFP4AFP8 behavior unchanged. - Adds exact two-expert WebGPU coverage with distinct per-expert scale products. ## Historical validation (initial Windows run) - Windows Release `onnxruntime_providers_webgpu` builds successfully. - Windows Release `onnxruntime_provider_test` builds successfully. - `MoETest.QMoETest_WebGPU_BlockFp8_128x128` passes locally on D3D12. - `clang-format --dry-run --Werror` and `git diff --check` pass. - Editor diagnostics report no errors in changed files. ## Target / follow-up - Targets `main` directly; the WebGPU MoE/QMoE prerequisite from #32447 is merged. This PR's review diff is limited to six FP8-related files, independent of the Qwen integration branch. - GenAI export wiring for MTP block-FP8 QMoE and full-model chat/performance measurements remain follow-up work. The INT4 RTN model measurements are not validation of the block-FP8 checkpoint. ## Historical validation (earlier main restack) - Linux Release static WebGPU provider and provider-test builds passed on current `main`. - All 14 `MoETest.QMoETest_WebGPU_*` tests passed on NVIDIA H200/Vulkan, including single-token decode, a 3x3 non-word-aligned expert-boundary case, a 129x129 four-scale-tile case, and negative integer-weight validation. - Focused `clang-format --dry-run --Werror` and `git diff --check` passed. ## Historical review follow-up (October 5) Added a separate-FC3 block-FP8 single-token test covering k=2 indirect expert routing, broadcast input, SwiGLU, and distinct per-expert FC1/FC2/FC3 scales. All 17 focused QMoE WebGPU tests passed on H200/Vulkan. All review threads have replies and are resolved; CI is rerunning for the new commit. ## Historical main update (October 5) Merged current `main` after #32878 landed and reconciled block-FP8 routing with its INT2/mixed-width QMoE logic. The review diff remains limited to the FP8 QMoE files. An isolated Linux static WebGPU provider/test build and all 14 focused block-FP8 plus INT2 compatibility tests passed on H200; the stacked tiled follow-up also passed 15 focused H200 tests. New-head CI is running. ## Latest validation (October 6, `40cebdeeef`) - Merged upstream `main` at `aac791ef2408`, resolved the QMoE schema/documentation conflicts, and confirmed GitHub reports this head as mergeable. - The Linux static WebGPU `onnxruntime_provider_test` build passed; all 26 focused `MoETest.QMoETest_WebGPU_*` tests passed on NVIDIA H200/Vulkan. The numerical block-FP8 reference now covers both FP16 and FP32 inputs/outputs, including FP8 decoding, accumulation, and biases. - Before FC1, FC2, or optional FC3 dispatch, the WebGPU path rejects workgroup dimensions that would be normalized and projections exceeding 32-bit shader addressing. Direct expert byte offsets and view sizes use 64-bit arithmetic; the shared limit also covers indirect expert indexing. The oversized-dispatch rejection has an H200 regression test. - The six review threads concerning FP32 coverage, dispatch normalization, and expert-offset overflow have replies and are resolved. New-head CI is still pending; earlier test counts and CI statements above describe their respective historical heads, not the current result. ## October 6 follow-up Reject orphan block-FP8 FC3 scales instead of silently ignoring them (09737a3); 27 focused H200 WebGPU QMoE tests pass. Refreshed generated ContribOperators.md to match the schema and address Windows GPU documentation CI. Fresh CI is pending. ## October 7 review follow-up Reject all four optional activation-scale inputs (17–20) in WebGPU block-FP8 mode instead of silently ignoring them. The GPU-only negative regression supplies each input separately. Updated the schema and generated operator docs to state that activation scales are not accepted. All 28 focused H200 WebGPU QMoE tests pass on eb51a28; new-head CI is pending. ## October 7 CI follow-up Merged current main (6030cc6), including upstream CUDA GQA fix #33179: the previous Linux/Windows CUDA and TensorRT builds failed because `is_int8` was used before its declaration in upstream CUDA code, not in this PR's WebGPU changes. The updated Linux static WebGPU provider test builds and all 28 focused H200 QMoE tests pass. Previous-head Python 3.14 and Windows QNN jobs failed ONNX materialized-node tests; Android minimal binary size exceeded its threshold by 1,494 bytes. New-head CI is running; these failures are not yet claimed fixed until the fresh checks report. ## October 7 upstream CI fix merged #33168 merged to main and corrects the ONNX wire-type to C-API-type conversion behind the old-head Windows QNN/Python 3.14 materialized-node failures. Merged main at fdf7f9d; static WebGPU tests build and all 28 focused H200 QMoE cases pass. The old-head Android binary-size check passed. New-head CI is running; no claim of CI green until the fresh checks finish. --------- Co-authored-by: Sayan Shaw <sayanshaw@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
kunal-vaishnavi
pushed a commit
that referenced
this pull request
Oct 9, 2026
## Summary - Replace the single-token QMoE gate full expert bubble sort with a stable parallel TopK reduction for normalized routing. Preserve descending-logit/index-ascending ties, softmax weight order, and the existing full sort when routing weights are not normalized (the denominator then needs all experts). - Cover boundary ties in 3-, 5-, 6-, and 512-expert pools with nonzero expert outputs. Skip the 512-expert case when either adapter workgroup limit is below 512. - Assert the selected gate indices directly for FP16/FP32 negative-infinity logits (`[0, -inf, -inf]` selects `[0, 1]`) and for all-equal logits with `k == num_experts` (`[0, 1, 2, 3]`). The negative-infinity assertion fails with the original comparator (it selects the invalid sentinel index 3). - Exercise FP16/FP32 nonnormalized routing with an excluded expert contributing appreciably to the all-expert softmax denominator. ## Validation - On head `45037061757c`, merged `origin/main` at `9f6baac7319a` (including #33168) without rebasing or gate conflicts. Rebuilt both the static test-all and split provider-test targets; all 28 focused WebGPU MoE tests passed on H200, including FP16/FP32 gate-index readback. Checked free GPU memory before the run; fresh CI for this head is in progress. - Main-based static WebGPU build and all 22 focused WebGPU QMoE tests passed on H200 after merging `main` at `aac791ef2408a266d8ee59dd5fa0ce11d1a5a0bb`, including the new tie, gate-index, and nonnormalized-routing regressions. - Qwen3.8-Flash-Next INT4 RTN/INT8-head export, 512 experts, k=10, FP16 I/O, one-token chunks, H200: 16-token prefill 11.64 s -> 0.90 s and steady decode ~0.674 s -> ~0.045 s/token. Baseline versus optimized logits were bit-identical (248,320 logits), and the first five generated token IDs matched. This full-model comparison used a QSA-enabled experimental build carrying the same single-token gate change because the export also requires the not-yet-merged SparseAttentionIndexer kernel; the isolated main-based plugin cannot load that export today. The PR itself does not depend on the QSA changes. - This PR changes only the single-token gate. Separate GPU-only multi-token prefill experiments are intentionally excluded pending model-quality validation. ## Main updates - October 5: merged `main` and added the follow-up requested on merged #32878: two experts use distinct INT2/INT4 packed weights as well as scales, with matching FP16-aware numerical references. The fixture explicitly rejects unsupported 8-bit encodings. - October 6: merged `main` at `aac791ef2408a266d8ee59dd5fa0ce11d1a5a0bb`; added direct gate-index assertions, 5-/6-expert pool tests, adapter-limit skip, and the full-denominator regression. That head was superseded by the fixture alignment fix below. - October 6 follow-up: pad the direct gate-index test fixture buffers to 4-byte boundaries for odd FP16 expert counts. On the prior head, macOS Dawn rejected a six-byte storage buffer and an eight-byte staging copy. The prior-head Windows job reported a lost device at the same gate-index test (0 ms), followed by cascading failures; its log does not establish whether the fixture caused the device loss. - October 7 follow-up: current-head arm64 Debug reproduced a real test-fixture readback error: `BufferManager::Download` rounds the copy to 16 bytes, but the `k=2` indices buffer was only 8 bytes. Align that buffer to 16 bytes. Windows novcpkg lost the WebGPU device at the same gate-index test and subsequently failed 1,538 unrelated tests; the static vcpkg aggregate reported the dependent test failure despite successful compilation. All 28 focused WebGPU MoE tests pass after the fix on H200. - October 7 main refresh: merged `main` at `9f6baac7319a` (including #33168), rebuilt the split static WebGPU provider-test target, and passed all 28 focused H200 WebGPU MoE tests on head `45037061757c`. Prior-head CI failures do not describe this new head; current-head checks are running. ## October 8 main refresh Merged `origin/main` at `864d87126864c64651f58a9b51d9006ad71b2c78` (including #33203) into the gate branch, resolving the test include conflict by retaining both sets of includes. Updated the direct gate-buffer test to call upstream `WebGpuDevice(0)` after the context-ID API change. On head `813f7999e55de5a8d29a7a02dfc869d514f327b4`, the static WebGPU provider-test build succeeded and all 34 focused H200/Vulkan `MoETest.QMoETest_WebGPU*` cases passed. Fresh CI is queued; no new-head CI result is claimed yet. ## Matched synthetic whole-inference benchmark (October 8) H200/Vulkan ONNX Runtime `perf_test`, synthetic single-token block-FP8 QMoE with 256 experts and top-4 routing, three 100-run trial averages: merged #32794 baseline **7.09 ms/run** median-trial mean versus the integrated gate branch **5.56 ms/run** (~**1.28x**). Individual trial averages ranged **5.71–7.12 ms** for baseline and **5.46–5.89 ms** for the integrated branch; these overlapping ranges warrant a variance caveat. This measures whole inference, **not** the isolated gate kernel or a full checkpoint. The separate combined branch passed **52/52** focused tests; that result is not claimed as an isolated #33058-head test run. ## October 8 documentation follow-up Added [WebGPU QMoE single-token gate documentation](https://github.com/sayanshaw24/onnxruntime/blob/8d38e0e31be0216e7d3b937578efa113fccee4ca/docs/contrib_ops/webgpu/qmoe_gate.md) on head `8d38e0e31b`, covering normalized top-k, tie and -inf handling, nonnormalized fallback, adapter limits, and benchmark scope. The previous-head Windows documentation-build job stopped while compiling unrelated CUDA sources before document generation, so its dependent validation failure does not establish a generated-doc mismatch. The previous-head Android job failed during NDK setup with a corrupt downloaded archive. Fresh-head CI is running. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Sayan Shaw <sayanshaw@microsoft.com>
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.
Description
Adds WebGPU execution support for INT2 and mixed-width integer
QMoE, following the operator contract introduced for CUDA in #32761.MatMulNBitsimplementation.The underlying WebGPU
MatMulNBitsshaders already implement 2-bit unpacking, so no new shader path is required.Validation
k=2indirect selection, different experts across two tokens, separate FC3, and unaligned-row rejection.Follow-up
The original mixed-width tests distinguished experts using different scales but identical packed-weight values. The separately reviewable follow-up in #33058 gives each expert distinct packed-weight values and updates the numerical reference to cover weight-offset errors. The new tests passed on H200; this follow-up is not part of the merged #32878 commit.