fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to prevent overflow (#26732)#29599
fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to prevent overflow (#26732)#29599RobertoReale wants to merge 4 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 |
| @@ -0,0 +1,116 @@ | |||
| #!/usr/bin/env python3 | |||
|
@RobertoReale, could you run |
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.
|
@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
|
|
@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? |
|
@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;
|
@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 @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 @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 |
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
f16tof32in 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
f16for 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 andBlockwiseMatMulNBits32) accumulate inf32(workgroup_shared/inter_results), with the result downcast to the output type only at the final write.matmul-shaders.ts(naive MatMul):valuesaccumulator promoted tof32.3rd-party/matmul_packed_webgpu.ts(tiled MatMul, vec4 + scalar variants):accpromoted tof32. 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 whenenable f16;is active, even if the accumulator variable is declared asf32.Native WebGPU EP (
onnxruntime/contrib_ops/webgpu/quantization/):matmul_nbits.wgsl.template:inter_resultsand the final reduction accumulate inf32along K; downcast at the output write. This aligns the kernel withmatmul_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 workgroupinter_resultsaccumulate inf32; bias and SiLU are applied inf32; 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 inf32.SDP8AIindp4a_matmul_common.wgsl.templateis unchanged - it is an exact integerdot4I8Packeddot product; the overflow path was the accumulation across K tiles.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 usef32.Testing
output_element_t = f16(withenable f16;) andf32variants, including the vector conversion constructors at the final stores.cd js/web && npx tsc --noEmit: clean; prettier: clean.onnxruntime-webbuilt from this branch - word-for-word transcript parity with the WASM EP, ~2x faster (methodology and results in this comment).matmul_4bits_test.cc) was offered in review; awaiting maintainer preference on in-PR vs follow-up.Performance note
f32accumulators 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