Skip to content

fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to prevent overflow (#26732) #29599

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

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

fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to prevent overflow (#26732) #29599

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4d83b9b
fix(webgpu): use f32 accumulators in fp16 MatMul/MatMulNBits to preve…
RobertoReale Jul 7, 2026
2f2a4b3
Address review feedback: extend f32 accumulation to fused native kern…
RobertoReale Jul 8, 2026
53c9320
Extend f32 accumulation to the generic dp4a_matmul kernel
RobertoReale Jul 8, 2026
2cf1e37
Fix vec4 constructor type mismatch in dp4a qualcomm store path
RobertoReale Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 27 additions & 15 deletions js/web/lib/wasm/jsep/webgpu/ops/3rd-party/matmul_packed_webgpu.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,19 +67,23 @@ const calculateResultSnippet = (transposeA: boolean, innerElementSize: number) =
let ACached2 = mm_Asub[k * innerElementSize + 2][localRow];
${innerElementSize === 3 ? '' : 'let ACached3 = mm_Asub[k * innerElementSize + 3][localRow];'}
for (var i = 0; i < rowPerThread; i = i + 1) {
acc[i] = BCached0 * ACached0[i] + acc[i];
acc[i] = BCached1 * ACached1[i] + acc[i];
acc[i] = BCached2 * ACached2[i] + acc[i];
${innerElementSize === 3 ? '' : 'acc[i] = BCached3 * ACached3[i] + acc[i];'}
// Explicit f32 casts on each operand are required: Dawn/D3D12 re-demotes
// temporaries to f16 when 'enable f16;' is active (issue #26732).
acc[i] = vec4<f32>(BCached0) * f32(ACached0[i]) + acc[i];
acc[i] = vec4<f32>(BCached1) * f32(ACached1[i]) + acc[i];
acc[i] = vec4<f32>(BCached2) * f32(ACached2[i]) + acc[i];
${innerElementSize === 3 ? '' : 'acc[i] = vec4<f32>(BCached3) * f32(ACached3[i]) + acc[i];'}
}`;
} else {
return `
for (var i = 0; i < rowPerThread; i = i + 1) {
let ACached = mm_Asub[tileRow + i][k];
acc[i] = BCached0 * ACached.x + acc[i];
acc[i] = BCached1 * ACached.y + acc[i];
acc[i] = BCached2 * ACached.z + acc[i];
${innerElementSize === 3 ? '' : 'acc[i] = BCached3 * ACached.w + acc[i];'}
// Explicit f32 casts on each operand are required: Dawn/D3D12 re-demotes
// temporaries to f16 when 'enable f16;' is active (issue #26732).
acc[i] = vec4<f32>(BCached0) * f32(ACached.x) + acc[i];
acc[i] = vec4<f32>(BCached1) * f32(ACached.y) + acc[i];
acc[i] = vec4<f32>(BCached2) * f32(ACached.z) + acc[i];
${innerElementSize === 3 ? '' : 'acc[i] = vec4<f32>(BCached3) * f32(ACached.w) + acc[i];'}
}`;
}
};
Expand Down Expand Up @@ -140,7 +144,9 @@ fn main(@builtin(local_invocation_id) localId : vec3<u32>,
let num_tiles = ${splitK ? `${Math.ceil(splitedDimInner / tileInner)}` : '(uniforms.dim_inner - 1) / tileInner + 1'};
var kStart = ${splitK ? `i32(globalId.z) * ${splitedDimInner}` : '0'};

var acc: array<vec4<${type}>, rowPerThread>;
// Accumulate in f32 to prevent fp16 overflow in long dot products (issue #26732).
// Tiles (mm_Asub/mm_Bsub) stay in ${type}: no shared-memory or bandwidth regression.
var acc: array<vec4<f32>, rowPerThread>;

// Loop over shared dimension.
let tileRowB = localRow * ${rowPerThreadB};
Expand Down Expand Up @@ -177,7 +183,7 @@ fn main(@builtin(local_invocation_id) localId : vec3<u32>,
}

for (var innerRow = 0; innerRow < rowPerThread; innerRow = innerRow + 1) {
mm_write(batch, globalRow + innerRow, globalCol, acc[innerRow]);
mm_write(batch, globalRow + innerRow, globalCol, vec4<${type}>(acc[innerRow]));
}
}`;
};
Expand Down Expand Up @@ -268,8 +274,10 @@ export const makeMatMulPackedSource = (
: `mm_Asub[localRow + innerRow * ${workgroupSize[1]}][k];`
}
for (var innerCol = 0; innerCol < colPerThread; innerCol = innerCol + 1) {
// Explicit f32 casts required: Dawn/D3D12 re-demotes temporaries to f16
// when 'enable f16;' is active (issue #26732).
acc[innerRow][innerCol] = acc[innerRow][innerCol] +
ACached * BCached[innerCol];
f32(ACached) * f32(BCached[innerCol]);
}
}
}
Expand All @@ -279,7 +287,7 @@ export const makeMatMulPackedSource = (
let gRow = globalRowStart + localRow + innerRow * ${workgroupSize[1]};
for (var innerCol = 0; innerCol < colPerThread; innerCol = innerCol + 1) {
let gCol = globalColStart + localCol + innerCol * ${workgroupSize[0]};
mm_write(batch, gRow, gCol, acc[innerRow][innerCol]);
mm_write(batch, gRow, gCol, ${type}(acc[innerRow][innerCol]));
}
}
`
Expand Down Expand Up @@ -328,7 +336,9 @@ for (var t = 0; t < num_tiles; t = t + 1) {
for (var innerRow = 0; innerRow < rowPerThread; innerRow = innerRow + 1) {
${readDataFromSubASnippet(transposeA)}
for (var innerCol = 0; innerCol < colPerThread; innerCol = innerCol + 1) {
acc[innerRow][innerCol] = acc[innerRow][innerCol] + ACached * BCached[innerCol];
// Explicit f32 casts required: Dawn/D3D12 re-demotes temporaries to f16
// when 'enable f16;' is active (issue #26732).
acc[innerRow][innerCol] = acc[innerRow][innerCol] + f32(ACached) * f32(BCached[innerCol]);
}
}
}
Expand All @@ -339,7 +349,7 @@ for (var t = 0; t < num_tiles; t = t + 1) {
for (var innerRow = 0; innerRow < rowPerThread; innerRow = innerRow + 1) {
for (var innerCol = 0; innerCol < colPerThread; innerCol = innerCol + 1) {
mm_write(batch, globalRow + innerRow, globalCol + innerCol,
acc[innerRow][innerCol]);
${type}(acc[innerRow][innerCol]));
}
}
`;
Expand All @@ -362,7 +372,9 @@ fn main(@builtin(local_invocation_id) localId : vec3<u32>,
};
var kStart = ${splitK ? `i32(globalId.z) * ${splitedDimInner}` : '0'};

