fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to prevent overflow (#26732)

[RobertoReale]

Summary

Fixes the numerical overflow that causes fp16 and q4f16 models (Gemma 3 270M, Whisper Q4, SmolLM2, kokoro-js) to produce garbage output on WebGPU, by promoting the MatMul / MatMulNBits accumulators from f16 to f32 in both the JSEP shader generators and the native WebGPU EP WGSL templates.

Root cause

The WGSL shaders generated for MatMul and MatMulNBits accumulate the dot-product inner loop in the output element type, which is f16 for fp16/q4f16 models. With hidden dimension D >= 2048, sums of fp16 products routinely exceed 65504 (f16 max), saturating to +Inf and then propagating NaN through subsequent LayerNorm / Softmax operations.

CPU/WASM is unaffected because MLAS accumulates internally in f32 even with f16 inputs. Note: issue #26732 was auto-closed by the stale bot without a runtime fix (the reported symptom was mitigated by re-exporting the affected models on the HF Hub); the underlying runtime issue is still present and was still being hit (e.g. kokoro-js, see the last comments on the issue). Maintainers may want to reopen it.

Fix details

JSEP (js/web/lib/wasm/jsep/webgpu/ops/):

• matmulnbits.ts: both kernels (default and BlockwiseMatMulNBits32) accumulate in f32 (workgroup_shared / inter_results), with the result downcast to the output type only at the final 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. The shared-memory tiles (mm_Asub/mm_Bsub) stay in the input dtype, so shared-memory usage and bandwidth are unchanged.

Explicit vec4<f32>() casts on each multiply operand are required - Dawn/D3D12 re-demotes temporaries to f16 when enable f16; is active, even if the accumulator variable is declared as f32.

Native WebGPU EP (onnxruntime/contrib_ops/webgpu/quantization/):

• matmul_nbits.wgsl.template: inter_results and the final reduction accumulate in f32 along K; downcast at the output write. This aligns the kernel with matmul_nbits_wide_tile.wgsl.template, which already accumulates in f32.
• matmul_nbits_mlp.wgsl.template / matmul_nbits_qkv.wgsl.template (fused MLP / QKV, added per review feedback): gate/up and q/k/v projection sums plus the workgroup inter_results accumulate in f32; bias and SiLU are applied in f32; downcast only at the final store.
• dp4a_matmul_small_m.wgsl.template / dp4a_matmul.wgsl.template: the cross-tile accumulators (inter_results, lane_outputs / lane_output1..4) accumulate in f32. SDP8AI in dp4a_matmul_common.wgsl.template is unchanged - it is an exact integer dot4I8Packed dot product; the overflow path was the accumulation across K tiles.
• Intentionally excluded: subgroup_matrix_matmul_nbits_* (accumulator precision is tied to the hardware cooperative-matrix configuration negotiated on the C++ side - see review thread).

Zero bandwidth regression: weights and activations remain in their original dtype (f16 / q4) in global GPU memory. Only register/workgroup-level accumulators use f32.

Testing

• WGSL validation: every modified shader pattern was compile-checked with naga (wgpu-native) in both output_element_t = f16 (with enable f16;) and f32 variants, including the vector conversion constructors at the final stores.
• cd js/web && npx tsc --noEmit: clean; prettier: clean.
• End-to-end: whisper-small q4 through onnxruntime-web built from this branch - word-for-word transcript parity with the WASM EP, ~2x faster (methodology and results in this comment).
• Manual browser test on Chrome 121+ via transformers.js: Gemma 3 270M FP16 and SmolLM2 360M Q4F16 on WebGPU.
• A numerical regression test (large-K fp16 case vs f32 reference in matmul_4bits_test.cc) was offered in review; awaiting maintainer preference on in-PR vs follow-up.

Performance note

f32 accumulators have negligible ALU overhead on modern GPUs. Memory bandwidth - the actual bottleneck in LLM inference - is unaffected. The codebase already uses f32 accumulation in the wide-tile MatMulNBits kernel, and the dp4a 8-bit path already promotes to f32 before scaling (mul_precision = f32) for the same overflow reason.

Fixes #26732
Related: #26367

[RobertoReale]

@microsoft-github-policy-service agree

[RobertoReale]

Additional real-world validation

I validated this fix end-to-end against a production consumer of onnxruntime-web (my Chrome extension, Voice Message Transcriber, which runs Whisper via @huggingface/transformers in the browser and currently force-disables WebGPU for q4 models specifically because of #26732).

Method: since this fix only touches the JSEP TypeScript path (js/web/lib/wasm/jsep/...), no wasm binary rebuild was needed. I sparse-cloned js/web + js/common, checked out this PR's branch, built onnxruntime-web from source (pulling the prebuilt wasm artifacts unchanged), and swapped the resulting dist/ into the extension's node_modules/onnxruntime-web in place of the published 1.22.0-dev build. I then temporarily removed the extension's q4/q8 WebGPU exclusion and ran the same voice message through both WASM and WebGPU on Windows, clearing the transcription cache between runs so each device actually re-ran inference.

Result (whisper-small, q4, ~36s of Italian audio):

Device	Time	Output
WASM (baseline)	26.7s	correct, full transcript
WebGPU (this PR's onnxruntime-web build)	13.2s	word-for-word identical text

No [Music]/hallucination, no NaN garbage, no crash — console confirmed device: webgpu used throughout with no fallback to WASM. ~2x speedup with zero accuracy regression, consistent with the "zero bandwidth regression" claim in the PR description (only register/workgroup accumulators promoted to f32).

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

[tianleiwu]

@RobertoReale, could you run lintrunner -a.
See

onnxruntime/docs/Coding_Conventions_and_Standards.md

Line 156 in a06675e

This project uses [lintrunner](https://github.com/suo/lintrunner) for linting. It provides a consistent linting experience locally and in CI. You can install the dependencies and initialize with

[tianleiwu]

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) and gate_inter_results/up_inter_results : array<array<output_element_t, ...>>
• matmul_nbits_qkv.wgsl.template — sum = q_output_element_t(0) and q/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.

[tianleiwu]

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) and gate_inter_results/up_inter_results : array<array<output_element_t, ...>>
• matmul_nbits_qkv.wgsl.template: sum = q_output_element_t(0) and q/k/v_inter_results
• dp4a_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.

[tianleiwu]

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.

[daijh]

Hi everyone,
Forcing f32 accumulators for f16 data types can cause register pressure on Intel GPUs, leading to significant performance degradation.
Could we reconsider this approach, or perhaps gate it by vendor device?
In the meantime, I will try to reproduce the reported overflow on my end.

[daijh]

Fixes the numerical overflow that causes fp16 and q4f16 models (Gemma 3 270M, Whisper Q4, SmolLM2, kokoro-js) to produce garbage output on WebGPU, by promoting the MatMul / MatMulNBits accumulators from f16 to f32 in both the JSEP shader generators and the native WebGPU EP WGSL template.

Could you share the links to these models?

[RobertoReale]

@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):

• matmul_nbits_mlp.wgsl.template — gate_sum/up_sum, gate/up_inter_results, bias and SiLU activation now computed in f32.
• matmul_nbits_qkv.wgsl.template — compute_projection_sum and q/k/v_inter_results in f32.
• dp4a_matmul_small_m.wgsl.template — inter_results and the final reduction in f32. SDP8AI itself is exact (integer dot4I8Packed), so the overflow path was the cross-tile accumulation over K; dp4a_matmul_common.wgsl.template is intentionally unchanged since it is shared.
• While auditing for the same pattern I found that the generic dp4a_matmul.wgsl.template also accumulates lane_outputs/lane_output1..4 in output_element_t across the whole K loop — fixed the same way in 53c9320.

Explicitly excluded: the subgroup_matrix_matmul_nbits_* kernels declare subgroup_matrix_result<f16, ...>; accumulator precision there is tied to the cooperative-matrix configuration negotiated on the C++ side, so moving to an f16×f16→f32 config is a larger change (and MMA units typically accumulate internally at higher precision already). I'd prefer to handle that in a follow-up if overflow is ever reported on those paths.

scripts/verify_f16_fix.py dropped in 2f2a4b3, which also resolves the lintrunner RUFF-FORMAT finding (the remaining diff is TS — prettier-clean — plus WGSL templates; if CI lint still flags anything I'll fix it). On regression coverage: agreed a numerical test is the right long-term answer — e.g. a large-K fp16 MatMulNBits case with inputs arranged so partial sums exceed 65504 while the true result stays representable, compared against an f32 reference. Happy to add that to onnxruntime/test/contrib_ops/matmul_4bits_test.cc in this PR or as a follow-up — whichever you prefer.

[RobertoReale]

@daijh Thanks for taking a look. Model links (these are the reports collected in #26732 / #26367):

• Gemma 3 270M fp16 — https://huggingface.co/onnx-community/gemma-3-270m-it-ONNX, transformers.js with dtype: "fp16", device: "webgpu" → <unused56> garbage tokens.
• SmolLM2 360M q4f16 — https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct (ONNX weights in the onnx-community mirror), dtype: "q4f16".
• Whisper q4 — https://huggingface.co/onnx-community/whisper-small and https://huggingface.co/onnx-community/whisper-large-v3-turbo, dtype: "q4" → hallucination loops / [Music]. This is the one I validated end-to-end above (word-for-word parity with the WASM EP after the fix, ~2x faster).
• kokoro-js (Kokoro-82M TTS) — https://huggingface.co/onnx-community/Kokoro-82M-v1.0-ONNX, fp16 → corrupted audio.

On register pressure — a few data points for the discussion:

1. The promotion is accumulator-only: buffer traffic, shared-memory tile loads and dequantized weights all stay fp16/q4. Per thread it is a handful of registers (e.g. 16×f32 instead of 16×f16 lane outputs in the generic dp4a kernel); the workgroup inter_results arrays grow by 2 bytes/element.
2. There is precedent in-tree: matmul_nbits_wide_tile.wgsl.template already accumulates in f32, and the dp4a 8-bit path already promotes to f32 before scaling for exactly this overflow reason (mul_precision = f32 in dp4a_matmul_common.wgsl.template).
3. fp32 models are unaffected — with output_element_t == f32 the added casts are identity.
4. The failure mode this fixes is not a minor accuracy loss but Inf/NaN propagation that makes fp16/q4f16 output unusable on WebGPU regardless of speed.

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.

[RobertoReale]

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 is_qualcomm store path of dp4a_matmul.wgsl.template - WGSL has no implicit conversions, so building vec4<output_element_t> directly from the now-f32 lane scalars is rejected when output_element_t is f16. Fixed in 2cf1e37 by constructing a vec4<f32> first and using the whole-vector conversion constructor. All modified patterns (MLP, QKV, dp4a small-M, dp4a generic, plus the final-store conversions) now compile clean under naga in both f16 and f32 variants; PR description updated to reflect the current scope.