Skip to content

[webgpu] Implement SubGroupMatrix based MatMulNBits for Metal - #23729

Merged
Guenther Schmuelling (guschmue) merged 15 commits into
mainfrom
user/sushraja/subgroupMatrix
Feb 21, 2025
Merged

[webgpu] Implement SubGroupMatrix based MatMulNBits for Metal#23729
Guenther Schmuelling (guschmue) merged 15 commits into
mainfrom
user/sushraja/subgroupMatrix

Conversation

@sushraja-msft

@sushraja-msft Sushanth Rajasankar (sushraja-msft) commented Feb 17, 2025

Copy link
Copy Markdown
Contributor

Description

Recent progress with SubGroupMatrix prototype in Dawn https://issues.chromium.org/issues/348702031, exposes SIMD-Group Matrix Functions to webgpu. This shader implements a matmulnbits using that primitive.

Observed perf gains, in terms of LLM inference speed, prefill perf for Phi 3.5 for a 1K token prefill see 3x improvement. 5.4s from 15s.

With Changes

./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       5.42498e+06                    <<< SubGroupMatrix 5.4s
	avg (tokens/s): 184.517
	p50 (us):       5.41982e+06
	stddev (us):    12023.8
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       91138.5
	avg (tokens/s): 10.9723
	p50 (us):       89488.5
	stddev (us):    35136.2
	n:              635 * 1 token(s)

Baseline

./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       1.45507e+07                     <<< Baseline 14.5s
	avg (tokens/s): 68.7938
	p50 (us):       1.45413e+07
	stddev (us):    22208.9
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       94109.8
	avg (tokens/s): 10.6259
	p50 (us):       89660
	stddev (us):    61579
	n:              635 * 1 token(s)

@github-actions github-actions Bot 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.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/contrib_ops/webgpu/quantization/matmul_nbits.cc
Comment thread onnxruntime/contrib_ops/webgpu/quantization/subgroup_matrix_matmul_nbits.cc Outdated
Comment thread onnxruntime/core/providers/webgpu/webgpu_context.h

@github-actions github-actions Bot 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.

You can commit the suggested changes from lintrunner.

Comment thread onnxruntime/core/providers/coreml/model/model.mm Outdated
Comment thread onnxruntime/core/providers/coreml/model/model.mm Outdated
@guschmue

Guenther Schmuelling (guschmue) commented Feb 20, 2025

Copy link
Copy Markdown
Contributor

the ort web pipeline compiles webgpu ep with emscripten which fails with:
shader_helper.cc:355:45: error: no member named 'ChromiumExperimentalSubgroupMatrix' in 'wgpu::FeatureName'
355 | if (device_.HasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix))

Possible the headerfile that comes with emscripten doesn't know that featurename yet

Maybe use
#if !defined(wasm)

@sushraja-msft

Copy link
Copy Markdown
Contributor Author

the ort web pipeline compiles webgpu ep with emscripten which fails with: shader_helper.cc:355:45: error: no member named 'ChromiumExperimentalSubgroupMatrix' in 'wgpu::FeatureName' 355 | if (device_.HasFeature(wgpu::FeatureName::ChromiumExperimentalSubgroupMatrix))

Possible the headerfile that comes with emscripten doesn't know that featurename yet

Maybe use #if !defined(wasm)

done !

@fs-eire

Copy link
Copy Markdown
Contributor

/azp run Android CI Pipeline,iOS CI Pipeline,ONNX Runtime React Native CI Pipeline,CoreML CI Pipeline,Linux DNNL CI Pipeline,Linux MIGraphX CI Pipeline,Linux ROCm CI Pipeline

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines successfully started running 7 pipeline(s).

@guschmue
Guenther Schmuelling (guschmue) deleted the user/sushraja/subgroupMatrix branch February 21, 2025 17:23
Guenther Schmuelling (guschmue) pushed a commit that referenced this pull request Mar 6, 2025
### Description
Recent progress with SubGroupMatrix prototype in Dawn
https://issues.chromium.org/issues/348702031, exposes SIMD-Group Matrix
Functions to webgpu. This shader implements a matmulnbits using that
primitive.

Observed perf gains, in terms of LLM inference speed, prefill perf for
Phi 3.5 for a 1K token prefill see 3x improvement. 5.4s from 15s.

With Changes
```
./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       5.42498e+06                    <<< SubGroupMatrix 5.4s
	avg (tokens/s): 184.517
	p50 (us):       5.41982e+06
	stddev (us):    12023.8
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       91138.5
	avg (tokens/s): 10.9723
	p50 (us):       89488.5
	stddev (us):    35136.2
	n:              635 * 1 token(s)

```
Baseline
```
./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       1.45507e+07                     <<< Baseline 14.5s
	avg (tokens/s): 68.7938
	p50 (us):       1.45413e+07
	stddev (us):    22208.9
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       94109.8
	avg (tokens/s): 10.6259
	p50 (us):       89660
	stddev (us):    61579
	n:              635 * 1 token(s)
```
Ashrit Shetty (ashrit-ms) pushed a commit that referenced this pull request Mar 17, 2025
### Description
Recent progress with SubGroupMatrix prototype in Dawn
https://issues.chromium.org/issues/348702031, exposes SIMD-Group Matrix
Functions to webgpu. This shader implements a matmulnbits using that
primitive.

Observed perf gains, in terms of LLM inference speed, prefill perf for
Phi 3.5 for a 1K token prefill see 3x improvement. 5.4s from 15s.

With Changes
```
./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       5.42498e+06                    <<< SubGroupMatrix 5.4s
	avg (tokens/s): 184.517
	p50 (us):       5.41982e+06
	stddev (us):    12023.8
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       91138.5
	avg (tokens/s): 10.9723
	p50 (us):       89488.5
	stddev (us):    35136.2
	n:              635 * 1 token(s)

```
Baseline
```
./model_benchmark -i ~/Phi-3.5-mini-instruct-onnx-web -l 1000
Batch size: 1, prompt tokens: 1001, tokens to generate: 128
Prompt processing (time to first token):
	avg (us):       1.45507e+07                     <<< Baseline 14.5s
	avg (tokens/s): 68.7938
	p50 (us):       1.45413e+07
	stddev (us):    22208.9
	n:              5 * 1001 token(s)
Token generation:
	avg (us):       94109.8
	avg (tokens/s): 10.6259
	p50 (us):       89660
	stddev (us):    61579
	n:              635 * 1 token(s)
```
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.

6 participants