var acc : array<array<${type}, colPerThread>, rowPerThread>;
// Accumulate in f32 to prevent fp16 overflow in long dot products (issue #26732).
// Tiles (mm_Asub/mm_Bsub) stay in ${type}: no shared-memory or bandwidth regression.
var acc : array<array<f32, colPerThread>, rowPerThread>;
${matmulSnippet}
}
`;
Expand Down
11 changes: 8 additions & 3 deletions js/web/lib/wasm/jsep/webgpu/ops/matmul-shaders.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,10 @@ export const createNaiveMatmulProgramInfo = (
];
appendActivationUniforms(activationAttributes, uniforms);

// Accumulate in f32 to prevent fp16 overflow in long dot products (issue #26732).
// Explicit f32 casts on each operand are required: Dawn/D3D12 re-demotes
// temporaries to f16 when 'enable f16;' is active.
const accType = components === 1 ? 'f32' : `vec${components}<f32>`;
const calcResult = (): string => {
let calcStr = `var a_data: ${a.type.value};`;
for (let i = 0; i < aComponents; i++) {
Expand All @@ -123,7 +127,7 @@ export const createNaiveMatmulProgramInfo = (

for (let j = 0; j < aComponents; j++) {
calcStr += `
values[${i}] = fma(${b.type.value}(a_data${aComponents === 1 ? '' : `[${j}]`}), b_data${j}, values[${i}]);\n`;
values[${i}] = fma(${accType}(a_data${aComponents === 1 ? '' : `[${j}]`}), ${accType}(b_data${j}), values[${i}]);\n`;
}
}
return calcStr;
Expand Down Expand Up @@ -155,12 +159,13 @@ export const createNaiveMatmulProgramInfo = (
${b.indicesSet('b_indices', b.rank - 2, 0)}
${b.indicesSet('b_indices', b.rank - 1, 0)}
let b_offset = ${b.indicesToOffset('b_indices')};
var values: array<${output.type.value}, ${outputNumber}>;
var values: array<${accType}, ${outputNumber}>;
for (var k: u32 = 0u; k < uniforms.K; k = k + ${aComponents}) {
${calcResult()}
}
for (var i = 0u; i < ${outputNumber}u; i++) {
var value = values[i];
// Downcast to the output type only at the final write.
var value = ${output.type.value}(values[i]);
${processBias}
${applyActivation}
let cur_indices = ${output.type.indices}(batch, row + i, col);
Expand Down
26 changes: 17 additions & 9 deletions js/web/lib/wasm/jsep/webgpu/ops/matmulnbits.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,8 +180,8 @@ export const createMatMulNBitsProgramInfo = (
(_, i) =>
`${
aComponents === 1
? `a_data${pass > 0 ? pass : ''}[${i}] * b_dequantized_values[${i}]`
: `dot(a_data${pass > 0 ? pass : ''}[${i}], b_dequantized_values[${i}])`
? `f32(a_data${pass > 0 ? pass : ''}[${i}]) * f32(b_dequantized_values[${i}])`
: `dot(vec${aComponents}<f32>(a_data${pass > 0 ? pass : ''}[${i}]), vec${aComponents}<f32>(b_dequantized_values[${i}]))`
}`,
).join(' + ')};
`;
Expand Down Expand Up @@ -242,8 +242,13 @@ export const createMatMulNBitsProgramInfo = (
var b_dequantized_values: ${qDqDataType};`;
return calcStr;
};
// Accumulate in f32 to prevent fp16 overflow in long dot products (issue #26732).
// Weights and activations stay in their original dtype; explicit f32 casts on each
// operand are required because Dawn/D3D12 re-demotes temporaries to f16 when
// 'enable f16;' is active. The result is downcast to the output type only at the write.
const accType = components === 1 ? 'f32' : `vec${components}<f32>`;
return `
var<workgroup> workgroup_shared: array<${output.type.value}, ${outputNumber * workgroupSize}>;
var<workgroup> workgroup_shared: array<${accType}, ${outputNumber * workgroupSize}>;
${shaderHelper.declareVariables(...inputVariables, output)}
${shaderHelper.mainStart([workgroupSize, 1, 1])}
let output_indices = ${output.offsetToIndices(`(global_idx / ${workgroupSize}) * ${outputNumber}`)};
Expand All @@ -267,13 +272,13 @@ export const createMatMulNBitsProgramInfo = (
workgroupBarrier();

if (local_id.x < ${outputNumber}) {
var output_value: ${output.type.value} = ${output.type.value}(0);
var output_value: ${accType} = ${accType}(0);
var workgroup_shared_offset: u32 = local_id.x;
for (var b: u32 = 0u; b < ${workgroupSize}u; b++) {
output_value += workgroup_shared[workgroup_shared_offset];
workgroup_shared_offset += ${outputNumber};
}
${output.setByIndices(`${output.type.indices}(batch, row, col + local_id.x)`, 'output_value')};
${output.setByIndices(`${output.type.indices}(batch, row, col + local_id.x)`, `${output.type.value}(output_value)`)};
}
}`;
};
Expand Down Expand Up @@ -368,7 +373,8 @@ export const createMatMulNBitsBlockSize32ProgramInfo = (

return `
var<workgroup> sub_a: array<${a.type.value}, ${aLengthPerTile}>;
var<workgroup> inter_results: array<array<${output.type.value}, ${workgroupX}>, ${workgroupY}>;
// Accumulate in f32 to prevent fp16 overflow in long dot products (issue #26732).
var<workgroup> inter_results: array<array<f32, ${workgroupX}>, ${workgroupY}>;
${shaderHelper.declareVariables(...inputVariables, output)}
${shaderHelper.mainStart([workgroupX, workgroupY, 1])}
let output_indices = ${output.offsetToIndices(`workgroup_index * ${workgroupY}`)};
Expand Down Expand Up @@ -444,9 +450,11 @@ export const createMatMulNBitsBlockSize32ProgramInfo = (
(_, i) => `${dataType}(b_value_lower[${i}]), ${dataType}(b_value_upper[${i}])`,
).join(', ')});
let b_dequantized_values = (b_quantized_values - mat2x4<${dataType}>(${Array(8).fill('zero_point').join(',')})) * scale;
// Explicit f32 casts on each operand are required: Dawn/D3D12 re-demotes
// temporaries to f16 when 'enable f16;' is active (issue #26732).
inter_results[local_id.y][local_id.x] += ${Array.from(
{ length: 2 },
(_, i) => `${`dot(a_data${i}, b_dequantized_values[${i}])`}`,
(_, i) => `${`dot(vec4<f32>(a_data${i}), vec4<f32>(b_dequantized_values[${i}]))`}`,
).join(' + ')};
}
word_offset += ${8 / aComponents};`;
Expand All @@ -458,13 +466,13 @@ export const createMatMulNBitsBlockSize32ProgramInfo = (
}

if (local_idx < ${workgroupY}) {
var output_value: ${output.type.value} = ${output.type.value}(0);
var output_value: f32 = f32(0);
for (var b = 0u; b < ${workgroupX}; b++) {
output_value += inter_results[local_idx][b];
}
if (col + local_idx < uniforms.output_shape[2])
{
${output.setByIndices(`${output.type.indices}(batch, row, col + local_idx)`, 'output_value')}
${output.setByIndices(`${output.type.indices}(batch, row, col + local_idx)`, `${output.type.value}(output_value)`)}
}
}
}`;
Expand Down
Loading