[WebGPU] Subgroups feature is not enabled for ort-web WebGPU EP

[Honry]

From my local debugging when running a GenAI model with ort-web WebGPU EP, I find the context.HasFeature(wgpu::FeatureName::Subgroups) in following file always return false. Thus the ApplyFlashAttention() will never be invoked.

onnxruntime/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

Line 984 in e57dc2a

context.HasFeature(wgpu::FeatureName::Subgroups) &&

However, I can get this feature from dev tool:

The ort-web I built was based on commit c1cbe08

[Honry]

Cc/ @guschmue, @huningxin

[fs-eire]

Could you share a full repo? A fiddler link would be ideal

[Honry]

Could you share a full repo? A fiddler link would be ideal

@fs-eire, this is the debug log I added: Honry@c08eaf6.

And this is the full repo, I built the webgpu ep locally then applied it to my LLM demo, you can visit it via:
https://honry.github.io/webnn-developer-preview/demos/text-generation/?provider=webgpu&model=tinyllama&ort=test

During session creation, the adatper has the Subgroups feature (its id is 17). While when I query the device_feautures_ set, it doesn't contain the Subgroups feature.

Then during inferencing, you will see the context.HasFeature(wgpu::FeatureName::Subgroups) always return false.

(Sorry I don't have a Fiddler, hope above information would help.)

[fs-eire]

Found the bug and fired an issue to upstream: https://issues.chromium.org/issues/435879324

[fs-eire]

Workaround for now:

after the wasm build (before run npm run build in js/web), go to file ort-wasm-simd-threaded.async.mjs, find the following lines: (xx may be different)

"dual-source-blending": "16",
xx: 17

modify the xx into "subgroups", and do the rest build steps.

[Honry]

Thanks @fs-eire! The workaround works.

But there comes a new issue, the result is incorrect after applying the FlashAttention. (Previously without FA, I can get expected result)

Try my refreshed demo.

[fs-eire]

Is the issue reproducible with Node.js binding (using the same JavaScript code and model)? This will help me to understand whether the problem is a web specific problem or not

[Honry]

Is the issue reproducible with Node.js binding (using the same JavaScript code and model)? This will help me to understand whether the problem is a web specific problem or not

Yes, I can reproduce it with Node.js binding. (FA is also applied)

[Honry]

Hold on, there's something wrong in my Node.js binding example. I can't even get expected result with FA disabled.

Then I turn to try the ORT GenAI example with the local built python of onnxruntime with WebGPU EP, it works with FA enabled.

[guschmue]

this is the model ?
https://huggingface.co/webnn/TinyLlama-1.1B-Chat-v1.0-onnx