microsoft / microsoft/onnxruntime
[Web] JSEP DequantizeLinear: wrong results for uint8 per-tensor quantization (bad vec4 indexing + unsigned subtraction wraparound)
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
`DequantizeLinear` on the JSEP WebGPU backend (`js/web/lib/wasm/jsep/webgpu/ops/quantize-linear.ts`) produces incorrect results for per-tensor-quantized `uint8`/`int8` input in two independent ways:
1. **Wrong input read for tensors larger than 4 elements.** The per-tensor fast path vectorizes output into `vec4`s, but the packed-input read still divides the *already-vectorized* `global_idx` by 4 a second time, so it re-reads the same 32-bit input word for every group of 4 output vec4s instead of advancing through the buffer.
2. **Unsigned integer wraparound for `uint8` with a nonzero zero-point.** The shader subtracts `zero_point_value` from `x_value` while both are still in the packed integer type (`vec4` for `uint8`) *before* casting to float. Since WGSL `u32` subtraction wraps instead of going negative, any element where `x < zero_point` produces a huge positive value (~2^32 scale) instead of the correct negative float.
Both bugs are specific to this JSEP file. The CPU kernel (`onnxruntime/core/providers/cpu/quantization/quantize_linear.cc`) and the newer native WebGPU EP (`onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc`) both compute in a signed/float type before subtracting, and the native EP has an explicit branch for packed vs. vectorized indexing that JSEP is missing.
### To reproduce
Two minimal cases, each isolating one bug (opset 13, `DequantizeLinear(x, x_scale, x_zero_point)`, no `axis`/`block_size` set — default per-tensor quantization).
### Case A — indexing bug (zero_point = 0, so no wraparound can mask/confound it)
```
x = uint8 [10, 20, 30, 40, 50, 60, 70, 80] (dims: [8])
x_scale = float32 1.0 (scalar)
x_zero_point = uint8 0 (scalar)
Expected y = [10, 20, 30, 40, 50, 60, 70, 80]
Actual y = [10, 20, 30, 40, 10, 20, 30, 40] (second vec4 re-reads the first input word)
```
### Case B — unsigned subtraction bug (exactly 4 elements, so Case A's indexing bug can't fire)
```
x = uint8 [0, 64, 128, 255] (dims: [4])
x_scale = float32 0.5 (scalar)
x_zero_point = uint8 128 (scalar)
Expected y = [-64, -32, 0, 63.5]
Actual y ≈ [2147483584, 2147483616, 0, 63.5] (~2^31, from u32 wraparound before the float cast)
```
A JS repro against `onnxruntime-web` with `executionProviders: ['webgpu']` (JSEP build) reproduces both; the same graph run on `wasm` gives the expected output in both cases.
## Expected behavior
Output should match `y = (x - zero_point) * scale` computed with sign-correct promotion, matching the WASM/CPU backend and the ONNX spec, for tensors of any size.
## Root cause (code pointers, current `main`)
`js/web/lib/wasm/jsep/webgpu/ops/quantize-linear.ts`:
- Lines ~99–102: per-tensor quantization with packed `int8`/`uint8` input sets `components = 4` (vectorized output) whenever `getMaxComponents(outputSize) === 4`, which is true for the overwhelming majority of tensor sizes.
- Line ~139: `let input = ${input.getByOffset('global_idx / 4')};` — this divides by 4 to step from a packed byte index to a packed word index, but `global_idx` here is already a **vec4-of-output index** (dispatch is over `outputSize / components`), not a per-element index. It needs to be `global_idx` directly (one word already holds exactly 4 packed bytes = 1 output vec4), not `global_idx / 4`.
- Line ~210: `${output.type.value}(x_value - zero_point_value) * scale_value` — the subtraction happens before the cast to `output.type.value` (float). For `uint8`, `x_value`/`zero_point_value` are `vec4`/`u32` from `unpack4xU8`, so this is unsigned subtraction that wraps instead of producing a negative intermediate.
For comparison, the native WebGPU EP gets both right:
- `onnxruntime/core/providers/webgpu/quantization/quantize_linear.cc` line ~41–51: branches on `output.NumComponents()` and uses `x.GetByOffset("global_idx / 4")` only for the *unvectorized* (`components == 1`) case; the vectorized case correctly uses `x.GetByOffset("global_idx")`.
- Same file, line ~150: `(output_value_t(x_value) - scale_value_t(zero_point_value)) * scale_value` — casts to the float value type before subtracting.
## Suggested fix
In `quantize-linear.ts`:
1. When `components === 4` (the vectorized per-tensor path) and input is packed, read the input word with `input.getByOffset('global_idx')` instead of `global_idx / 4`.
2. Cast `x_value` and `zero_point_value` to the output float type before subtracting, e.g. `${output.type.value}(x_value) - ${output.type.value}(zero_point_value)) * scale_value`, mirroring the native WebGPU EP's `quantize_linear.cc`.
## Additional context
- Neither case above is covered by the existing test suite. The only per-tensor `uint8` case in `js/web/test/data/ops/dequantizelinear.jsonc` uses `dims: [4]` (too small to trigger bug 1) with `x_zero_point = 0` and no values below the zero point in the failing direction relevant to bug 2 in that shape/coverage. The ONNX node test enabled for `webgpu` (`test_dequantizelinear`) is also `dims: [4]`.
- This may be related to previously reported GPU/WASM divergence on quantized models: microsoft/onnxruntime#25227, huggingface/transformers.js#1512.
- JSEP is currently documented as being in maintenance mode (bug fixes and security fixes only) in favor of the native WebGPU EP, so this may only need a fix if JSEP builds are still expected to be correctness-supported in the interim.
### Urgency
I ran into this issue working on a hobby project of mine. I'm working on a work-around myself. But I believe this is an important issue to fix, since it's causing silent garbage values being produced.
### ONNX Runtime Installation
Built from Source
### ONNX Runtime Version or Commit ID
a7df32cf6087a11884042a2a95526d72100e3b95
### Execution Provider
'webgpu' (WebGPU)
Contributor guide
Research direction
Start in js/web/lib/wasm/jsep/webgpu/ops/quantize-linear.ts and inspect the per-tensor packed/vectorized paths, using the native quantize_linear.cc implementation as a comparison. Add regression coverage to js/web/test/data/ops/dequantizelinear.jsonc and run the WebGPU test_dequantizelinear node test. Done means the two reported uint8 cases match the expected outputs without changing the CPU or native WebGPU paths.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- typescript
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100