[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>