mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
[JS/WebGPU] Multihead attention improvements (#20286)
### Description Enabled more usecases ### 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:
Satya Kumar Jandhyala
GitHub
parent
76461c8f4d
commit
d42ac7f0c6
11 changed files with 228 additions and 136 deletions
|
|
@ -405,7 +405,8 @@ export class WebGpuBackend {
|
|||
*/
|
||||
run(program: ProgramInfo, inputTensorViews: readonly TensorView[], outputIndices: readonly number[],
|
||||
createKernelOutput: (index: number, dataType: number, dims: readonly number[]) => TensorView,
|
||||
createIntermediateOutput: (dataType: number, dims: readonly number[]) => TensorView): TensorView[] {
|
||||
createIntermediateOutput: (dataType: number, dims: readonly number[]) => TensorView,
|
||||
outputCount: number): TensorView[] {
|
||||
TRACE_FUNC_BEGIN(program.name);
|
||||
// create info for inputs
|
||||
const inputDatas: GpuData[] = [];
|
||||
|
|
@ -438,7 +439,7 @@ export class WebGpuBackend {
|
|||
// value -3 is used for placeholder output. So -3, -2, -1 and 0, 1, 2, ... are valid
|
||||
// output indices. see type definition of ComputeContextInputsOutputsMapping for more details.
|
||||
if (!Number.isInteger(validatedOutputIndices[i]) || validatedOutputIndices[i] < -3 ||
|
||||
validatedOutputIndices[i] >= outputs.length) {
|
||||
validatedOutputIndices[i] >= outputCount) {
|
||||
throw new Error(`Invalid output index: ${validatedOutputIndices[i]}`);
|
||||
}
|
||||
if (validatedOutputIndices[i] === -3) {
|
||||
|
|
|
|||
|
|
@ -120,7 +120,8 @@ class ComputeContextImpl implements ComputeContext {
|
|||
const gpuDataId = bufferSize > 0 ? this.backend.gpuDataManager.create(bufferSize).id : 0;
|
||||
return new TensorViewImpl(this.module, dataType, gpuDataId, dims);
|
||||
};
|
||||
return this.backend.run(program, mappedInputs, outputIndices, createKernelOutput, createTemporaryOutput);
|
||||
return this.backend.run(
|
||||
program, mappedInputs, outputIndices, createKernelOutput, createTemporaryOutput, this.outputCount);
|
||||
}
|
||||
|
||||
output(index: number, dims: readonly number[]): number {
|
||||
|
|
|
|||
|
|
@ -22,7 +22,7 @@ import {instanceNorm} from './ops/instance-norm';
|
|||
import {layerNorm} from './ops/layer-norm';
|
||||
import {matMul} from './ops/matmul';
|
||||
import {matMulNBits, parseMatMulNBitsAttributes} from './ops/matmulnbits';
|
||||
import {multiHeadAttention, parseMultiHeadAttentionAttributes} from './ops/multi-head-attentiion';
|
||||
import {multiHeadAttention, parseMultiHeadAttentionAttributes} from './ops/multihead-attentiion';
|
||||
import {pad} from './ops/pad';
|
||||
import * as pool from './ops/pool';
|
||||
import {range} from './ops/range';
|
||||
|
|
|
|||
|
|
@ -3,9 +3,10 @@
|
|||
|
||||
import {DataType} from '../../../wasm-common';
|
||||
import {TensorView} from '../../tensor-view';
|
||||
import {ComputeContext, GpuDataType, ProgramUniform} from '../types';
|
||||
import {ComputeContext, GpuDataType, ProgramInputTensorInfoDependency, ProgramUniform} from '../types';
|
||||
|
||||
import {castToF32, fillVector, getMaxComponents, inputVariable, outputVariable, ShaderHelper, sumVector, tensorTypeToWsglStorageType, tensorTypeToWsglValueType, UniformDataElementType, UniformsArrayType} from './common';
|
||||
import {createTensorShapeVariables, getMaxComponents, inputVariable, outputVariable, ShaderHelper, tensorTypeToWsglStorageType, tensorTypeToWsglValueType, UniformDataElementType, UniformsArrayType} from './common';
|
||||
import {createConcatProgramInfo} from './concat';
|
||||
|
||||
export const enum AttentionQkvFormat {
|
||||
unknown, // enum value not set, or depends on qkv projection implementation details
|
||||
|
|
@ -203,9 +204,6 @@ const validateAttentionInputs = (inputs: readonly TensorView[], attributes: Atte
|
|||
if (past) {
|
||||
throw new Error('past is not supported');
|
||||
}
|
||||
if (relativePositionBias) {
|
||||
throw new Error('relativePositionBias is not supported');
|
||||
}
|
||||
|
||||
return {
|
||||
batchSize,
|
||||
|
|
@ -231,7 +229,7 @@ const validateAttentionInputs = (inputs: readonly TensorView[], attributes: Atte
|
|||
};
|
||||
};
|
||||
|
||||
export const computeInPlaceSoftmax = (context: ComputeContext, input: TensorView, n: number, d: number) => {
|
||||
const createInPlaceSoftmaxProgramInfo = (_context: ComputeContext, input: TensorView, n: number, d: number) => {
|
||||
const components = getMaxComponents(d);
|
||||
let WG = 64;
|
||||
const dComp = d / components;
|
||||
|
|
@ -240,92 +238,105 @@ export const computeInPlaceSoftmax = (context: ComputeContext, input: TensorView
|
|||
} else if (dComp / 8 < 64) {
|
||||
WG = Math.ceil(dComp / 8);
|
||||
}
|
||||
const elementsPerWG = Math.ceil(d / components / WG);
|
||||
const elementsPerThread = Math.ceil(d / components / WG);
|
||||
const programUniforms: ProgramUniform[] = [
|
||||
{type: input.dataType, data: 1 / d}, {type: DataType.uint32, data: dComp},
|
||||
{type: DataType.uint32, data: elementsPerWG}
|
||||
{type: DataType.uint32, data: elementsPerThread}
|
||||
];
|
||||
const dataType = tensorTypeToWsglStorageType(input.dataType, components);
|
||||
const f32Type = tensorTypeToWsglValueType(DataType.float, components);
|
||||
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => {
|
||||
const inputHelper = outputVariable('x', input.dataType, input.dims, components);
|
||||
let threadMaxValue = 'thread_max_vector';
|
||||
if (components === 2) {
|
||||
threadMaxValue = 'max(thread_max_vector.x, thread_max_vector.y)';
|
||||
} else if (components === 4) {
|
||||
threadMaxValue =
|
||||
'max(max(thread_max_vector.x, thread_max_vector.y), max(thread_max_vector.z, thread_max_vector.w))';
|
||||
}
|
||||
const elemValueType = tensorTypeToWsglValueType(input.dataType);
|
||||
const uniforms: UniformsArrayType = [
|
||||
{name: 'd_inv', type: elemValueType as UniformDataElementType}, {name: 'd_comp', type: 'u32'},
|
||||
{name: 'elements_per_wg', type: 'u32'}
|
||||
{name: 'elements_per_thread', type: 'u32'}
|
||||
];
|
||||
|
||||
return `
|
||||
var<workgroup> wgMax: array<f32, ${WG}>;
|
||||
var<workgroup> wgSum: array<f32, ${WG}>;
|
||||
var<workgroup> thread_max: array<f32, ${WG}>;
|
||||
var<workgroup> thread_sum: array<f32, ${WG}>;
|
||||
${shaderHelper.registerUniforms(uniforms).declareVariables(inputHelper)}
|
||||
${shaderHelper.mainStart([
|
||||
WG, 1, 1
|
||||
])}
|
||||
let localOffset = local_idx * uniforms.elements_per_wg;
|
||||
let offset: u32 = workgroup_id.x * uniforms.d_comp + localOffset;
|
||||
let local_offset = local_idx * uniforms.elements_per_thread;
|
||||
let offset = workgroup_id.x * uniforms.d_comp + local_offset;
|
||||
|
||||
var thread_max_vector = ${fillVector('f32', components, '-3.402823e+38f')};
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {
|
||||
thread_max_vector = max(${castToF32(elemValueType, components, 'x[offset + i]')}, thread_max_vector);
|
||||
var thread_max_vector = ${inputHelper.type.value}(-3.402823e+38f);
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_thread && i + local_offset < uniforms.d_comp; i++) {
|
||||
thread_max_vector = max(${f32Type}(x[offset + i]), thread_max_vector);
|
||||
}
|
||||
wgMax[local_idx] = ${threadMaxValue};
|
||||
thread_max[local_idx] = ${(() => {
|
||||
switch (components) {
|
||||
case 1:
|
||||
return 'thread_max_vector';
|
||||
case 2:
|
||||
return 'max(thread_max_vector.x, thread_max_vector.y)';
|
||||
case 4:
|
||||
return 'max(max(thread_max_vector.x, thread_max_vector.y), max(thread_max_vector.z, thread_max_vector.w))';
|
||||
default:
|
||||
throw new Error(`Unsupported components: ${components}`);
|
||||
}
|
||||
})()};
|
||||
workgroupBarrier();
|
||||
|
||||
var maxValue = -3.402823e+38f;
|
||||
var max_value: f32 = -3.402823e+38f;
|
||||
for (var i = 0u; i < ${WG}; i++) {
|
||||
maxValue = max(wgMax[i], maxValue);
|
||||
max_value = max(thread_max[i], max_value);
|
||||
}
|
||||
|
||||
var sumVector = ${fillVector('f32', components, '0')};
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {
|
||||
sumVector += exp(${castToF32(elemValueType, components, 'x[offset + i]')} - maxValue);
|
||||
var sum_vector = ${inputHelper.type.value}(${0});
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_thread && i + local_offset < uniforms.d_comp; i++) {
|
||||
sum_vector += exp(${f32Type}(x[offset + i]) - max_value);
|
||||
}
|
||||
wgSum[local_idx] = ${sumVector('sumVector', components)};
|
||||
thread_sum[local_idx] = ${(() => {
|
||||
switch (components) {
|
||||
case 1:
|
||||
return 'sum_vector';
|
||||
case 2:
|
||||
return 'sum_vector.x + sum_vector.y';
|
||||
case 4:
|
||||
return 'sum_vector.x + sum_vector.y + sum_vector.z + sum_vector.w';
|
||||
default:
|
||||
throw new Error(`Unsupported components: ${components}`);
|
||||
}
|
||||
})()};
|
||||
workgroupBarrier();
|
||||
|
||||
var sum: f32 = 0;
|
||||
for (var i = 0u; i < ${WG}; i++) {
|
||||
sum += wgSum[i];
|
||||
sum += thread_sum[i];
|
||||
}
|
||||
|
||||
if (sum == 0) {
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {
|
||||
x[offset + i] = ${fillVector(elemValueType, components, 'uniforms.d_inv')};
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_thread && i + local_offset < uniforms.d_comp; i++) {
|
||||
x[offset + i] = ${inputHelper.type.value}(uniforms.d_inv);
|
||||
}
|
||||
} else {
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_wg && i + localOffset < uniforms.d_comp; i++) {
|
||||
let f32input = ${castToF32(elemValueType, components, 'x[offset + i]')};
|
||||
x[offset + i] = ${inputHelper.type.value}(exp(f32input - maxValue) / sum);
|
||||
for (var i: u32 = 0; i < uniforms.elements_per_thread && i + local_offset < uniforms.d_comp; i++) {
|
||||
var f32input = ${f32Type}(x[offset + i]);
|
||||
x[offset + i] = ${inputHelper.type.value}(exp(f32input - max_value) / sum);
|
||||
}
|
||||
}
|
||||
}`;
|
||||
};
|
||||
|
||||
context.compute(
|
||||
{
|
||||
name: 'AttentionProbsSoftmax',
|
||||
shaderCache: {hint: `${WG};${dataType};${components}`},
|
||||
getShaderSource,
|
||||
getRunData: () => ({outputs: [], dispatchGroup: {x: n}, programUniforms}),
|
||||
},
|
||||
{inputs: [input], outputs: []});
|
||||
return {
|
||||
name: 'AttentionProbsSoftmax',
|
||||
shaderCache: {hint: `${WG};${dataType};${components}`},
|
||||
getShaderSource,
|
||||
getRunData: () => ({outputs: [], dispatchGroup: {x: n}, programUniforms}),
|
||||
};
|
||||
};
|
||||
|
||||
const computeAttentionProbs =
|
||||
(context: ComputeContext, q: TensorView, key: TensorView, _bias: TensorView|undefined,
|
||||
const createAttentionProbsProgramInfo =
|
||||
(_context: ComputeContext, q: TensorView, key: TensorView, relativePositionBias: TensorView|undefined,
|
||||
parameters: AttentionParameters, attributes: AttentionAttrs) => {
|
||||
const probsShape = [
|
||||
parameters.batchSize, parameters.numHeads, parameters.sequenceLength,
|
||||
parameters.kvSequenceLength + parameters.pastSequenceLength
|
||||
];
|
||||
const probsShape =
|
||||
[parameters.batchSize, parameters.numHeads, parameters.sequenceLength, parameters.totalSequenceLength];
|
||||
|
||||
// TODO: handle mask
|
||||
|
||||
const alpha = attributes.scale === 0 ? 1.0 / Math.sqrt(parameters.headSize) : attributes.scale;
|
||||
|
|
@ -340,20 +351,33 @@ const computeAttentionProbs =
|
|||
const programUniforms: ProgramUniform[] = [
|
||||
{type: DataType.uint32, data: parameters.sequenceLength}, {type: DataType.uint32, data: vectorizedHeadSize},
|
||||
{type: DataType.uint32, data: parameters.totalSequenceLength},
|
||||
{type: DataType.uint32, data: parameters.kvSequenceLength}, {type: q.dataType, data: alpha}
|
||||
{type: DataType.uint32, data: parameters.numHeads}, {type: DataType.uint32, data: parameters.kvSequenceLength},
|
||||
{type: q.dataType, data: alpha}
|
||||
];
|
||||
|
||||
const inputs = [q, key];
|
||||
const inputDependencies: ProgramInputTensorInfoDependency[] = ['type', 'type'];
|
||||
if (relativePositionBias) {
|
||||
inputDependencies.push('rank');
|
||||
programUniforms.push(...createTensorShapeVariables(relativePositionBias.dims));
|
||||
}
|
||||
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => {
|
||||
const qInput = inputVariable('q', q.dataType, q.dims, components);
|
||||
const kInput = inputVariable('key', key.dataType, key.dims, components);
|
||||
const inputVars = [qInput, kInput];
|
||||
const relativePositionBiasInput = relativePositionBias ?
|
||||
inputVariable('relative_position_bias', relativePositionBias.dataType, relativePositionBias.dims.length) :
|
||||
undefined;
|
||||
if (relativePositionBiasInput) {
|
||||
inputVars.push(relativePositionBiasInput);
|
||||
}
|
||||
const output = outputVariable('output', q.dataType, probsShape);
|
||||
const dataType = tensorTypeToWsglStorageType(q.dataType);
|
||||
|
||||
const uniforms: UniformsArrayType = [
|
||||
{name: 'M', type: 'u32'}, {name: 'K', type: 'u32'}, {name: 'N', type: 'u32'},
|
||||
{name: 'kv_sequence_length', type: 'u32'}, {name: 'alpha', type: dataType as UniformDataElementType}
|
||||
{name: 'num_heads', type: 'u32'}, {name: 'kv_sequence_length', type: 'u32'},
|
||||
{name: 'alpha', type: dataType as UniformDataElementType}
|
||||
];
|
||||
return `
|
||||
const beta: ${dataType} = 1.0;
|
||||
|
|
@ -361,7 +385,7 @@ const computeAttentionProbs =
|
|||
|
||||
var<workgroup> tileQ: array<${qInput.type.storage}, ${TILE_SIZE * TILE_SIZE}>;
|
||||
var<workgroup> tileK: array<${qInput.type.storage}, ${TILE_SIZE * TILE_SIZE}>;
|
||||
${shaderHelper.registerUniforms(uniforms).declareVariables(qInput, kInput, output)}
|
||||
${shaderHelper.registerUniforms(uniforms).declareVariables(...inputVars, output)}
|
||||
${shaderHelper.mainStart([
|
||||
TILE_SIZE, TILE_SIZE, 1
|
||||
])}
|
||||
|
|
@ -372,7 +396,7 @@ const computeAttentionProbs =
|
|||
let qOffset = uniforms.M * uniforms.K * headIdx + m * uniforms.K;
|
||||
let kOffset = uniforms.kv_sequence_length * uniforms.K * headIdx + n * uniforms.K;
|
||||
|
||||
var value = ${fillVector(dataType, components)};
|
||||
var value = ${qInput.type.value}(0);
|
||||
for (var w: u32 = 0u; w < uniforms.K; w += TILE_SIZE) {
|
||||
if (global_id.y < uniforms.M && w + local_id.x < uniforms.K) {
|
||||
tileQ[TILE_SIZE * local_id.y + local_id.x] = q[qOffset + local_id.y * uniforms.K + w + local_id.x];
|
||||
|
|
@ -382,7 +406,7 @@ const computeAttentionProbs =
|
|||
}
|
||||
workgroupBarrier();
|
||||
|
||||
for (var k: u32 = 0u; k<TILE_SIZE && w+k < uniforms.K; k++) {
|
||||
for (var k: u32 = 0u; k < TILE_SIZE && w+k < uniforms.K; k++) {
|
||||
value += tileQ[TILE_SIZE * local_id.y + k] * tileK[TILE_SIZE * local_id.x + k];
|
||||
}
|
||||
|
||||
|
|
@ -392,33 +416,47 @@ const computeAttentionProbs =
|
|||
let headOffset = headIdx * uniforms.M * uniforms.N;
|
||||
if (global_id.y < uniforms.M && global_id.x < uniforms.N) {
|
||||
let outputIdx = headOffset + global_id.y * uniforms.N + global_id.x;
|
||||
output[outputIdx] = ${sumVector('value', components)} * uniforms.alpha;
|
||||
var sum = ${(() => {
|
||||
switch (components) {
|
||||
case 1:
|
||||
return 'value';
|
||||
case 2:
|
||||
return 'value.x + value.y';
|
||||
case 4:
|
||||
return 'value.x + value.y + value.z + value.w';
|
||||
default:
|
||||
throw new Error(`Unsupported components: ${components}`);
|
||||
}
|
||||
})()};
|
||||
output[outputIdx] = sum * uniforms.alpha;
|
||||
${(() => {
|
||||
if (relativePositionBiasInput) {
|
||||
return `
|
||||
let batch = workgroup_id.z / uniforms.num_heads;
|
||||
let head = workgroup_id.z % uniforms.num_heads;
|
||||
var indices = ${relativePositionBiasInput.type.indices}(batch, head, global_id.y, global_id.x);
|
||||
output[outputIdx] += ${relativePositionBiasInput.getByIndices('indices')};`;
|
||||
}
|
||||
return '';
|
||||
})()}
|
||||
}
|
||||
}`;
|
||||
};
|
||||
|
||||
const probs = context.compute(
|
||||
{
|
||||
name: 'AttentionProbs',
|
||||
shaderCache: {hint: `${components}`, inputDependencies: ['type', 'type']},
|
||||
getRunData: () => ({
|
||||
outputs: [{dims: probsShape, dataType: q.dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: dispatch,
|
||||
programUniforms
|
||||
}),
|
||||
getShaderSource,
|
||||
},
|
||||
{inputs, outputs: [-1]})[0];
|
||||
|
||||
computeInPlaceSoftmax(
|
||||
context, probs, parameters.batchSize * parameters.numHeads * parameters.sequenceLength,
|
||||
parameters.totalSequenceLength);
|
||||
|
||||
return probs;
|
||||
return {
|
||||
name: 'AttentionProbs',
|
||||
shaderCache: {hint: `${components}`, inputDependencies},
|
||||
getRunData: () => ({
|
||||
outputs: [{dims: probsShape, dataType: q.dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: dispatch,
|
||||
programUniforms
|
||||
}),
|
||||
getShaderSource,
|
||||
};
|
||||
};
|
||||
|
||||
const computeVxAttentionScore =
|
||||
(context: ComputeContext, probs: TensorView, v: TensorView, params: AttentionParameters) => {
|
||||
|
||||
const createVxAttentionScoreProgramInfo =
|
||||
(_context: ComputeContext, probs: TensorView, v: TensorView, params: AttentionParameters) => {
|
||||
const outputShape = [params.batchSize, params.sequenceLength, params.vHiddenSize];
|
||||
const TILE_SIZE = 12;
|
||||
const dispatch = {
|
||||
|
|
@ -432,6 +470,7 @@ const computeVxAttentionScore =
|
|||
{type: DataType.uint32, data: params.vHiddenSize}
|
||||
];
|
||||
|
||||
const inputDependencies: ProgramInputTensorInfoDependency[] = ['type', 'type'];
|
||||
const getShaderSource = (shaderHelper: ShaderHelper) => {
|
||||
const probsHelper = inputVariable('probs', probs.dataType, probs.dims);
|
||||
const vHelper = inputVariable('v', v.dataType, v.dims);
|
||||
|
|
@ -482,27 +521,62 @@ const computeVxAttentionScore =
|
|||
}`;
|
||||
};
|
||||
|
||||
return context.compute(
|
||||
{
|
||||
name: 'AttentionScore',
|
||||
shaderCache: {inputDependencies: ['type', 'type']},
|
||||
getRunData: () => ({
|
||||
outputs: [{dims: outputShape, dataType: probs.dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: dispatch,
|
||||
programUniforms
|
||||
}),
|
||||
getShaderSource,
|
||||
},
|
||||
{inputs: [probs, v], outputs: [0]})[0];
|
||||
return {
|
||||
name: 'AttentionScore',
|
||||
shaderCache: {inputDependencies},
|
||||
getRunData: () => ({
|
||||
outputs: [{dims: outputShape, dataType: probs.dataType, gpuDataType: GpuDataType.default}],
|
||||
dispatchGroup: dispatch,
|
||||
programUniforms
|
||||
}),
|
||||
getShaderSource,
|
||||
};
|
||||
};
|
||||
|
||||
export const applyAttention =
|
||||
(context: ComputeContext, q: TensorView, k: TensorView, v: TensorView, _maskIndex: TensorView|undefined,
|
||||
_past: TensorView|undefined, _pastKey: TensorView|undefined, _pastValue: TensorView|undefined,
|
||||
_past: TensorView|undefined, pastKey: TensorView|undefined, pastValue: TensorView|undefined,
|
||||
relativePositionBias: TensorView|undefined, parameters: AttentionParameters, attributes: AttentionAttrs) => {
|
||||
const probs = computeAttentionProbs(context, q, k, relativePositionBias, parameters, attributes);
|
||||
// Concatinate pastKey and K to produce presentKey.
|
||||
const presentKeyShape =
|
||||
[parameters.batchSize, parameters.numHeads, parameters.totalSequenceLength, parameters.headSize];
|
||||
const concatKeyInputs = pastKey ? [pastKey, k] : [k];
|
||||
const key = (context.outputCount > 1 || pastKey) ?
|
||||
context.compute(
|
||||
createConcatProgramInfo(concatKeyInputs, 2, presentKeyShape, k.dataType),
|
||||
{inputs: concatKeyInputs, outputs: [context.outputCount > 1 ? 1 : -1]})[0] :
|
||||
k;
|
||||
|
||||
computeVxAttentionScore(context, probs, v, parameters);
|
||||
// Concatinate pastValue and V to produce presentValue.
|
||||
const presentValueShape =
|
||||
[parameters.batchSize, parameters.numHeads, parameters.totalSequenceLength, parameters.headSize];
|
||||
const concatValueInputs = pastValue ? [pastValue, v] : [v];
|
||||
const value = (context.outputCount > 2 || pastValue) ?
|
||||
context.compute(
|
||||
createConcatProgramInfo(concatValueInputs, 2, presentValueShape, v.dataType),
|
||||
{inputs: concatValueInputs, outputs: [context.outputCount > 2 ? 2 : -1]})[0] :
|
||||
v;
|
||||
const inputsK = [q, key];
|
||||
if (relativePositionBias) {
|
||||
inputsK.push(relativePositionBias);
|
||||
}
|
||||
|
||||
// Run AttentionProbs
|
||||
const probs = context.compute(
|
||||
createAttentionProbsProgramInfo(context, q, key, relativePositionBias, parameters, attributes),
|
||||
{inputs: inputsK, outputs: [-1]})[0];
|
||||
|
||||
// Run Softmax
|
||||
context.compute(
|
||||
createInPlaceSoftmaxProgramInfo(
|
||||
context, probs, parameters.batchSize * parameters.numHeads * parameters.sequenceLength,
|
||||
parameters.totalSequenceLength),
|
||||
{inputs: [probs], outputs: []});
|
||||
|
||||
// Run AttrionScore
|
||||
const inputsV = [probs, value];
|
||||
context.compute(
|
||||
createVxAttentionScoreProgramInfo(context, probs, value, parameters), {inputs: inputsV, outputs: [0]});
|
||||
};
|
||||
|
||||
const prepare = (context: ComputeContext, parameters: AttentionParameters) => {
|
||||
|
|
|
|||
|
|
@ -313,7 +313,7 @@ export const castToF32 = (dataType: string, components: number, value: string) =
|
|||
return `f32(${value})`;
|
||||
}
|
||||
|
||||
return `vec${components}f(${value})`;
|
||||
return `vec${components}f32(${value})`;
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ const assignOutputData = (inputs: readonly IndicesHelper[], output: IndicesHelpe
|
|||
return codeLines.join('\n');
|
||||
};
|
||||
|
||||
const createConcatProgramInfo =
|
||||
export const createConcatProgramInfo =
|
||||
(inputs: readonly TensorView[], adjustedAxis: number, outputShape: number[], dataType: DataType): ProgramInfo => {
|
||||
const outputSize = ShapeUtil.size(outputShape);
|
||||
|
||||
|
|
|
|||
|
|
@ -11,15 +11,18 @@ import {applyAttention, AttentionAttrs, AttentionMaskType, AttentionParameters,
|
|||
import {inputVariable, outputVariable, ShaderHelper, UniformsArrayType} from './common';
|
||||
import {createTransposeProgramInfo, TransposeAttributes} from './transpose';
|
||||
|
||||
const getInput = (inputs: readonly TensorView[], i: number) =>
|
||||
(inputs.length > i) && (inputs[i].dims.length > 0) && (ShapeUtil.size(inputs[i].dims)) > 0 ? inputs[i] : undefined;
|
||||
|
||||
const validateInputs = (inputs: readonly TensorView[], attributes: AttentionAttrs): AttentionParameters => {
|
||||
const query = inputs[0];
|
||||
const key = inputs[1];
|
||||
const value = inputs[2];
|
||||
const bias = inputs[3];
|
||||
const keyPaddingMask = inputs[4];
|
||||
const relativePositionBias = inputs[5];
|
||||
const pastKey = inputs[6];
|
||||
const pastValue = inputs[7];
|
||||
const key = getInput(inputs, 1);
|
||||
const value = getInput(inputs, 2);
|
||||
const bias = getInput(inputs, 3);
|
||||
const keyPaddingMask = getInput(inputs, 4);
|
||||
const relativePositionBias = getInput(inputs, 5);
|
||||
const pastKey = getInput(inputs, 6);
|
||||
const pastValue = getInput(inputs, 7);
|
||||
|
||||
// Abbreviation and Meanings:
|
||||
// B: batch_size
|
||||
|
|
@ -73,6 +76,16 @@ const validateInputs = (inputs: readonly TensorView[], attributes: AttentionAttr
|
|||
if (pastKey.dims.length !== 4) {
|
||||
throw new Error('Input "past_key" is expected to have 4 dimensions');
|
||||
}
|
||||
if (pastKey.dims[0] !== batchSize || pastKey.dims[1] !== attributes.numHeads || pastKey.dims[3] !== headSize) {
|
||||
throw new Error('Input "past_key" shape (batch_size, num_heads, past_sequence_length, head_size)');
|
||||
}
|
||||
if (pastValue.dims[0] !== batchSize || pastValue.dims[1] !== attributes.numHeads ||
|
||||
pastValue.dims[3] !== headSize) {
|
||||
throw new Error('Input "past_value" shape (batch_size, num_heads, past_sequence_length, head_size)');
|
||||
}
|
||||
if (pastKey.dims[2] !== pastValue.dims[2]) {
|
||||
throw new Error('Input "past_key" and "past_value" shall have same dim 2 (past_sequence_length)');
|
||||
}
|
||||
if (pastValue.dims.length !== 4) {
|
||||
throw new Error('Input "past_value" is expected to have 4 dimensions');
|
||||
}
|
||||
|
|
@ -186,23 +199,20 @@ const validateInputs = (inputs: readonly TensorView[], attributes: AttentionAttr
|
|||
|
||||
const totalSequenceLength = pastSequenceLength + kvSequenceLength;
|
||||
const broadcastResPosBias = false;
|
||||
// if (extraAddQk) {
|
||||
// if (extraAddQk.dims[0] === 1) {
|
||||
// broadcastResPosBias = true;
|
||||
// }
|
||||
// }
|
||||
|
||||
if (keyPaddingMask) {
|
||||
throw new Error('Key padding mask is not supported');
|
||||
}
|
||||
|
||||
if (relativePositionBias) {
|
||||
throw new Error('extraAddQk is not supported');
|
||||
}
|
||||
if (pastKey) {
|
||||
throw new Error('pastKey is not supported');
|
||||
}
|
||||
if (pastValue) {
|
||||
throw new Error('pastValue is not supported');
|
||||
if (relativePositionBias.dims.length !== 4) {
|
||||
throw new Error('Input "relative_position_bias" is expected to have 4 dimensions');
|
||||
}
|
||||
if ((relativePositionBias.dims[0] !== batchSize && relativePositionBias.dims[0] !== 1) ||
|
||||
relativePositionBias.dims[1] !== attributes.numHeads || relativePositionBias.dims[2] !== sequenceLength ||
|
||||
relativePositionBias.dims[3] !== totalSequenceLength) {
|
||||
throw new Error('Input "relative_position_bias" shape (batch_size, 1, sequence_length, kv_sequence_length)');
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
|
|
@ -305,38 +315,44 @@ const maybeTransposeToBNSHAndAddBias =
|
|||
|
||||
export const multiHeadAttention = (context: ComputeContext, attributes: AttentionAttrs): void => {
|
||||
const params = validateInputs(context.inputs, attributes);
|
||||
|
||||
if (context.inputs[0].dims.length === 5) {
|
||||
const query = context.inputs[0];
|
||||
const key = getInput(context.inputs, 1);
|
||||
const value = getInput(context.inputs, 2);
|
||||
const bias = getInput(context.inputs, 3);
|
||||
const keyPaddingMask = getInput(context.inputs, 4);
|
||||
const relativePositionBias = getInput(context.inputs, 5);
|
||||
const pastKey = getInput(context.inputs, 6);
|
||||
const pastValue = getInput(context.inputs, 7);
|
||||
if (query.dims.length === 5) {
|
||||
throw new Error('Packed QKV is not implemented');
|
||||
}
|
||||
|
||||
if (context.inputs[1]?.dims.length === 5) {
|
||||
if (key?.dims.length === 5) {
|
||||
throw new Error('Packed KV is not implemented');
|
||||
}
|
||||
|
||||
// applyAttention expects BNSH inputs
|
||||
const kvBNSH = context.inputs[1] && context.inputs[2] && context.inputs[1].dims.length === 4 &&
|
||||
context.inputs[2].dims.length === 4;
|
||||
const kvBNSH = key && value && key.dims.length === 4 && value.dims.length === 4;
|
||||
|
||||
const Q = maybeTransposeToBNSHAndAddBias(
|
||||
context, params.batchSize, params.numHeads, params.sequenceLength, params.headSize, context.inputs[0],
|
||||
context.inputs[3], 0);
|
||||
context, params.batchSize, params.numHeads, params.sequenceLength, params.headSize, query, bias, 0);
|
||||
|
||||
if (kvBNSH) {
|
||||
return applyAttention(
|
||||
context, Q, context.inputs[1], context.inputs[2], context.inputs[4], undefined, undefined, undefined,
|
||||
context.inputs[5], params, attributes);
|
||||
context, Q, key, value, keyPaddingMask, undefined, undefined, undefined, relativePositionBias, params,
|
||||
attributes);
|
||||
}
|
||||
if (!key || !value) {
|
||||
throw new Error('key and value must be provided');
|
||||
}
|
||||
|
||||
const K = maybeTransposeToBNSHAndAddBias(
|
||||
context, params.batchSize, params.numHeads, params.kvSequenceLength, params.headSize, context.inputs[1],
|
||||
context.inputs[3], params.hiddenSize);
|
||||
context, params.batchSize, params.numHeads, params.kvSequenceLength, params.headSize, key, bias,
|
||||
params.hiddenSize);
|
||||
|
||||
const V = maybeTransposeToBNSHAndAddBias(
|
||||
context, params.batchSize, params.numHeads, params.kvSequenceLength, params.vHeadSize, context.inputs[2],
|
||||
context.inputs[3], 2 * params.hiddenSize);
|
||||
context, params.batchSize, params.numHeads, params.kvSequenceLength, params.vHeadSize, value, bias,
|
||||
2 * params.hiddenSize);
|
||||
|
||||
applyAttention(
|
||||
context, Q, K, V, context.inputs[4], undefined, context.inputs[6], context.inputs[7], context.inputs[5], params,
|
||||
attributes);
|
||||
context, Q, K, V, keyPaddingMask, undefined, pastKey, pastValue, relativePositionBias, params, attributes);
|
||||
};
|
||||
|
|
@ -1369,7 +1369,7 @@
|
|||
"matmul-broadcast.jsonc",
|
||||
"mul.jsonc",
|
||||
"mul_int32.jsonc",
|
||||
"multi-head-attention.jsonc",
|
||||
"multihead-attention.jsonc",
|
||||
//"neg.jsonc",
|
||||
"neg-int32.jsonc",
|
||||
"not.jsonc",
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include "multi_head_attention.h"
|
||||
#include "multihead_attention.h"
|
||||
#include "core/providers/js/js_data_types.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
Loading…
Reference in a new issue