[Web] Upgrading from 1.20.1 to 1.21.* breaks Segment Anything models on WebGPU

[xenova]

Describe the issue

After upgrading to onnxruntime-web version 1.21.0-dev.20241205-d27fecd3d3 in this commit, Segment Anything models no longer work on WebGPU. WASM is correct (albeit slow). You can use https://huggingface.co/Xenova/slimsam-77-uniform for testing, as shown below.

Correct behaviour:

Incorrect behaviour:

To reproduce

1. Check out the demo source code: https://github.com/huggingface/transformers.js-examples/blob/main/segment-anything-webgpu/index.js
2. Update the import to https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.1.1 and see that everything works correctly.
3. Update the import to https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.1.2 and see that the output is incorrect.

Urgency

Since Segment Anything is quite a popular Transformers.js model, this is pretty important.

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.21.0-dev.20241205-d27fecd3d3

Execution Provider

'webgpu' (WebGPU)

[xenova]

cc @fs-eire @guschmue

[fs-eire]

After some investigation, I found the root cause of the issue is #22755. If I revert the change and rebuild, the SAM demo will work.

There should be some problem of the ScatterND WebGPU implementation. Need further investigation.

[xenova]

Thanks for the update! Please let me know if you put out a build with the change reverted (or whether you can put out a build that has it fixed).

[prathikr]

@xenova will look into this in the next day or two. If I cannot fix by then I will revert the PR to mitigate the issue while continuing to investigate. I appreciate your patience.

[xenova]

Thanks @prathikr! If you could publish a dev build with the PR reverted, that could work too! Hopefully it can also include #23026 too.

[xenova]

Also, if it helps debug, it appears to be an issue w/ fp16 or q8 (not fp32): huggingface/transformers.js#1134 (comment).

[fs-eire]

@xenova please try the latest dev package: onnxruntime-web@1.21.0-dev.20250108-d8e8d4fac0