saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[js/webgpu] Fix attention shader recompilation issue (#21770)

Browse source 
### Description
<!-- Describe your changes. -->

This PR fixes the `AttentionProbsSoftmax` recompilation issue when
executing the phi3 model. With this fix, it will further improve the
phi3 performance.

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
This commit is contained in:
Jiajia Qin 2024-08-18 08:15:15 +08:00 committed by GitHub
parent 49fc168eed
commit c4ade796d6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
js/web/lib/wasm/jsep/webgpu/ops/attention.ts
View file

@ -262,9 +262,7 @@ const createInPlaceSoftmaxProgramInfo = (_context: ComputeContext, input: Tensor
let WG = 64;
const dComp = d / components;
if (dComp < WG) {
WG = 1;
} else if (dComp / 8 < 64) {
WG = Math.ceil(dComp / 8);
WG = 32;
}
const elementsPerThread = Math.ceil(d / components / WG);
const programUniforms: ProgramUniform[] = [
@ -274,7 +272,7 @@ const createInPlaceSoftmaxProgramInfo = (_context: ComputeContext, input: Tensor
];
const dataType = tensorTypeToWsglStorageType(input.dataType, components);
const f32Type = tensorTypeToWsglValueType(DataType.float, components);
const inputDependencies: ProgramInputTensorInfoDependency[] = ['type'];
const getShaderSource = (shaderHelper: ShaderHelper) => {
const inputHelper = outputVariable('x', input.dataType, input.dims, components);
const elemValueType = tensorTypeToWsglValueType(input.dataType);
@ -353,7 +351,7 @@ const createInPlaceSoftmaxProgramInfo = (_context: ComputeContext, input: Tensor
return {
name: 'AttentionProbsSoftmax',
shaderCache: { hint: `${WG};${dataType};${components}` },
shaderCache: { hint: `${WG};${dataType};${components}`, inputDependencies },
getShaderSource,
getRunData: () => ({ outputs: [], dispatchGroup: { x: n }, programUniforms }),
};