microsoft / microsoft/onnxruntime

[Web] WebGPU EP: intermediate buffer reused before its consumer reads it — corrupted Div output (int64 Sub→Unsqueeze→Cast→Div chain + parallel ReduceMean branch)

Open
#32,134 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ep:WebGPU platform:web stale
Dominant language
C++
Stars
21.9k
Forks
4.2k
Avg merge
4d 11h
Merged PRs (30d)
184

Description

## Describe the issue

In a graph with two independent branches, the WebGPU EP produces corrupted output for one
branch: the float32 output buffer of a `Cast` (int64→float32) is reused for a later
same-size tensor of the *other* branch before the `Div` consuming the Cast output has read
it. The corrupted `Div` output numerically equals `other_tensor_values / divisor` — i.e. the
Div kernel ran correctly but read a clobbered input buffer.

Minimal structure (7 nodes, no weights; full builder script attached):

- branch 1: `P (int64 [1,194]) → Sub(1) → Unsqueeze(-1) → Cast(float32) → Div(timescale
initializer [1,1,32]) → D [1,194,32]` (the RoPE position-angle pattern emitted by the
PyTorch ONNX exporter for SmolVLA-style models)
- branch 2: `X (float32 [1,194,960]) → Pow(2) → ReduceMean(-1) → Add(eps)` (RMSNorm head)

Observed: `D`'s leading rows contain `(mean(X²)+eps) / timescale` instead of
`positions / timescale`. Removing the `Add` (keeping only Pow→ReduceMean) makes `D` correct;
so does replacing branch 1's int64 chain with a float32 graph input (same Div, same
initializer) — suggesting the aliasing is tied to the int64 emulation treating the Cast
output as reusable.

Properties that point at execution-plan/buffer-pool state rather than a kernel bug:
- deterministic and bit-identical across runs,
- identical in ort-web 1.22.0 and 1.27.0,
- identical at every `graphOptimizationLevel` (disabled/basic/extended/all),
- reproduces on two unrelated Vulkan stacks (AMD RDNA4/RADV and SwiftShader) in Chrome 151,
- wasm EP and Python onnxruntime CPU are exact on the same graphs,
- exposing intermediate tensors as graph outputs (which changes buffer lifetimes) moves or
hides the corruption.

Found while porting a SmolVLA policy to onnxruntime-web: in the full 2094-node model the
same mechanism corrupts the RoPE angles and cascades NaN through all transformer layers.

## To reproduce

1. Run the builder script below (`python build_webgpu_aliasing_repro.py`, needs only `onnx`/`numpy`) — builds
`synth_depth{1..5}.onnx` and `synth_workaround.onnx`.
2. Run each model on identical inputs with EPs `wasm` and `webgpu` (any host page;
`X = uniform(-100,100)`, `P = arange(1,195)`), diff output `D`.
3. `synth_depth1/2`: webgpu == wasm (exact). `synth_depth3/4/5` (Add present): webgpu output
`D` corrupted as described, e.g. `D[0..k] = 3279.0, 2458.9, 1843.9, ...` (=
`mean(X²)+eps` scaled by successive timescale entries) where the reference is
`0/timescale = 0`. `synth_workaround` (float input instead of int64 chain): exact.

## Urgency
No hard deadline; blocks WebGPU deployment of transformer models whose exports contain the
(common) int64 position → Cast → Div RoPE pattern. Workaround exists (feed positions as
float inputs), so not urgent for us.

## System information
- ONNX Runtime Installation: Released Package (npm onnxruntime-web)
- ONNX Runtime Version: 1.27.0 (also reproduced on 1.22.0)
- Execution Provider: 'webgpu' (WebGPU)
- Browser: Chrome 151.0.7922.71, Linux (Arch)
- GPU: AMD Radeon AI PRO R9700 (RADV, Mesa 26.1.7) — also reproduces on SwiftShader
(--headless=new --enable-unsafe-webgpu --enable-features=Vulkan)

---
Disclosure: this bug was investigated and the report drafted with AI assistance (Claude Code). The repro was verified by actually running it as described above on the listed configurations.

build_webgpu_aliasing_repro.py

