BiasSplitGelu fix for batch size > 1

dakenf · Aug 31, 2023 · d3cbe9e · d3cbe9e
1 parent 92aeb75
commit d3cbe9e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/js/web/lib/wasm/jsep/webgpu/ops/bias-split-gelu.ts b/js/web/lib/wasm/jsep/webgpu/ops/bias-split-gelu.ts
@@ -52,8 +52,10 @@ const createBiasSplitGeluProgramInfo = (metadata: ProgramMetadata, inputs: reado
   ${shaderHelper.mainStart()}
     ${shaderHelper.guardAgainstOutOfBoundsWorkgroupSizes(outputSize)}
     let biasIdx = global_idx % halfChannels;
-    let valueLeft = input[global_idx] + bias[biasIdx];
-    let valueRight = input[global_idx + halfChannels] + bias[biasIdx + halfChannels];
+    let batchIndex = global_idx / halfChannels;
+    let inputOffset = biasIdx + batchIndex * halfChannels * 2;
+    let valueLeft = input[inputOffset] + bias[biasIdx];
+    let valueRight = input[inputOffset + halfChannels] + bias[biasIdx + halfChannels];
     let geluRight = valueRight * 0.5 * (erf_vf32(valueRight / M_SQRT2) + 1);
 
     ${output.setByOffset('global_idx', 'valueLeft * geluRight')}