Skip to content

[WebGPU] LinearAttention: increase tile_v when subgroups are available - #28519

Merged
Guenther Schmuelling (guschmue) merged 6 commits into
microsoft:mainfrom
daijh:linear-attn-dev
May 22, 2026
Merged

[WebGPU] LinearAttention: increase tile_v when subgroups are available#28519
Guenther Schmuelling (guschmue) merged 6 commits into
microsoft:mainfrom
daijh:linear-attn-dev

Conversation

@daijh

@daijh Jianhui Dai (daijh) commented May 15, 2026

Copy link
Copy Markdown
Contributor

Description

  • Scale tile_v by 4x when subgroup is enabled and the vectorized dimension has enough columns, improving data reuse.
  • Gate the tile_v expansion on seq_length >= 16, where the prefill benefit outweighs increased register pressure.
  • Remove redundant zero-initialization of the state tile (WGSL default- initializes vars to zero).

Intel Panther Lake (xe-3lpg)

Model Prefill Baseline (TPS) Optimized (TPS) Change
Qwen3.5-0.8B 128 1534.30 1681.00 9.56%
Qwen3.5-0.8B 1024 3267.30 3917.60 19.90%
Qwen3.5-0.8B 4096 2864.50 3563.40 24.40%
Qwen3.5-2B 128 1295.60 1344.60 3.78%
Qwen3.5-2B 1024 2177.10 2344.00 7.67%
Qwen3.5-2B 4096 1942.60 2247.00 15.67%
Qwen3.5-4B 128 701.30 736.20 4.98%
Qwen3.5-4B 1024 946.90 1036.30 9.44%
Qwen3.5-4B 4096 824.10 912.00 10.67%

Motivation and Context

See above.

@daijh
Jianhui Dai (daijh) marked this pull request as draft May 15, 2026 03:44
- Scale tile_v by 4x when subgroup is enabled and the vectorized
  dimension has enough columns, improving data reuse.
- Remove redundant zero-initialization of the state tile (WGSL default-
  initializes private vars to zero).
@guschmue Guenther Schmuelling (guschmue) added the ep:WebGPU ort-web webgpu provider label May 18, 2026
@daijh
Jianhui Dai (daijh) marked this pull request as ready for review May 19, 2026 08:40
@daijh

Copy link
Copy Markdown
Contributor Author

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR optimizes the WebGPU LinearAttention kernel by increasing the V-dimension tiling when subgroup reductions are available (to improve data reuse), gating that expansion to longer sequences, and removing redundant per-invocation state zeroing in the WGSL shader.

Changes:

  • Adjust tile_v selection logic and expand it (up to 4×) when subgroups are enabled and seq_length >= 16.
  • Move subgroup_min_size determination earlier and pass it into the shader program configuration.
  • Remove explicit zero-initialization of the per-thread state tile in the WGSL template.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File Description
onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc Adjusts tile selection/expansion heuristics and subgroup configuration passed into the shader.
onnxruntime/contrib_ops/webgpu/bert/linear_attention.wgsl.template Removes explicit state zero-initialization (relies on WGSL default initialization).

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc
@daijh

Copy link
Copy Markdown
Contributor Author

Jiajia Qin (@qjia7) Thanks. Fixed the comment, please take another look.

Comment thread onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc Outdated
@guschmue
Guenther Schmuelling (guschmue) merged commit 1053327 into microsoft:main May 22, 2026
92 of 93 checks passed
@daijh
Jianhui Dai (daijh) deleted the linear-attn-dev branch May 25, 2026 02:48
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ep:WebGPU ort-web webgpu provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants