Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[webgpu] Use subgroup for matmulnbits #23224

Open
wants to merge 17 commits into
base: main
Choose a base branch
from
Open

[webgpu] Use subgroup for matmulnbits #23224

wants to merge 17 commits into from
+293 −127

Conversation

qjia7
Copy link
Contributor

@qjia7 qjia7 commented Dec 30, 2024 •
edited
Loading

Description

This PR applies subgroup to implement matmulnbits when tile_m > 1 for intel devices.
With this PR, prefill for 500 tokens prompt for phi3 becomes 3.5s from 8.5s on intel Meteor Lake.

qjia7 added 17 commits December 27, 2024 11:42
@qjia7
[webgpu] Add kernel type to profile info
3b32acc
@qjia7
relax some limitations
8e0b27f
@qjia7
[webgpu] Use override shape in shader key
c0dd1db
@qjia7
subgroup support only for subgroup size = 32
a66dd01 
In this version, the local_id.x = 8, local_id.y = 4.
To load A data, each thread needs to access memory twice so that the
tile size A is 64 x 4.

In order to get the correct shuffle A when acculate the inter_results,
we have to unconditionally get a_data low and a data high. Then use
select to decide which data to use for current thread.

So this method doubles the shuffle commands.
@qjia7
support subgroup with size = 32
d3a196b 
Use local_id.x = 4, local_id.y = 8 so that we only need to suffle once
compared with previous method.
@qjia7
subgroup size = 32, workgroup size = 64
99fe697
@qjia7
put a/scale data into workgroup memory
0e85048
@qjia7
support subgroup size = 8/16/32/64
6feb70f
@qjia7
workaround the issue that subgroupShuffle is very slow
55fc6bb 
When the index is a variable
@qjia7
add sg_size check
128cd7d
@qjia7
Revert "add sg_size check"
0775e4e 
This reverts commit 128cd7d.
@qjia7
apply subgroup for tile_m > 1
454fd0d
@qjia7
check subgroup feature
16fc075
@qjia7
only apply subgroup for intel GPUs
622bdd0
@qjia7
Revert "[webgpu] Use override shape in shader key"
729df58 
This reverts commit c0dd1db.
@qjia7
Revert "[webgpu] Add kernel type to profile info"
dff1961 
This reverts commit 3b32acc.
@qjia7
add more limitations to use subgroup
8328e14
@qjia7 qjia7 changed the title [Not for Review] [webgpu] Test subgroup for matmulnbits [webgpu] Use subgroup for matmulnbits Jan 2, 2025
@qjia7 qjia7 marked this pull request as ready for review January 2, 2025 04:47
@qjia7
Copy link
Contributor Author

qjia7 commented Jan 2, 2025

@guschmue @fs-eire Please take a look. Currently, this PR is only applied on Intel devices since I don't see perf improvement on NV device on my hand. Need some time to investigate the reason.

Please also let me know whether it works on your xe devices.
If you want to check other GPUs like, mac, you can just change

onnxruntime/onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc

Line 572 in 8328e14

const bool use_subgroup = context.Device().HasFeature(wgpu::FeatureName::Subgroups) && context.AdapterInfo().vendor == std::string_view{"intel"} && components_a == 4 && block_size == 32;

to

  const bool use_subgroup = context.Device().HasFeature(wgpu::FeatureName::Subgroups) && components_a == 4 && block_size == 32;

cc @sushraja-msft

@guschmue
Copy link
Contributor

guschmue commented Jan 2, 2025

/azp run ONNX Runtime Web CI Pipeline,Windows GPU CI Pipeline,Linux Android Emulator QNN CI Pipeline

@guschmue
Copy link
Contributor

guschmue commented Jan 2, 2025

/azp run Linux CPU CI Pipeline,Linux CPU Minimal Build E2E CI Pipeline,Linux GPU CI Pipeline,Linux GPU TensorRT CI Pipeline,Linux OpenVINO CI Pipeline,Linux QNN CI Pipeline,MacOS CI Pipeline,Windows ARM64 QNN CI Pipeline,Windows CPU CI Pipeline

@azure-pipelines Azure Pipelines
Copy link

Azure Pipelines successfully started running 2 pipeline(s).

@guschmue
Copy link
Contributor

guschmue commented Jan 2, 2025

/azp run Windows GPU TensorRT CI Pipeline,onnxruntime-binary-size-checks-ci-pipeline,orttraining-linux-ci-pipeline,orttraining-linux-gpu-ci-pipeline,orttraining-ortmodule-distributed,Windows x64 QNN CI Pipeline,Big Models

@guschmue
Copy link
Contributor

guschmue commented Jan 2, 2025

/azp run Windows GPU CUDA CI Pipeline,Windows GPU DML CI Pipeline,Windows GPU Doc Gen CI Pipeline

@azure-pipelines Azure Pipelines
Copy link

Azure Pipelines successfully started running 4 pipeline(s).

@azure-pipelines Azure Pipelines
Copy link

Azure Pipelines successfully started running 3 pipeline(s).

@azure-pipelines Azure Pipelines
Copy link

Azure Pipelines successfully started running 9 pipeline(s).

@guschmue guschmue added the ep:WebGPU ort-web webgpu provider label Jan 2, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
ep:WebGPU ort-web webgpu provider
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
2 participants
@qjia7 @guschmue