```python
"""Minimal repro for the onnxruntime-web WebGPU EP buffer-aliasing bug found in M0.

Builds synth_depth{1..5}.onnx plus synth_workaround.onnx (see docs/m0-findings.md).
Each graph has two independent branches:

D-branch: P(int64 [1,194]) -> Sub(1) -> Unsqueeze(-1) -> Cast(float) -> Div(timescale) -> D
X-branch: X(float [1,194,960]) -> Pow(2) [-> ReduceMean(-1) -> Add(eps) -> Sqrt -> Reciprocal]

On the WebGPU EP (ort-web 1.22.0 and 1.27.0, Chrome 151, RADV and SwiftShader alike), as soon
as the X-branch contains the Add (depth >= 3), output D is corrupted: the Cast-output buffer is
reused for the ReduceMean+eps tensor before Div reads it, so D == (mean(X^2)+eps)/timescale for
the leading rows. wasm EP and CPU are always correct. depth 1-2 are clean.

synth_workaround.onnx feeds the positions as a float input F instead of the int64 chain —
bit-exact on WebGPU, demonstrating the aliasing needs the int64->float Cast.

Run each model with the wasm and webgpu EPs on identical inputs and diff output D.
"""
import numpy as np
import onnx
from onnx import helper, numpy_helper, TensorProto as TP
from pathlib import Path

out_dir = Path(__file__).parent
ts = (10000.0 ** ((2.0 / 64) * np.arange(32, dtype=np.float32))).reshape(1, 1, 32)

BRANCH = [
("Pow", ["X", "two"], "pw", [1, 194, 960]),
("ReduceMean", ["pw", "axes"], "mn", [1, 194, 1]),
("Add", ["mn", "eps"], "ad", [1, 194, 1]),
("Sqrt", ["ad"], "sq", [1, 194, 1]),
("Reciprocal", ["sq"], "rc", [1, 194, 1]),
]
INITS = [
numpy_helper.from_array(np.array(1, dtype=np.int64), "one"),
numpy_helper.from_array(np.array([-1], dtype=np.int64), "axm1"),
numpy_helper.from_array(ts, "ts"),
numpy_helper.from_array(np.array(2.0, dtype=np.float32), "two"),
numpy_helper.from_array(np.array([-1], dtype=np.int64), "axes"),
numpy_helper.from_array(np.array(1e-6, dtype=np.float32), "eps"),
]

def build(name, d_branch_nodes, extra_inputs, depth):
nodes = list(d_branch_nodes)
for op, ins, out, _ in BRANCH[:depth]:
kw = {"keepdims": 1} if op == "ReduceMean" else {}
nodes.append(helper.make_node(op, ins, [out], **kw))
last, lshape = BRANCH[depth - 1][2], BRANCH[depth - 1][3]
graph = helper.make_graph(nodes, name,
[helper.make_tensor_value_info("X", TP.FLOAT, [1, 194, 960])] + extra_inputs,
[helper.make_tensor_value_info("D", TP.FLOAT, [1, 194, 32]),
helper.make_tensor_value_info(last, TP.FLOAT, lshape)],
initializer=INITS)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
onnx.checker.check_model(model)
onnx.save(model, out_dir / f"{name}.onnx")
print("built", name)

int64_chain = [
helper.make_node("Sub", ["P", "one"], ["s"]),
helper.make_node("Unsqueeze", ["s", "axm1"], ["u"]),
helper.make_node("Cast", ["u"], ["f"], to=TP.FLOAT),
helper.make_node("Div", ["f", "ts"], ["D"]),
]
for depth in range(1, 6):
build(f"synth_depth{depth}", int64_chain,
[helper.make_tensor_value_info("P", TP.INT64, [1, 194])], depth)

build("synth_workaround", [helper.make_node("Div", ["F", "ts"], ["D"])],
[helper.make_tensor_value_info("F", TP.FLOAT, [1, 194, 1])], 3)
```

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running build_webgpu_aliasing_repro.py and comparing the generated models with the wasm and webgpu execution providers. Investigate the WebGPU execution-plan and buffer-pool state around the int64 Sub→Unsqueeze→Cast→Div branch and the parallel ReduceMean branch; done means WebGPU matches wasm without corrupted D output across the supplied depth models.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, javascript, python
Domain
backend, machine-learning, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.