onnxruntime
29feb3ba - [WebGPU] Tile large block-FP8 QMoE expert matmuls (#33075)

Commit
10 days ago
[WebGPU] Tile large block-FP8 QMoE expert matmuls (#33075) ## Summary Follow up on #32794 with an optimized **routed expert** matmul for block-scaled FP8 QMoE. This is separate from the dense MatMul kernel in #33074. - Stage selected-expert E4M3 weights and FP16 activations into shared tiles and use supported 16x16x16 FP16-input subgroup-matrix instructions with FP32 accumulation. FP8 is decoded before the matrix instruction; this is **not native FP8 matrix hardware**. - Handle per-128x128 scales, signed FP8 codes, FC1/FC2 bias, odd-byte expert boundaries, and incomplete matrix/scale tiles. - Only dispatch the matrix path for direct expert batches with at least 128 rows and 256 input/output channels on adapters exposing the required subgroup configuration. Preserve the original scalar path for single-token decode, smaller routed groups, and other adapters. A 16-row synthetic case regressed with the matrix path, so it is deliberately excluded. ## Validation - Linux Release static WebGPU provider, provider tests, and Python WebGPU runtime built successfully. - All 17 focused WebGPU QMoE tests passed on NVIDIA H200/Vulkan, including a new 128-row, 257-channel case covering signed FP8 values, distinct 128x128 block scales, bias, tail tiles, and non-word-aligned expert offset. - Synthetic end-to-end QMoE with 128 rows, 2560 hidden channels, 640 intermediate channels, two experts, and 30 repeated calls: matrix path 4.4–5.1 ms/call versus original scalar 6.17 ms/call in this harness (about 1.2–1.4x). This includes routing and other kernel work; it is **not isolated matrix-kernel timing**. At 16 rows the original scalar path remains enabled (about 3.7 ms/call). - `clang-format --dry-run --Werror` and `git diff --check` passed. No official block-FP8 checkpoint or end-to-end model-quality claims: the local checkpoint is BF16. Stack base is #32794 for review; merge the PRs individually when ready. ## Stack update (October 5) Restacked onto #32794 after its separate-FC3 regression was added. The review diff remains limited to the three QMoE optimization files. All 18 focused stacked WebGPU QMoE tests passed on H200/Vulkan; CI is rerunning for the updated base. ## Current main update (October 5) Merged the latest #32794 base, which itself now includes current `main` and merged INT2/mixed-width QMoE support. The review diff remains the three tiled-FP8 files. All 15 focused H200 block-FP8 and merged INT2 compatibility tests pass; new-head CI is running. ## Review follow-up (October 6) - Apply FP32 block scales to FP32 subgroup-matrix partial sums, not FP16 staged weights; a +448/-448 cancellation test with scale 200 checks finite output. - Keep matrix workgroup dimensions within the adapter limit, falling back to scalar when its dispatch is legal; reject oversized scalar dispatch through the updated #32794 base. - Add mixed 128+4 row routing with distinct expert FC1/FC2/FC3 weights, scales, and nonzero biases, an independent SwiGLU reference, and an adapter-aware assertion that all three eligible projections actually dispatch the matrix path. - After merging #32794 head `40cebdeeef`, built the Linux Release provider tests and passed all 30 focused WebGPU QMoE tests on H200/Vulkan; formatting and `git diff --check` passed. This does not imply other platform CI has completed. ## October 6 base update Merged updated FP8 QMoE base 09737a35ac, inheriting orphan FC3-scale validation and synchronized operator docs that address the Windows GPU documentation failure (stacked head 23922ce3a7). All 31 focused H200 WebGPU QMoE tests pass; fresh CI is pending. Previously drafted replies to three review threads have now been published. ## October 7 base update Merged updated #32794 base eb51a28fcd to inherit explicit rejection of unused FP8 activation-scale inputs. The new four-case regression and all 32 focused H200 WebGPU QMoE tests pass on stacked head 158c7d0abd. Fresh CI is pending. ## October 7 CI stack update Merged current main and updated #32794 (541e5c7606), including the already-merged upstream CUDA GQA `is_int8` build fix #33179. The static WebGPU provider tests build and all 32 focused H200 routed-QMoE tests pass. New-head CI is pending. ## October 7 upstream CI fix merged Restacked on updated #32794 at 79873aeb9a, including merged upstream #33168 ONNX node-test type conversion. Static WebGPU provider tests build and all 32 focused H200 routed-QMoE cases pass. Fresh CI is pending; previous-head failures are not evidence for this head. ## October 8 merged-base refresh and paired performance Merged current main after #32794 landed. The review diff remains the QMoE matrix-path implementation and its tests. In an integrated branch with #33058 and #33074, 52/52 H200 WebGPU QMoE and dense FP8 regressions passed. For a synthetic FP16-input, 2-expert block-FP8 QMoE (128 tokens, hidden 2560, intermediate 640), ONNX Runtime perf_test on H200/Vulkan measured 7.86 ms/run on the merged #32794 implementation versus 6.05 ms/run with the tiled path (median of three 100-run trial averages, ~1.30x). These are whole-session inference latencies, not isolated kernel timings or an official-FP8-checkpoint result. Current-head CI is pending. ## October 9 main refresh and subgroup ownership fix Merged main at 9a6089e176, including merged #33058 and its gate tests/documentation. The PR now targets the merged FP8 QMoE implementation on main; #32794 is no longer an unmerged stack prerequisite. Addressed the new subgroup-ownership review in 7567ecd231: the shared shader entry point exposes the actual `@builtin(subgroup_id)` as `subgroup_index`, and the QMoE matrix path uses it for accumulator initialization, left-operand loading, and result storage. Cooperative tile transfers still use local invocation indices. This matches the dense FP8 implementation in #33074. Validation on the updated head: - Linux Release WebGPU provider-test build succeeded. - 38/38 focused H200/Vulkan WebGPU QMoE tests passed, including the matrix-capability-aware mixed-routing/separate-FC3 regression and the merged gate cases. - 89/89 WGSL-template tests passed. - Captured GPU shaders use the subgroup-ID builtin for all three matrix tile offsets; no `lane / 32u` offset remains. - `clang-format --dry-run --Werror` and `git diff --check` passed. The subgroup-mapping fix was executed on H200, not an adapter with a noncontiguous subgroup mapping. Fresh GitHub CI is pending; the local results are not a claim that all platform CI has completed. --------- Co-authored-by: Sayan Shaw <sayanshaw@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading