feat(webgpu): add preferredMatmulAccumulatorPrecision provider option for MatMulNBits - #29599
feat(webgpu): add preferredMatmulAccumulatorPrecision provider option for MatMulNBits#29599Roberto (RobertoReale) wants to merge 13 commits into
Conversation
…nt overflow Problem: - fp16 and q4f16 models (Gemma 3 270M, Whisper Q4, SmolLM2, kokoro-js) produce garbage output on WebGPU. f16 max is ~65504; summing 2048+ dot-product terms overflows to +Inf, which propagates as NaN through LayerNorm/Softmax. - CPU/WASM path is unaffected (MLAS accumulates in f32 internally). Fix (JSEP path): - matmulnbits.ts: both kernels (default and BlockwiseMatMulNBits32) now accumulate in f32 (workgroup_shared / inter_results). Explicit vec<f32>() casts per operand are required: Dawn/D3D12 re-demotes temporaries to f16 when 'enable f16;' is active. Output downcast to f16 only at the write. - matmul-shaders.ts (naive MatMul): values accumulator promoted to f32. - 3rd-party/matmul_packed_webgpu.ts (tiled MatMul, vec4 + scalar variants): acc promoted to f32; tiles (mm_Asub/mm_Bsub) stay in f16, so there is no shared-memory or bandwidth regression. Fix (native WebGPU EP): - contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template: inter_results and the final reduction accumulate in f32 along K; downcast at output write. This aligns the kernel with matmul_nbits_wide_tile.wgsl.template, which already accumulates in f32. Testing: - scripts/verify_f16_fix.py: static source check - all checks PASS - cd js/web && npx tsc --noEmit: clean; prettier: clean - Manual browser test on Chrome 121+ with Gemma 3 270M FP16 and SmolLM2 360M Q4F16 via transformers.js Performance note: f32 accumulators only affect register-level computation; weights/activations remain f16/q4 in GPU memory, so memory bandwidth (the actual bottleneck) is unaffected. Fixes microsoft#26732 Related: microsoft#26367
|
@microsoft-github-policy-service agree |
Additional real-world validationI validated this fix end-to-end against a production consumer of Method: since this fix only touches the JSEP TypeScript path ( Result (whisper-small, q4, ~36s of Italian audio):
No Only whisper-small was tested this way so far; the fix is generic to the shared MatMul/MatMulNBits kernels rather than model-specific, so I'd expect the same result on tiny/large-v3-turbo (also q4) — happy to confirm those too if useful. Full write-up (methodology + how I'll re-enable WebGPU once this merges): https://github.com/RobertoReale/Voice-Message-Transcriber/blob/main/docs/webgpu-onnxruntime-fix.md |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
Roberto (@RobertoReale), could you run |
Tianlei Wu (tianleiwu)
left a comment
Tianlei Wu (tianleiwu)
left a comment
There was a problem hiding this comment.
Thanks for the fix — the root-cause analysis (f16 dot-product accumulation saturating past 65504 for D >= 2048) is correct and the JSEP-side changes (promoting values/acc/workgroup_shared/inter_results to f32 and downcasting only at the final write) look right. A few points before this is complete:
1. The native WebGPU EP fix is incomplete (high priority). Only matmul_nbits.wgsl.template is patched, but the same overflow-prone f16 accumulation pattern exists in the sibling fused kernels that ORT dispatches for the same models:
matmul_nbits_mlp.wgsl.template—gate_sum/up_sum = output_element_t(0)andgate_inter_results/up_inter_results : array<array<output_element_t, ...>>matmul_nbits_qkv.wgsl.template—sum = q_output_element_t(0)andq/k/v_inter_results : array<array<q_output_element_t, ...>>dp4a_matmul_small_m.wgsl.template—inter_results : array<array<output_element_t, ...>>
The MLP and QKV fused kernels are selected precisely for the fused MLP / attention-QKV subgraphs that Gemma 3 / SmolLM produce, so those paths can still emit garbage for the models this PR targets. Either extend the same f32-accumulation change to these templates or explicitly document why they are out of scope.
2. Shared-memory / occupancy claim is imprecise. "Zero bandwidth regression / no shared-memory regression" is accurate for the A/B tiles (they stay in the input dtype), but the accumulator workgroup arrays (inter_results, workgroup_shared) do double in size (f16 -> f32) for fp16/q4f16 outputs. On SLM-bound dispatch configs this can lower occupancy. Correctness rightly wins here, but the description/comments should acknowledge the accumulator SLM growth rather than claim zero regression.
3. Missing automated regression coverage. Validation currently relies on manual browser runs plus a static grep script. A numerical unit test (large-K fp16 MatMul / MatMulNBits compared against an f32 reference) in the existing WebGPU op test suite would guard against this regressing again and would exercise the actual runtime rather than source strings.
Inline comments below.
| // Accumulate partial sums along K in f32: with fp16 outputs, summing 2048+ f16 | ||
| // products overflows the f16 max (65504) and poisons the output with +Inf/NaN | ||
| // (issue #26732). Only the block-local `sum` stays in output_element_t. | ||
| var<workgroup> inter_results: array<array<f32, tile_size_k_vec>, tile_size>; |
There was a problem hiding this comment.
This correctly fixes the generic kernel, but the same f16-accumulation pattern is left unfixed in the fused native kernels that get dispatched for the very models this PR targets:
matmul_nbits_mlp.wgsl.template:gate_sum/up_sum = output_element_t(0)andgate_inter_results/up_inter_results : array<array<output_element_t, ...>>matmul_nbits_qkv.wgsl.template:sum = q_output_element_t(0)andq/k/v_inter_resultsdp4a_matmul_small_m.wgsl.template:inter_results : array<array<output_element_t, ...>>
The MLP/QKV kernels are selected for fused MLP / QKV subgraphs (common in Gemma 3 / SmolLM), so those paths can still overflow. Please extend the f32 accumulation to them or note explicitly why they are excluded.
| @@ -0,0 +1,116 @@ | |||
| #!/usr/bin/env python3 | |||
There was a problem hiding this comment.
Please drop this file from the PR. It introduces a new top-level scripts/ directory for a static regex check over shader source (not a runtime test), it already fails lintrunner (RUFF-FORMAT), and it will silently bit-rot the moment the shaders are refactored because it asserts on exact source substrings. If you want regression coverage, prefer a numerical unit test (large-K fp16 MatMul/MatMulNBits vs. an f32 reference) in the existing WebGPU op test suite, which exercises real runtime behavior.
|
Hi everyone, |
Could you share the links to these models? |
…els, drop verify script - matmul_nbits_mlp.wgsl.template: accumulate gate/up projection sums and workgroup inter_results in f32; apply bias and SiLU activation in f32; downcast to output_element_t only at the final store. - matmul_nbits_qkv.wgsl.template: accumulate q/k/v projection sums and workgroup inter_results in f32; downcast at the final store. - dp4a_matmul_small_m.wgsl.template: accumulate inter_results and the final reduction in f32 (SDP8AI partial products are exact integer dots; the cross-tile accumulation over K was the overflow path); downcast at the final store. dp4a_matmul_common.wgsl.template is shared with other kernels and is intentionally left unchanged. - Remove scripts/verify_f16_fix.py per review feedback (static source check, not a runtime test; also failed lintrunner RUFF-FORMAT). Weights, activations and all buffer traffic remain fp16/q4; only register/workgroup accumulators are promoted, so memory bandwidth is unchanged.
Same overflow pattern as dp4a_matmul_small_m: per-lane register accumulators (lane_outputs / lane_output1..4) summed output_element_t partial dot products across the whole K loop. Promote them to f32 and downcast to output_element_t only at the final store; SDP8AI itself and all buffer types are unchanged.
|
Tianlei Wu (@tianleiwu) Thanks for the review — addressed in 2f2a4b3 and 53c9320: Fused/native kernels extended to f32 accumulation (same pattern as the generic kernel: buffers and dequantized weights stay fp16/q4, only accumulators are promoted, downcast at the final store):
Explicitly excluded: the
|
|
Jianhui Dai (@daijh) Thanks for taking a look. Model links (these are the reports collected in #26732 / #26367):
On register pressure — a few data points for the discussion:
That said, if your measurements on Intel show a significant regression, I'm open to gating — either per-vendor, or (probably better, since overflow risk scales with reduction length) promoting only when K exceeds a threshold. Looking forward to your repro results and numbers; happy to iterate on whatever the data shows. |
WGSL has no implicit conversions: constructing vec4<output_element_t> directly from f32 scalar components is a compile error when output_element_t is f16. Build a vec4<f32> first, then use the whole-vector conversion constructor. Verified with naga (wgpu-native): the direct mixed-scalar constructor is rejected, the vector-to-vector conversion compiles for both f16 and f32 output types.
|
Follow-up on 53c9320: while compile-validating the changed shader patterns with naga (wgpu-native), I caught a type error I had introduced in the |
I tested these two models using the WebGPU build with the fp16 accumulator from #29611, but I am unable to reproduce the reported issues.
Could you verify if this still happens on a recent ONNX Runtime build and share a test page to help us reproduce the issue? |
|
Jianhui Dai (@daijh) Thanks for testing — I believe I can explain the non-repro: the fp16 Gemma weights on the Hub today are not the ones the bug was reported against. On 2025-12-02 the fp16 exports of Reproduction with the same model, pinned to the pre-mitigation revision
const gen = await pipeline('text-generation', 'onnx-community/gemma-3-270m-it-ONNX', {
dtype: 'fp16',
device: 'webgpu',
revision: '2950c41f44e1799fa9a42c2dc9fd2648fee8fb4f', // last commit before the clip mitigation
});I just ran this on stock
Same weights: correct on WASM, garbage on WebGPU. The current Self-contained test page (serve locally, open in Chrome;
|
Roberto (@RobertoReale) By the way, regarding the previous issue, were you using the legacy JSEP or the WebGPU EP? |
Legacy JSEP — transformers.js 3.8.1 bundles So I re-ran everything on the native WebGPU EP with a recent build — transformers.js 4.2.0 → 1. On the native EP, the pre-clip Gemma model fails before it can overflow
2. A minimal test that needs no model download — and why it passes on IntelSelf-contained page (below): a 119-byte embedded ONNX model, On my Intel Iris Xe: WebGPU returns exact 0 — no overflow, even though the generated That looks like evidence that f16 accumulators are fine — until you probe the same device with raw WGSL:
The Intel shader compiler evaluates straight-line f16 chains at higher precision and only rounds to f16 across loop iterations. ORT's native MatMul kernels are heavily unrolled, so on Intel they get f32 accumulation silently, from the driver — which is why the f16-accumulator builds look correct there. The JSEP kernels accumulate in And it does not transfer across vendors. The same probes via wgpu-native:
Six configurations, six behaviors. This is all WGSL-conformant — implementations are allowed to evaluate floating-point expressions with extra intermediate precision — and that's precisely the problem: whether an f16 accumulator overflows is an accident of vendor × driver × code shape. A kernel validated with f16 accumulators on Intel says nothing about NVIDIA/Qualcomm/Apple, and even on Intel it flips between D3D12 and Vulkan and between looped and unrolled codegen. The explicit f32 accumulator is the only way to make the guarantee portable — and on hardware whose compiler already promotes intermediates (Intel), it changes essentially nothing, which is consistent with the neutral perf numbers discussed earlier. (Side note for scope: the native EP's plain fp16 Test page — save as .html, serve locally, open in Chrome (
|
|
This is a great deep dive! |
|
Thanks Jianhui Dai (@daijh) — really appreciate you not blocking, and the offer to follow up. One thing I want to correct before we settle on the gating, because I think it changes the conclusion: this isn't NVIDIA-specific. The failure that #26732 was actually reported against runs on Intel too — my original garbage repro ( The reason your test passes on Intel while mine fails on the same vendor is the piece I probed in the deep-dive, and it's the key point for gating:
So the property "f16 accumulation is safe here" isn't On cost: I don't think there's a measured regression to gate around yet. The change is accumulator-only (weights/activations stay f16/q4 in VRAM, downcast at store), and the in-tree wide-tile kernel already accumulated in f32 with the comment "minimal performance impact compared to an f16 accumulator" — which #29611 is now removing. If you do have Panther Lake numbers showing f32 accumulation costs measurably on a specific fast path, I'd genuinely like to see them — that's exactly the input that should drive any narrowing. And that's the constructive version I'd propose instead of a vendor gate: default to f32 (correct everywhere), and opt down to f16 only where it's both proven safe and proven to cost. Your Last thing, tying the two PRs together: if #29611 lands the wide-tile accumulator as |
|
IMO, JSEP actually gets less maintenance. I'm really looking forward to a test page where a language model inference session causes an overflow on an Intel GPU via either D3D or Vulkan, because of a |
|
Thanks Jianhui Dai (@daijh) — completely agree on JSEP getting less maintenance than the C++ Native WebGPU EP, and that is why this PR updates both paths (the TS generators in On why As explored in the raw-WGSL probes (#issuecomment-4923584353), the reason However, on Intel via Vulkan ( Because This is why I am 100% supportive of your idea to use |
Roberto (RobertoReale)
changed the title
|
Tianlei Wu (@tianleiwu) Jiajia Qin (@qjia7) Hariharan Seshadri (@hariharans29) Jianhui Dai (@daijh) — I've corrected this PR's description, because its central claim was wrong. I had been citing Gemma 3 270M as the real-model proof of an f16 accumulator overflow, and asserting that the Clip nodes added to the Hub model on 2025-12-02 were a workaround for it. I finally checked that instead of arguing it, and it is false:
So #26732's root cause is a model-level fp16 residual overflow, and this PR does not fix it. I've removed What this PR still rests on: whisper-small q4, which is a genuine end-to-end accumulator overflow on the JSEP path (garbage on WebGPU → word-for-word WASM parity with only the accumulator promoted), plus the portability argument for strict-rounding backends. That is narrower than what I originally wrote, and Jianhui Dai (@daijh) was right to keep pushing for a real test case rather than accept the synthetic ones. I'd rather you review this on what it actually demonstrates. If you'd prefer I re-scope it further (e.g. JSEP-only, or hold the native-EP hunks until Jianhui Dai (@daijh)'s A/B packages report), say so and I'll do it. |
Every MatMulNBits test in matmul_4bits_test.cc generates its inputs from Gaussian(0, 0.25), which keeps the running partial sum along K in the single digits even at K=11008, so none of them can reach the f16 ceiling of 65504. Add a case where the partial sum crosses it while the inputs and the exact result stay inside f16 range: A=8, dequantized B=+112 over the first half of K and -112 over the second, K=8192, M=1 so the dispatch lands on matmul_nbits.wgsl.template. A lane peaks at 114688 halfway through K and cancels back to an exact 0. With an f32 accumulator the output is exactly 0; with an f16 accumulator it saturates to +Inf. The block-local sum stays in output_element_t and peaks at 28672, so the test isolates the accumulator carried across K.
|
Two things I owed this PR from Tianlei Wu (@tianleiwu)'s review. The unit test is in, It's a guard, not a detector: on Intel D3D12 dxc promotes unrolled f16 chains by itself, so it can pass there even with an f16 accumulator. That's stated in the comment above the test. Second, the description claimed zero regression, which was only true of global memory. The accumulator workgroup arrays do double for fp16 outputs. Corrected. |
…tion
Adds the WebGPU EP provider option "preferredMatmulAccumulatorPrecision"
("f16" / "f32") proposed in microsoft#29611, and makes every MatMulNBits kernel this
branch touches read it instead of hardcoding f32.
Plumbing follows the existing kvCacheQuantizationBits option: constant and
value strings in webgpu_provider_options.h, parsing in
webgpu_provider_factory.cc, config field plus accessor on the EP, and an
inline forwarder on ComputeContext so kernels can read it.
The five WGSL templates gain a compile-time `acc_f32` parameter selecting
`alias acc_element_t = f32` or `= output_element_t`, and the accumulators,
their reduction and the bias adds use that alias. Only the accumulator moves:
the block-local dot products keep the precision they have today, so the f16
setting rounds at the block boundary, which is the same shape measured in the
Gemma writeup. acc_f32 is added to every affected CacheHint so the two
variants cannot collide in the program cache.
The shipped default lives on a single line
(WebGpuExecutionProviderConfig::matmul_accumulator_precision_f32) and is f32
here, which keeps current behaviour; flipping it to f16 is a one-line change
if that is where the review lands.
Not compiled: no build tree available. Verified by running the in-tree
tools/python/wgsl_template generator over contrib_ops/webgpu, which parses all
five templates and emits both alias branches for each. The JSEP shaders are
untouched: they read env.webgpu and never see EP provider options, so this
option does not reach them.
Matches the default agreed on microsoft#29611: the option ships opt-in to f32, the native kernels keep accumulating in output_element_t unless asked otherwise.
The exclusion has been there since the first commit, on the grounds that ONNX Runtime's WebGPU backend returned "[Music]" instead of speech on q4. Re-measured against the pinned stack, which has not moved since day one (transformers 3.8.1 -> onnxruntime-web 1.22.0-dev.20250409): on whisper-small q4, three real voice messages of 26 / 33 / 39 s, WebGPU output is character-identical to WASM and 1.5x faster on an Intel Iris Xe, more on a discrete GPU. The failure does not reproduce. q8 stays excluded: its crash was never retested and no model here uses it. The crash flag and the WASM fallback are untouched. docs/webgpu-onnxruntime-fix.md rewritten. It claimed the cause was an f16 accumulator overflow that microsoft/onnxruntime#29599 would fix; both halves were wrong. whisper-small q4 contains zero fp16 tensors, so that PR is inert on this model, and the July A/B compared two builds a release cycle apart rather than the patch. What actually fixed it is not known - neither the packages nor the model changed - so the re-enable rests on the measurement, not on a theory.
Float16_LargeK_AccumulatorOverflow was written when this branch hardcoded the f32 accumulator, so it built the EP with no provider options. Now that the option ships defaulting to "f16", a plain DefaultWebGpuExecutionProvider() accumulates in output_element_t and the lane saturates at 114688 exactly as the comment describes: the test would fail on the shipped default, and fail for the right reason, which makes it useless as a guard. Build the EP with preferredMatmulAccumulatorPrecision = "f32" instead, using WebGpuExecutionProviderWithOptions the same way Float32_Large already does a few lines above. The test now guards the opt-in path, which is the one that is supposed to be exact.
The three JSEP shader generators (matmulnbits.ts, matmul-shaders.ts, 3rd-party/matmul_packed_webgpu.ts) go back to what main has. Two reasons, and the first one is the one that forces it. The provider option is a WebGPU EP option. JSEP shaders read env.webgpu and never see EP provider options, so the option cannot reach them. Leaving those hunks in would ship a PR where the native kernels are opt-in and the JSEP path is switched to f32 unconditionally for everyone, which is not a coherent thing to ask a reviewer to approve, and is not what was agreed on microsoft#29611. The second reason is that plumbing a parallel env.webgpu flag through JSEP to restore the symmetry means adding an option to a path microsoft#29716 is removing. If the JSEP path needs this before it goes away, it is a separate change with its own justification. This is the retraction of the whisper-small q4 evidence carried through to the diff: that measurement was the reason the JSEP hunks were here, and it did not hold up (whisper-small q4 has no fp16 tensor at all, so the hunks were inert on the model they were credited to).
Roberto (RobertoReale)
changed the title
|
Jiajia Qin (@qjia7) Jianhui Dai (@daijh) Tianlei Wu (@tianleiwu) Hariharan Seshadri (@hariharans29) — I've pushed the rework we settled on over on #29611, and rewritten the description to match. Now that #29611 has merged, that thread isn't the right place for this any more, so I'm moving it back here. This is now the Two things I decided rather than kept asking about. The JSEP hunks are gone — the files are back to The honest gap: I still have no ORT build tree here, so none of this has seen a compiler and the test has never run. The templates go through If the shape is wrong, say so and I'll redo it. |
There was a problem hiding this comment.
Pull request overview
Adds a new WebGPU Execution Provider option, preferredMatmulAccumulatorPrecision ("f16" default, "f32" opt-in), and plumbs it through the native WebGPU MatMulNBits family so kernels can select f16 vs f32 accumulation across the K loop. This is a targeted numerical-stability/portability knob for strict-rounding backends, while keeping the shipped default behavior unchanged.
Changes:
- Introduces and parses the new provider option, storing it on the WebGPU EP config and exposing it via
ComputeContext. - Threads an
acc_f32compile-time parameter through multiple MatMulNBits/DP4A WGSL templates and incorporates it into cache hints to avoid shader-cache collisions. - Adds a WebGPU-focused unit test that exercises accumulator overflow behavior at large K when f16 accumulation is used.
Reviewed changes
Copilot reviewed 17 out of 17 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/matmul_4bits_test.cc | Adds a regression test that requires the f32 accumulator option to avoid Inf/NaN on strict-rounding backends. |
| onnxruntime/core/providers/webgpu/webgpu_provider_options.h | Defines the new provider option key and its "f16"/"f32" values. |
| onnxruntime/core/providers/webgpu/webgpu_provider_factory.cc | Parses preferredMatmulAccumulatorPrecision into the EP config. |
| onnxruntime/core/providers/webgpu/webgpu_execution_provider.h | Adds config storage + accessor for the chosen accumulator precision. |
| onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc | Wires the parsed config into the EP instance. |
| onnxruntime/core/providers/webgpu/compute_context.h | Exposes the EP setting to kernels via ComputeContextBase. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.wgsl.template | Adds acc_f32 template param and switches the cross-K accumulator storage/reduction accordingly. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.h | Extends MatMulNBitsProgram to carry acc_f32 into shader generation. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc | Plumbs ComputeContext option into program construction and cache hint. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.wgsl.template | Adds acc_f32 param and updates accumulator storage for QKV fused decode path. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc | Passes option into template expansion and cache hints for the QKV decode program. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.wgsl.template | Adds acc_f32 param and updates accumulator storage for MLP fused decode path. |
| onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc | Passes option into template expansion and cache hints for the MLP decode program. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template | Adds acc_f32 param and changes accumulator types for DP4A matmul variants. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_small_m.wgsl.template | Adds acc_f32 param and adjusts cross-K accumulator and final store casting. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.h | Extends DP4A program classes to carry acc_f32 for shader generation. |
| onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul_nbits.cc | Plumbs ComputeContext option into DP4A program construction and cache hints. |
Suppressed comments (3)
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_qkv.cc:44
acc_f32_is used as a WGSL template parameter, but it is never initialized from theacc_f32constructor argument. That leaves the shader variant selection non-deterministic and can also desync the program-cache key from the generated shader code.
: Program{"MatMulNBitsQkvDecode"},
tile_size_(tile_size),
single_scale_weights_(single_scale_weights),
tile_size_k_vec_(tile_size_k_vec),
k_unroll_tiles_(k_unroll_tiles),
has_norm_(has_norm),
has_skip_input_(has_skip_input),
has_skip_output_(has_skip_output) {
onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits_mlp.cc:75
acc_f32_is referenced as a WGSL template parameter but is never initialized from theacc_f32constructor argument. This makes shader variant selection and caching non-deterministic.
bool acc_f32)
: Program{"MatMulNBitsMlpDecode"},
tile_size_(tile_size),
has_gate_bias_(has_gate_bias),
has_up_bias_(has_up_bias),
has_norm_input_(has_norm_input),
has_skip_input_(has_skip_input),
has_skip_output_(has_skip_output),
single_scale_weights_(single_scale_weights),
tile_size_k_vec_(tile_size_k_vec),
k_unroll_tiles_(k_unroll_tiles),
activation_kind_(activation_kind) {}
onnxruntime/contrib_ops/webgpu/quantization/dp4a_matmul.wgsl.template:508
- Bias addition is currently done after casting
lane_output*down tooutput_element_t, so enablingacc_f32doesn't actually keep the bias add in f32. Add bias inacc_element_tand downcast only once at the final store.
output.setByOffset(output_idx, vec4<output_element_t>(lane_output1) + bias_vec1);
output.setByOffset(output_idx+1, vec4<output_element_t>(lane_output2) + bias_vec2);
output.setByOffset(output_idx+2, vec4<output_element_t>(lane_output3) + bias_vec3);
output.setByOffset(output_idx+3, vec4<output_element_t>(lane_output4) + bias_vec4);
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| // Precision of the partial sums accumulated along K, selected by the | ||
| // "preferredMatmulAccumulatorPrecision" provider option. The block-local dot products stay in | ||
| // f32 either way; this only changes the accumulator carried across K, which is what can | ||
| // overflow the f16 max (~65504). | ||
| #if acc_f32 | ||
| alias acc_element_t = f32; | ||
| #else | ||
| alias acc_element_t = output_element_t; | ||
| #endif |
| output.setByOffset(output_idx, vec4<output_element_t>(vec4<acc_element_t>(lane_outputs[0], lane_outputs[1], lane_outputs[2], lane_outputs[3])) + bias_vec1); | ||
| output.setByOffset(output_idx+1, vec4<output_element_t>(vec4<acc_element_t>(lane_outputs[4], lane_outputs[5], lane_outputs[6], lane_outputs[7])) + bias_vec2); | ||
| output.setByOffset(output_idx+2, vec4<output_element_t>(vec4<acc_element_t>(lane_outputs[8], lane_outputs[9], lane_outputs[10], lane_outputs[11])) + bias_vec3); | ||
| output.setByOffset(output_idx+3, vec4<output_element_t>(vec4<acc_element_t>(lane_outputs[12], lane_outputs[13], lane_outputs[14], lane_outputs[15])) + bias_vec4); |
| @@ -70,7 +74,8 @@ class DP4AMatMulNBitsSmallMProgram final : public Program<DP4AMatMulNBitsSmallMP | |||
| has_weight_idx_(has_weight_idx), | |||
| has_weight_idx_indirect_(has_weight_idx_indirect), | |||
| single_scale_weights_(single_scale_weights), | |||
| broadcast_a_row_(broadcast_a_row) {} | |||
| broadcast_a_row_(broadcast_a_row), | |||
| acc_f32_(acc_f32) {} | |||
| MatMulNBitsProgram(uint32_t tile_size, uint32_t nbits, bool has_zero_points, bool has_bias, bool has_weight_idx, bool has_weight_idx_indirect, bool single_scale_weights, uint32_t tile_size_k_vec = 16, bool broadcast_a_row = false, bool acc_f32 = true) | ||
| : Program{"MatMulNBits"}, tile_size_(tile_size), nbits_(nbits), has_zero_points_(has_zero_points), has_bias_(has_bias), has_weight_idx_{has_weight_idx}, has_weight_idx_indirect_{has_weight_idx_indirect}, single_scale_weights_(single_scale_weights), tile_size_k_vec_(tile_size_k_vec), broadcast_a_row_(broadcast_a_row), acc_f32_(acc_f32) {} |
| @@ -106,15 +116,15 @@ fn compute_projection_sum(weight: q_b_value_t, | |||
| #if component_a == 1 | |||
| let a0 = load_a_vec4(a_offset); | |||
| let a1 = load_a_vec4(a_offset + 4); | |||
| sum += dot(a0, w0) + dot(a1, w1); | |||
| sum += dot(vec4<f32>(a0), vec4<f32>(w0)) + dot(vec4<f32>(a1), vec4<f32>(w1)); | |||
| #elif component_a == 2 | |||
| @@ -76,8 +86,8 @@ fn compute_gate_up_sums(b_global: u32, kidx: u32, idx: u32, k_offset: u32) -> ve | |||
| let gate_b_value = gate_b.getByOffset(b_global * uniforms.K_of_b + k_offset); | |||
| let up_b_value = up_b.getByOffset(b_global * uniforms.K_of_b + k_offset); | |||
|
|
|||
| var gate_sum = output_element_t(0); | |||
| var up_sum = output_element_t(0); | |||
| var gate_sum = f32(0); | |||
| var up_sum = f32(0); | |||
| var a_offset = idx * (8 / component_a) * component_b; | |||
| #if component_b == 1 | |||
| let gate_b_value_lower = vec4<output_element_t>(unpack4xU8(gate_b_value & 0x0F0F0F0Fu)) - vec4<output_element_t>(default_zero_point); | |||
| @@ -91,18 +101,18 @@ fn compute_gate_up_sums(b_global: u32, kidx: u32, idx: u32, k_offset: u32) -> ve | |||
| #if component_a == 1 | |||
| let a0 = vec4<output_element_t>(tile_A[a_offset], tile_A[a_offset + 1], tile_A[a_offset + 2], tile_A[a_offset + 3]); | |||
| let a1 = vec4<output_element_t>(tile_A[a_offset + 4], tile_A[a_offset + 5], tile_A[a_offset + 6], tile_A[a_offset + 7]); | |||
| gate_sum += dot(a0, gate_b0) + dot(a1, gate_b1); | |||
| up_sum += dot(a0, up_b0) + dot(a1, up_b1); | |||
| gate_sum += dot(vec4<f32>(a0), vec4<f32>(gate_b0)) + dot(vec4<f32>(a1), vec4<f32>(gate_b1)); | |||
| up_sum += dot(vec4<f32>(a0), vec4<f32>(up_b0)) + dot(vec4<f32>(a1), vec4<f32>(up_b1)); | |||
| #elif component_a == 2 | |||
| let a0 = vec4<output_element_t>(tile_A[a_offset], tile_A[a_offset + 1]); | |||
| let a1 = vec4<output_element_t>(tile_A[a_offset + 2], tile_A[a_offset + 3]); | |||
| gate_sum += dot(a0, gate_b0) + dot(a1, gate_b1); | |||
| up_sum += dot(a0, up_b0) + dot(a1, up_b1); | |||
| gate_sum += dot(vec4<f32>(a0), vec4<f32>(gate_b0)) + dot(vec4<f32>(a1), vec4<f32>(gate_b1)); | |||
| up_sum += dot(vec4<f32>(a0), vec4<f32>(up_b0)) + dot(vec4<f32>(a1), vec4<f32>(up_b1)); | |||
| #elif component_a == 4 | |||
| let a0 = tile_A[a_offset]; | |||
| let a1 = tile_A[a_offset + 1]; | |||
| gate_sum += dot(a0, gate_b0) + dot(a1, gate_b1); | |||
| up_sum += dot(a0, up_b0) + dot(a1, up_b1); | |||
| gate_sum += dot(vec4<f32>(a0), vec4<f32>(gate_b0)) + dot(vec4<f32>(a1), vec4<f32>(gate_b1)); | |||
| up_sum += dot(vec4<f32>(a0), vec4<f32>(up_b0)) + dot(vec4<f32>(a1), vec4<f32>(up_b1)); | |||
Jiajia Qin (qjia7)
left a comment
Jiajia Qin (qjia7)
left a comment
There was a problem hiding this comment.
Thanks for narrowing this change to an opt-in native WebGPU EP option. Exposing the accumulator choice through the provider configuration is preferable to changing precision globally. I found several correctness issues and some missing shader-path coverage.
1. Initialize acc_f32_ in the QKV and MLP decode programs
MatMulNBitsQkvDecodeProgram and MatMulNBitsMlpDecodeProgram accept bool acc_f32, but their constructors do not assign it to acc_f32_.
GenerateShaderCode() later uses that member:
WGSL_TEMPLATE_PARAMETER(acc_f32, acc_f32_)This means shader generation reads an uninitialized value. It can also cause the generated WGSL variant to disagree with the acc_f32 value included in CacheHint.
Please add acc_f32_(acc_f32) to both constructor initializer lists.
This is a pre-merge blocker.
2. Fix the undefined output_element_t alias in the QKV template
The false branch in the QKV template currently uses:
#if acc_f32
alias acc_element_t = f32;
#else
alias acc_element_t = output_element_t;
#endifHowever, the program declares outputs named q_output, k_output, and v_output. Their generated aliases should be q_output_element_t, k_output_element_t, and v_output_element_t. There is no output named output that would generate output_element_t.
Consequently, the default accumulator branch can produce invalid WGSL.
Please use an available output alias—likely q_output_element_t if all three outputs are guaranteed to have the same element type—or introduce an explicit common alias.
This is a pre-merge blocker.
3. Preserve f32 precision through the DP4A bias addition
The DP4A bias paths convert accumulated values to output_element_t before adding the bias:
vec4<output_element_t>(vec4<acc_element_t>(...)) + bias_vecWhen acc_f32=true and the output type is fp16, this rounds or saturates the accumulated result to fp16 before adding the bias. The f32 accumulator option is therefore only partially effective.
Please convert the bias to acc_element_t, perform the addition in accumulator precision, and convert to the output type only at the final store:
f32 accumulator + f32 bias -> final output conversion
The same issue appears in both DP4A output-store blocks.
4. Consider exposing this as a boolean provider option
preferredMatmulAccumulatorPrecision = "f16" | "f32" is more general than the behavior implemented here. The actual choice is binary:
- enable f32 accumulation; or
- use the normal accumulator type, which follows the output element type.
Consider naming the option:
enableMatmulNBitsF32Accumulator
with the provider key:
ep.webgpuexecutionprovider.enableMatmulNBitsF32Accumulator
Including MatmulNBits clarifies that this affects the MatMulNBits family rather than every WebGPU MatMul kernel.
Describing the false case as "f16" is also misleading when the output is fp32. The false behavior means “follow the output type,” not “force f16.”
5. Make constructor defaults consistent with the provider default
Some program constructors currently use:
bool acc_f32 = truewhile the provider configuration defaults the option to false.
Existing call sites may pass the value explicitly, but the mismatch creates a future footgun. A new call site could silently enable f32 accumulation without consulting the provider configuration.
Please either:
- remove the default argument and require every caller to provide the option; or
- change the default to
false.
Removing the default is safer because it makes incomplete option plumbing a compile-time error.
6. Clarify the precision of block-local dot products
The QKV and MLP templates make their block-local helpers return f32 and explicitly convert operands to vec4<f32>, including when acc_f32=false.
That appears broader than the PR description, which states that block-local dot products retain their previous precision and only the accumulator carried across K changes.
If always-f32 block-local dot products are intentional, please update the PR description and shader comments and provide performance and numerical validation for the default path. Otherwise, these helpers should preserve output-type computation when acc_f32=false.
7. Add focused tests for every accumulator-dependent shader path
Please add small WebGPU tests that instantiate and execute both acc_f32=false and acc_f32=true for every affected implementation.
The coverage should include:
- generic MatMulNBits
- fused QKV
- fused MLP
- DP4A regular path
- DP4A small-M path
- wide-tile path
- subgroup-matrix path
- bias and no-bias variants where applicable
- default/output-type accumulation and f32 accumulation
Small deterministic inputs should be sufficient. Each case should verify:
- session and pipeline creation succeeds;
- the selected WGSL variant compiles;
- dispatch completes;
- output is finite and numerically reasonable;
- changing the provider option selects the correct cached shader variant.
This coverage would detect:
- undefined aliases in an individual WGSL branch;
- missing
acc_f32_initialization; - incomplete provider-option propagation;
- an affected implementation that ignores the option;
- cache hints that do not distinguish accumulator variants;
- shader compilation failures limited to wide-tile or subgroup-matrix paths.
The existing large-K cancellation test can remain as the focused numerical regression demonstrating why f32 accumulation is needed. The smaller tests would complement it by ensuring every generated shader variant compiles and executes.
Verdict
I would request changes before approval.
The primary blockers are:
- uninitialized
acc_f32_in the QKV and MLP programs; - undefined
output_element_tin the QKV default branch; - premature downcasting in the DP4A bias path;
- missing compilation and runtime coverage across all affected implementations, including wide-tile and subgroup-matrix paths.
The provider-option design, constructor defaults, and block-local dot-product precision should also be resolved or documented before merge.
| for (var i = 0u; i < 16u; i++) | ||
| { | ||
| lane_outputs[i] += SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i)); | ||
| lane_outputs[i] += acc_element_t(SDP8AI(own_a0, subgroupShuffle(own_b0, i), own_a1, subgroupShuffle(own_b1, i), subgroupShuffle(own_scale_b, i) * own_scale_a, subgroupShuffle(zero, i))); |
There was a problem hiding this comment.
Is it better if we change SDP8AI definition to return acc_element_t type instead of output_element_t type so that the whole code changes are small?
| } | ||
|
|
||
| fn compute_gate_up_sums(b_global: u32, kidx: u32, idx: u32, k_offset: u32) -> vec2<output_element_t> { | ||
| fn compute_gate_up_sums(b_global: u32, kidx: u32, idx: u32, k_offset: u32) -> vec2<f32> { |
There was a problem hiding this comment.
Why f32 not acc_element_t?
| idx: u32) -> q_output_element_t { | ||
| var sum = q_output_element_t(0); | ||
| idx: u32) -> f32 { | ||
| var sum = f32(0); |
…in QKV template, bias downcast ordering in dp4a, SDP8AI return type, constructor defaults
matmul_nbits_qkv.wgsl.template has no single "output" binding (q_output/k_output/v_output
instead), so acc_element_t's f16 branch aliased to an undefined output_element_t and would
fail shader compilation on the shipped default ("f16"). Aliases to q_output_element_t instead,
since all three outputs share one dtype (single T1 constraint).
dp4a_matmul.wgsl.template added bias after downcasting the accumulator to output_element_t,
which could saturate to f16 range before the bias add even when acc_f32 selected an f32
accumulator, defeating the point of the option. Bias is now added in acc_element_t and the
result downcast to output_element_t only at the store, in both the is_qualcomm and
non-qualcomm paths.
SDP8AI in dp4a_matmul_common.wgsl.template downcast its result to output_element_t before
returning, so for n_bits==8 (where mul_precision is f32) an individual dot-product term could
already saturate before ever reaching the accumulator, regardless of acc_f32. It now returns
acc_element_t directly; every call site already wraps the call in acc_element_t(...), so this
needed no caller changes.
MatMulNBitsProgram and DP4AMatMulNBitsSmallMProgram defaulted their acc_f32 constructor
parameter to true, inconsistent with the shipped provider-option default of "f16". Current call
sites pass the value explicitly, but the default is corrected to match.
Validated with tools/python/wgsl_gen.py (exit 0, both alias branches emitted for all five
templates) and with naga/wgpu-native on the four changed patterns in both acc_f32 states, plus
a negative control confirming the harness rejects an undefined identifier.
|
Pushed (
On the bias ordering in Jiajia Qin (@qjia7) on On Still not compiled — no build tree here. Checked with |
Jiajia Qin (qjia7)
left a comment
Jiajia Qin (qjia7)
left a comment
There was a problem hiding this comment.
Review frame
- Exact head reviewed:
b5e48ef776313b6a1e96c62a42c0aee460b93d7d. - Feature validity: valid. A session-scoped opt-in to f32 accumulation is a reasonable portability fallback while retaining the current f16 default.
- Risk/scope: deep. This adds a provider option and changes precision, shader generation, dispatch variants, program-cache identity, and public WebGPU configuration.
- Direction gate: the provider-option direction is appropriate, but the current implementation does not yet provide a consistent or publicly reachable contract.
qjia7 comment status
| Item | Status at the latest head |
|---|---|
Initialize acc_f32_ in QKV and MLP |
Not resolved. Both members remain uninitialized. |
QKV undefined output_element_t |
Resolved. It now uses q_output_element_t, and Q/K/V share T1. |
| DP4A bias addition precision | Resolved. Bias is converted to acc_element_t and the result is downcast only at the store. |
| Boolean, MatMulNBits-scoped option | Not resolved. The general string option remains. |
| Constructor defaults | Resolved as requested: the remaining defaults now match the provider default (false); the regular DP4A constructor requires the argument. |
| Block-local MLP/QKV dot-product precision | Not resolved. The default path still computes these helpers explicitly in f32. |
| Coverage of accumulator-dependent paths | Not resolved. The only new runtime case covers the generic f32 path. |
GitHub still reports all three qjia7 inline threads as unresolved. The SDP8AI request is fixed in code; the MLP and QKV f32 threads remain substantively open.
Confirmed findings
C1: Initialize the QKV and MLP acc_f32_ members
MatMulNBitsQkvDecodeProgram accepts acc_f32, but its initializer list ends at has_skip_output_(has_skip_output) (matmul_nbits_qkv.cc:44). MatMulNBitsMlpDecodeProgram similarly ends at activation_kind_(activation_kind) (matmul_nbits_mlp.cc:75). Both later pass the uninitialized member to WGSL_TEMPLATE_PARAMETER, while their CacheHint uses the correctly initialized local value from the caller.
This is undefined behavior and can compile an f16 shader under an f32 cache key, or the reverse. Please add acc_f32_(acc_f32) to both initializer lists. This remains a merge blocker.
C2: Make the f32 option effective for optimized MatMulNBits dispatches
ApplyMatMulNBits selects subgroup-matrix and wide-tile programs before it reads MatmulAccumulatorPrecisionF32() for the generic program (matmul_nbits.cc:228-326). The wide-tile shader still declares results : array<output_element_t, ...>, and subgroup-matrix configurations can use f16 results. Consequently, common M >= 4 shapes can ignore an explicit "f32" request.
Please either make those paths honor the option, or skip/fall back from a path that cannot honor it when f32 is requested. If the intended contract is deliberately narrower, rename the option to describe that scope and document exactly which paths may ignore it. Simply excluding these paths leaves the current provider-option contract misleading.
S1: Expose and forward the option through the supported WebGPU APIs
The native parser accepts the new config key, but WebGpuExecutionProviderOption in js/common/lib/inference-session.ts has no corresponding property. The browser bridge in js/web/lib/wasm/session-options.ts only forwards explicitly handled fields, so an extra property is ignored; the Node bridge in js/node/src/session_options_helper.cc rejects it as unrecognized. The new C++ test constructs the EP internally and therefore does not cover either public path.
Please add the typed option plus validation/forwarding in the browser and Node bridges, with a public-path test. Otherwise the primary onnxruntime-web users cannot enable the feature.
S2: Preserve the f16 default in the fused block-local helpers
Relative to main, compute_gate_up_sums and compute_projection_sum were changed from output-element computation to explicit f32 returns and vec4<f32> operands. This happens even when acc_f32=false, contradicting the stated unchanged default and potentially changing default numerics and performance.
Roberto's explanation that these sums are block-local and bounded explains why overflow is less concerning there, but it does not justify changing their default precision. Please use acc_element_t consistently for the helper return/operands, so the option selects f32 and the default retains output-type computation.
Roberto's latest reply
The reply is partially valid:
- QKV alias fix: valid.
- DP4A bias-order fix: valid in both Qualcomm and non-Qualcomm store paths.
SDP8AIreturn-type fix: valid; it now preserves f32 through the scaled term when requested.- Constructor-default fix: valid.
- Always-f32 MLP/QKV helper explanation: incomplete for compatibility; it does not address the default behavior change.
The reply does not mention or fix the uninitialized QKV/MLP members, public option reachability, optimized-path bypass, or remaining coverage gap.
Test coverage
Float16_LargeK_AccumulatorOverflow deterministically selects the generic MatMulNBits path and requests f32, which is useful numerical coverage. It does not exercise the default/f16 variant, fused MLP, fused QKV, either DP4A path, public option forwarding, or cache separation. Existing MLP/QKV fusion tests and accuracy-level-4 MatMulNBits tests should be parameterized/reused with both option values instead of adding duplicate WebGPU-only test bodies.
For wide-tile/subgroup-matrix, first settle C2: either test that f32 is honored/falls back, or narrow the option's contract. wgsl_gen.py succeeds and git diff --check is clean, but the PR has no compile/test CI beyond CLA, and the author states that the C++ and runtime test have not been built or run.
Verdict
Request changes. The feature and provider-option direction are valid, but C1, C2, S1, and S2 block merge. qjia7's alias, DP4A bias, SDP8AI, and constructor-default points are fixed; initialization, helper precision, option scope, and coverage remain open. The option naming/boolean shape should be resolved while the API is still new. Formatting is cleanup only.
qjia7's review of b5e48ef: - MatMulNBitsQkvDecodeProgram and MatMulNBitsMlpDecodeProgram took acc_f32 but never assigned it, so shader generation read an uninitialized member while the cache hint used the caller's value. Both now initialize it. - The wide-tile program is selected before the generic one and ignored the option entirely, so M >= 4 prefill shapes silently kept an f16 accumulator. It now takes acc_f32, aliases acc_element_t, adds the bias in accumulator precision and downcasts only at the store, with acc_f32 in its cache hint. - The subgroup-matrix path cannot honour an f32 request: every fp16 entry in supported_subgroup_matrix_configs has resultComponentType F16 and the kernels declare subgroup_matrix_result<f16, ...>. It now declines itself when f32 is requested for an fp16 output instead of ignoring the option. - compute_gate_up_sums and compute_projection_sum were changed to unconditional f32 in an earlier commit, which altered the default path relative to main. They follow acc_element_t again, so acc_f32=false is identical to main. - Expose preferredMatmulAccumulatorPrecision through the JS APIs: typed option in js/common, validation and forwarding in js/web, accepted key in js/node. - Add Float16_AccumulatorPrecisionOption_AllPaths, which runs the generic, wide-tile and dp4a dispatches with the option in both states.
|
Pushed C1 was undefined behaviour and you're right that the cache hint made it worse: both decode programs took C2: the wide-tile program is selected before the generic one, so the shapes it takes were ignoring the option outright. It now carries For subgroup matrix I took the other branch of your suggestion and declined the path. Every fp16 entry in S2: correct, and I should have fixed it when you first raised it instead of explaining why it was safe. The helpers follow S1: typed option in On the name, I'd keep New test Still nothing compiled here: |
microsoft#31703 rewrote the MatMulNBits wide-tile compute stage: the A tile is stored as pairs of vecs and the reduction now has three variants picked from subgroup_min_size (full-tile subgroupShuffle band, chunked bands, or a direct read). Re-apply the acc_element_t accumulator on top of that shape. - results is array<acc_element_t, kTileM>, and all three reduction branches widen the dot products; the subgroupShuffle itself stays in the workgroup-memory element type, so the broadcast payload does not grow. - MatMulNBitsWideTileProgram carries both subgroup_min_size and acc_f32, and the cache hint carries both, so the variants cannot collide. - With acc_f32 = false the casts are identities and the generated shader is main's, checked by diffing wgsl_gen output against 1bc68c1.
|
Now that #31703 has landed I've merged The wide-tile hunk needed redoing rather than re-applying. Jianhui Dai (@daijh)'s reduction has three variants selected from With the default the casts are identities, and I checked that instead of assuming it: generating both templates with One thing I had wrong: I expected shuffling the already-widened Worth flagging for the cost side: on fp16 |
Summary
Adds the WebGPU EP provider option
preferredMatmulAccumulatorPrecision("f16"default,"f32"opt-in) proposed by Jiajia Qin (@qjia7) on #29611, and routes every MatMulNBits kernel this PR touches through it instead of hardcoding the accumulator type. Weights and activations are untouched: only the register/workgroup accumulators change, so global memory traffic is identical either way.Default behaviour on
mainis unchanged by this PR.> This description was rewritten on 2026-08-06 (scope refreshed 2026-08-12). The PR started life as an unconditional f16→f32 accumulator promotion across both the JSEP generators and the native EP, justified by Gemma 3 270M and later by whisper-small q4. Both of those attributions were wrong and I retracted them (Gemma, whisper). What is left is below, and it is deliberately narrower.
What this is now
An opt-in, not a default change. The shipped default is
"f16", which is what #29611 landed. Setting"f32"promotes the accumulators carried along K in the six MatMulNBits kernels below. The block-local dot products keep the precision they have today either way, so the f16 setting rounds at the block boundary — the same shape that was measured in the Gemma writeup.Session-scoped, not a model attribute. It is an EP provider option, so nobody has to re-export a model to use it, and it also covers the fp16
MatMulpath, which has noaccuracy_levelto hang off.What I am not claiming
I have no real model that overflows one of these accumulators. Not Gemma 3 270M (the f16 partial sum inside all 126 MatMulNBits nodes peaks at ~324; the overflow there is the model-level fp16 residual stream). Not whisper-small q4 (that model has no fp16 tensor at all, so these hunks are inert on it; its real saturation is a
Powinside an unfused decomposed LayerNorm, filed separately as #31626). A/B runs on Qwen2.5-0.5B q4f16, gemma-3-270m and Phi-4-mini across two GPUs showed zero Inf and zero NaN either way, only rounding drift.What holds up is portability, and only that. WGSL permits extra intermediate precision: Intel's D3D12 compiler (dxc) promotes unrolled
acc +=chains to f32 on its own but rounds to f16 between loop iterations, while NVIDIA/Vulkan rounds strictly (probes). So "f16 accumulation is safe" is a property of vendor × backend × code shape, not of a vendor — which is why a vendor-string gate cannot express it, and why a kernel carrying an accumulator across a K loop has no headroom guarantee on strict-rounding configurations. That argues for a way back, which is what this option is. It does not argue for changing the default, and I am not asking for that.Scope
Native WebGPU EP (
onnxruntime/contrib_ops/webgpu/quantization/):matmul_nbits,matmul_nbits_wide_tile,matmul_nbits_mlp,matmul_nbits_qkv,dp4a_matmul,dp4a_matmul_small_m. Each template gains a compile-timeacc_f32parameter selectingalias acc_element_t = f32or the output element type; the accumulators, their reduction and the bias adds use that alias, and the downcast happens only at the store.acc_f32is added to every affectedCacheHintso the two variants cannot collide in the program cache.SDP8AIindp4a_matmul_common.wgsl.templatereturnsacc_element_t— it is an exact integerdot4I8Packeddot product, but it must not downcast before the value reaches the accumulator.subgroup_matrix_matmul_nbits_*cannot honour an f32 request: every fp16 entry insupported_subgroup_matrix_configshasresultComponentType == F16and the kernels declaresubgroup_matrix_result<f16, ...>. Rather than ignore the option there,CanApplySubgroupMatrixMatMulNBitsdeclines when"f32"is requested for an fp16 output, so the dispatch falls through to a kernel that does honour it. fp32 outputs already accumulate in f32 and are unaffected.kvCacheQuantizationBitsoption: constant and value strings inwebgpu_provider_options.h, parsing inwebgpu_provider_factory.cc, a config field plus accessor on the EP, and an inline forwarder onComputeContext.JS bindings: the option is typed in
js/common/lib/inference-session.ts, validated and forwarded to the EP injs/web/lib/wasm/session-options.ts, and accepted by the node bridge injs/node/src/session_options_helper.cc, soonnxruntime-webandonnxruntime-nodeusers can actually set it.The JSEP hunks have been removed (
matmulnbits.ts,matmul-shaders.ts,3rd-party/matmul_packed_webgpu.tsare back tomain). JSEP shaders readenv.webgpuand never see EP provider options, so the option cannot reach them; leaving them in would have meant shipping a PR where the native side is opt-in and the JSEP side is switched unconditionally for everyone. Adding a parallelenv.webgpuflag to restore the symmetry means putting an option into a path #29716 is removing, so if JSEP needs this before it goes away it should be its own change.Testing
MatMulNBits.Float16_LargeK_AccumulatorOverflowinmatmul_4bits_test.cc, built with the option set to"f32". A=8, dequantized B at +112 over the first half of K and -112 over the second, K=8192, M=1 so the dispatch lands onmatmul_nbits.wgsl.template. A lane peaks at 114688 halfway along K and cancels back to an exact 0, so f32 returns exactly 0 while f16 saturates to +Inf. Every other MatMulNBits test in that file draws its inputs from Gaussian(0, 0.25) and cannot reach 65504 at any K, which is why this class of bug had no coverage. It is a regression guard, not a universal detector: on Intel D3D12 it can pass with an f16 accumulator, and the comment above it says so.MatMulNBits.Float16_AccumulatorPrecisionOption_AllPathsruns the generic, wide-tile and dp4a dispatches with the option in both states on ordinary Gaussian inputs, so a variant that fails to compile, an unpropagated flag or a cache hint that cannot separate the two variants shows up as a test failure. The fused MLP and QKV decode kernels are reached through graph fusion rather than a MatMulNBits node and are not covered by it.tools/python/wgsl_gen.pyovercontrib_ops/webgpuparses all six templates and emits both alias branches for each. The changed shader shapes are compile-checked with naga in both accumulator states, with negative controls confirming the harness rejects an f32 accumulator stored to an f16 output without conversion.Performance
f32 accumulators cost registers and shared memory, not bandwidth — the accumulator workgroup arrays double for fp16 outputs, which can lower occupancy on SLM-bound dispatch configurations. Jianhui Dai (@daijh) measured that cost on Intel, and #29611 is the conclusion drawn from it. Since the default here stays
"f16", nobody pays for it unless they ask.Related: #29611, #29716, #31626, #26732, #26367, #31703 (both touch
matmul_nbits_wide_tile.wgsl.template)