@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/backward/attention_backward.js

@@ -1,6 +1,5 @@
-import { CommandRecorder } from '../../command-recorder.js';
 import { getDevice } from '../../device.js';
-import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../../tensor.js';
 import { castF16ToF32, recordCastF16ToF32 } from '../cast.js';
 import { runMatmul, recordMatmul } from '../matmul.js';
@@ -15,24 +14,16 @@ async function ensureF32(tensor, recorder = null) {
   if (!recorder) {
     return castF16ToF32(tensor);
   }
-  const casted = await recordCastF16ToF32(recorder, tensor);
-  recorder.trackTemporaryBuffer(casted.buffer);
-  return casted;
+  return recordCastF16ToF32(recorder, tensor);
 }
 
-function createHeadSliceBuffers(recorder, headBytes, softmaxBytes) {
+function createHeadSliceBuffers(headBytes, softmaxBytes) {
   const qHeadBuf = acquireBuffer(headBytes, undefined, 'attn_q_head');
   const kHeadBuf = acquireBuffer(headBytes, undefined, 'attn_k_head');
   const vHeadBuf = acquireBuffer(headBytes, undefined, 'attn_v_head');
   const sHeadBuf = acquireBuffer(softmaxBytes, undefined, 'attn_s_head');
   const dHeadBuf = acquireBuffer(headBytes, undefined, 'attn_d_head');
 
-  recorder.trackTemporaryBuffer(qHeadBuf);
-  recorder.trackTemporaryBuffer(kHeadBuf);
-  recorder.trackTemporaryBuffer(vHeadBuf);
-  recorder.trackTemporaryBuffer(sHeadBuf);
-  recorder.trackTemporaryBuffer(dHeadBuf);
-
   return { qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf };
 }
 
@@ -49,6 +40,19 @@ function trackTensorBuffer(recorder, tensor) {
   recorder.trackTemporaryBuffer(tensor.buffer);
 }
 
+function releaseTensorBuffer(tensor) {
+  if (tensor?.buffer) {
+    releaseBuffer(tensor.buffer);
+  }
+}
+
+function maybeTrackOwnedTensor(ownedTensors, originalTensor, resolvedTensor) {
+  if (resolvedTensor !== originalTensor) {
+    ownedTensors.push(resolvedTensor);
+  }
+  return resolvedTensor;
+}
+
 async function runAttentionBackwardCore(
   q,
   k,
@@ -63,11 +67,23 @@ async function runAttentionBackwardCore(
     throw new Error('attention backward requires seqLen, numHeads, and headDim');
   }
 
-  const qTensor = await ensureF32(q, recorder);
-  const kTensor = await ensureF32(k, recorder);
-  const vTensor = await ensureF32(v, recorder);
-  const sTensor = await ensureF32(softmax, recorder);
-  const dTensor = await ensureF32(gradOutput, recorder);
+  const ownedInputTensors = [];
+  const ownedRecorderInputTensors = [];
+  const qTensor = !recorder
+    ? maybeTrackOwnedTensor(ownedInputTensors, q, await ensureF32(q))
+    : maybeTrackOwnedTensor(ownedRecorderInputTensors, q, await ensureF32(q, recorder));
+  const kTensor = !recorder
+    ? maybeTrackOwnedTensor(ownedInputTensors, k, await ensureF32(k))
+    : maybeTrackOwnedTensor(ownedRecorderInputTensors, k, await ensureF32(k, recorder));
+  const vTensor = !recorder
+    ? maybeTrackOwnedTensor(ownedInputTensors, v, await ensureF32(v))
+    : maybeTrackOwnedTensor(ownedRecorderInputTensors, v, await ensureF32(v, recorder));
+  const sTensor = !recorder
+    ? maybeTrackOwnedTensor(ownedInputTensors, softmax, await ensureF32(softmax))
+    : maybeTrackOwnedTensor(ownedRecorderInputTensors, softmax, await ensureF32(softmax, recorder));
+  const dTensor = !recorder
+    ? maybeTrackOwnedTensor(ownedInputTensors, gradOutput, await ensureF32(gradOutput))
+    : maybeTrackOwnedTensor(ownedRecorderInputTensors, gradOutput, await ensureF32(gradOutput, recorder));
 
   const headElements = seqLen * headDim;
   const headBytes = headElements * dtypeBytes(qTensor.dtype);
@@ -77,171 +93,247 @@ async function runAttentionBackwardCore(
   const gradQBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_q');
   const gradKBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_k');
   const gradVBuf = acquireBuffer(totalBytes, undefined, 'attn_grad_v');
+  let completed = false;
 
-  if (!recorder) {
-    for (let h = 0; h < numHeads; h += 1) {
-      const qOffset = h * headBytes;
-      const kOffset = h * headBytes;
-      const vOffset = h * headBytes;
-      const dOffset = h * headBytes;
-      const sOffset = h * softmaxBytes;
+  try {
+    if (!recorder) {
+      for (let h = 0; h < numHeads; h += 1) {
+        const qOffset = h * headBytes;
+        const kOffset = h * headBytes;
+        const vOffset = h * headBytes;
+        const dOffset = h * headBytes;
+        const sOffset = h * softmaxBytes;
 
-      const qHeadBuf = acquireBuffer(headBytes, undefined, 'attn_q_head');
-      const kHeadBuf = acquireBuffer(headBytes, undefined, 'attn_k_head');
-      const vHeadBuf = acquireBuffer(headBytes, undefined, 'attn_v_head');
-      const sHeadBuf = acquireBuffer(softmaxBytes, undefined, 'attn_s_head');
-      const dHeadBuf = acquireBuffer(headBytes, undefined, 'attn_d_head');
+        const qHeadBuf = acquireBuffer(headBytes, undefined, 'attn_q_head');
+        const kHeadBuf = acquireBuffer(headBytes, undefined, 'attn_k_head');
+        const vHeadBuf = acquireBuffer(headBytes, undefined, 'attn_v_head');
+        const sHeadBuf = acquireBuffer(softmaxBytes, undefined, 'attn_s_head');
+        const dHeadBuf = acquireBuffer(headBytes, undefined, 'attn_d_head');
+        let sTransposed = null;
+        let dV = null;
+        let vTransposed = null;
+        let dS = null;
+        let dQK = null;
+        let dQ = null;
+        let dQKTransposed = null;
+        let dK = null;
 
-      const sliceEncoder = getDevice().createCommandEncoder();
-      sliceEncoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
-      sliceEncoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
-      sliceEncoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
-      sliceEncoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
-      sliceEncoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
-      getDevice().queue.submit([sliceEncoder.finish()]);
+        try {
+          const sliceEncoder = getDevice().createCommandEncoder();
+          sliceEncoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
+          sliceEncoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
+          sliceEncoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
+          sliceEncoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
+          sliceEncoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
+          getDevice().queue.submit([sliceEncoder.finish()]);
 
-      const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
-        qHeadBuf,
-        kHeadBuf,
-        vHeadBuf,
-        sHeadBuf,
-        dHeadBuf,
-        seqLen,
-        headDim
-      );
+          const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
+            qHeadBuf,
+            kHeadBuf,
+            vHeadBuf,
+            sHeadBuf,
+            dHeadBuf,
+            seqLen,
+            headDim
+          );
 
-      const sTransposed = await runTranspose(sHead, seqLen, seqLen);
-      const dV = await runMatmul(sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        bDtype: 'f32',
-      });
+          sTransposed = await runTranspose(sHead, seqLen, seqLen);
+          dV = await runMatmul(sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            bDtype: 'f32',
+          });
 
-      const vTransposed = await runTranspose(vHead, seqLen, headDim);
-      const dS = await runMatmul(dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
-        transposeB: false,
-        bDtype: 'f32',
-      });
-      const dQK = causal
-        ? await runBackwardKernel(
-          'attention_backward',
-          sHead,
-          dS,
-          16,
-          (view) => {
-            view.setUint32(0, seqLen, true);
-            view.setUint32(4, seqLen, true);
-            view.setUint32(8, 1, true);
-          }
-        )
-        : await runSoftmaxBackward(sHead, dS, { rows: seqLen, cols: seqLen });
+          vTransposed = await runTranspose(vHead, seqLen, headDim);
+          dS = await runMatmul(dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
+            transposeB: false,
+            bDtype: 'f32',
+          });
+          dQK = causal
+            ? await runBackwardKernel(
+              'attention_backward',
+              sHead,
+              dS,
+              16,
+              (view) => {
+                view.setUint32(0, seqLen, true);
+                view.setUint32(4, seqLen, true);
+                view.setUint32(8, 1, true);
+              }
+            )
+            : await runSoftmaxBackward(sHead, dS, { rows: seqLen, cols: seqLen });
 
-      const dQ = await runMatmul(dQK, kHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        alpha: scale,
-        bDtype: 'f32',
-      });
-      const dQKTransposed = await runTranspose(dQK, seqLen, seqLen);
-      const dK = await runMatmul(dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        alpha: scale,
-        bDtype: 'f32',
-      });
+          dQ = await runMatmul(dQK, kHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            alpha: scale,
+            bDtype: 'f32',
+          });
+          dQKTransposed = await runTranspose(dQK, seqLen, seqLen);
+          dK = await runMatmul(dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            alpha: scale,
+            bDtype: 'f32',
+          });
 
-      const copyEncoder = getDevice().createCommandEncoder();
-      copyEncoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
-      copyEncoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
-      copyEncoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
-      getDevice().queue.submit([copyEncoder.finish()]);
-    }
-  } else {
-    const encoder = recorder.getEncoder();
-    for (let h = 0; h < numHeads; h += 1) {
-      const qOffset = h * headBytes;
-      const kOffset = h * headBytes;
-      const vOffset = h * headBytes;
-      const dOffset = h * headBytes;
-      const sOffset = h * softmaxBytes;
+          const copyEncoder = getDevice().createCommandEncoder();
+          copyEncoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
+          copyEncoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
+          copyEncoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
+          getDevice().queue.submit([copyEncoder.finish()]);
+          await getDevice().queue.onSubmittedWorkDone();
+        } finally {
+          releaseTensorBuffer(sTransposed);
+          releaseTensorBuffer(dV);
+          releaseTensorBuffer(vTransposed);
+          releaseTensorBuffer(dS);
+          releaseTensorBuffer(dQK);
+          releaseTensorBuffer(dQ);
+          releaseTensorBuffer(dQKTransposed);
+          releaseTensorBuffer(dK);
+          releaseBuffer(qHeadBuf);
+          releaseBuffer(kHeadBuf);
+          releaseBuffer(vHeadBuf);
+          releaseBuffer(sHeadBuf);
+          releaseBuffer(dHeadBuf);
+        }
+      }
+    } else {
+      const encoder = recorder.getEncoder();
+      for (let h = 0; h < numHeads; h += 1) {
+        const qOffset = h * headBytes;
+        const kOffset = h * headBytes;
+        const vOffset = h * headBytes;
+        const dOffset = h * headBytes;
+        const sOffset = h * softmaxBytes;
 
-      const { qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf } = createHeadSliceBuffers(
-        recorder,
-        headBytes,
-        softmaxBytes
-      );
+        const { qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf } = createHeadSliceBuffers(
+          headBytes,
+          softmaxBytes
+        );
+        const headBuffers = [qHeadBuf, kHeadBuf, vHeadBuf, sHeadBuf, dHeadBuf];
+        let sTransposed = null;
+        let dV = null;
+        let vTransposed = null;
+        let dS = null;
+        let dQK = null;
+        let dQ = null;
+        let dQKTransposed = null;
+        let dK = null;
 
-      encoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
-      encoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
-      encoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
-      encoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
-      encoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
+        try {
+          encoder.copyBufferToBuffer(qTensor.buffer, qOffset, qHeadBuf, 0, headBytes);
+          encoder.copyBufferToBuffer(kTensor.buffer, kOffset, kHeadBuf, 0, headBytes);
+          encoder.copyBufferToBuffer(vTensor.buffer, vOffset, vHeadBuf, 0, headBytes);
+          encoder.copyBufferToBuffer(sTensor.buffer, sOffset, sHeadBuf, 0, softmaxBytes);
+          encoder.copyBufferToBuffer(dTensor.buffer, dOffset, dHeadBuf, 0, headBytes);
 
-      const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
-        qHeadBuf,
-        kHeadBuf,
-        vHeadBuf,
-        sHeadBuf,
-        dHeadBuf,
-        seqLen,
-        headDim
-      );
+          const { qHead, kHead, vHead, sHead, dHead } = createHeadTensors(
+            qHeadBuf,
+            kHeadBuf,
+            vHeadBuf,
+            sHeadBuf,
+            dHeadBuf,
+            seqLen,
+            headDim
+          );
 
-      const sTransposed = await recordTranspose(recorder, sHead, seqLen, seqLen);
-      const dV = await recordMatmul(recorder, sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        bDtype: 'f32',
-      });
+          sTransposed = await recordTranspose(recorder, sHead, seqLen, seqLen);
+          dV = await recordMatmul(recorder, sTransposed, dHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            bDtype: 'f32',
+          });
 
-      const vTransposed = await recordTranspose(recorder, vHead, seqLen, headDim);
-      const dS = await recordMatmul(recorder, dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
-        transposeB: false,
-        bDtype: 'f32',
-      });
-      const dQK = causal
-        ? await recordBackwardKernel(
-          recorder,
-          'attention_backward',
-          sHead,
-          dS,
-          16,
-          (view) => {
-            view.setUint32(0, seqLen, true);
-            view.setUint32(4, seqLen, true);
-            view.setUint32(8, 1, true);
-          }
-        )
-        : await recordSoftmaxBackward(recorder, sHead, dS, { rows: seqLen, cols: seqLen });
+          vTransposed = await recordTranspose(recorder, vHead, seqLen, headDim);
+          dS = await recordMatmul(recorder, dHead, vTransposed.buffer, seqLen, seqLen, headDim, {
+            transposeB: false,
+            bDtype: 'f32',
+          });
+          dQK = causal
+            ? await recordBackwardKernel(
+              recorder,
+              'attention_backward',
+              sHead,
+              dS,
+              16,
+              (view) => {
+                view.setUint32(0, seqLen, true);
+                view.setUint32(4, seqLen, true);
+                view.setUint32(8, 1, true);
+              }
+            )
+            : await recordSoftmaxBackward(recorder, sHead, dS, { rows: seqLen, cols: seqLen });
 
-      const dQ = await recordMatmul(recorder, dQK, kHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        alpha: scale,
-        bDtype: 'f32',
-      });
-      const dQKTransposed = await recordTranspose(recorder, dQK, seqLen, seqLen);
-      const dK = await recordMatmul(recorder, dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
-        transposeB: false,
-        alpha: scale,
-        bDtype: 'f32',
-      });
+          dQ = await recordMatmul(recorder, dQK, kHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            alpha: scale,
+            bDtype: 'f32',
+          });
+          dQKTransposed = await recordTranspose(recorder, dQK, seqLen, seqLen);
+          dK = await recordMatmul(recorder, dQKTransposed, qHead.buffer, seqLen, headDim, seqLen, {
+            transposeB: false,
+            alpha: scale,
+            bDtype: 'f32',
+          });
 
-      encoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
-      encoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
-      encoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
+          encoder.copyBufferToBuffer(dQ.buffer, 0, gradQBuf, qOffset, headBytes);
+          encoder.copyBufferToBuffer(dK.buffer, 0, gradKBuf, kOffset, headBytes);
+          encoder.copyBufferToBuffer(dV.buffer, 0, gradVBuf, vOffset, headBytes);
+        } catch (error) {
+          releaseTensorBuffer(sTransposed);
+          releaseTensorBuffer(dV);
+          releaseTensorBuffer(vTransposed);
+          releaseTensorBuffer(dS);
+          releaseTensorBuffer(dQK);
+          releaseTensorBuffer(dQ);
+          releaseTensorBuffer(dQKTransposed);
+          releaseTensorBuffer(dK);
+          releaseBuffer(qHeadBuf);
+          releaseBuffer(kHeadBuf);
+          releaseBuffer(vHeadBuf);
+          releaseBuffer(sHeadBuf);
+          releaseBuffer(dHeadBuf);
+          throw error;
+        }
 
-      trackTensorBuffer(recorder, sTransposed);
-      trackTensorBuffer(recorder, dV);
-      trackTensorBuffer(recorder, vTransposed);
-      trackTensorBuffer(recorder, dS);
-      trackTensorBuffer(recorder, dQK);
-      trackTensorBuffer(recorder, dQ);
-      trackTensorBuffer(recorder, dQKTransposed);
-      trackTensorBuffer(recorder, dK);
+        for (const buffer of headBuffers) {
+          recorder.trackTemporaryBuffer(buffer);
+        }
+        trackTensorBuffer(recorder, sTransposed);
+        trackTensorBuffer(recorder, dV);
+        trackTensorBuffer(recorder, vTransposed);
+        trackTensorBuffer(recorder, dS);
+        trackTensorBuffer(recorder, dQK);
+        trackTensorBuffer(recorder, dQ);
+        trackTensorBuffer(recorder, dQKTransposed);
+        trackTensorBuffer(recorder, dK);
+      }
+    }
+    if (recorder) {
+      for (const tensor of ownedRecorderInputTensors) {
+        trackTensorBuffer(recorder, tensor);
+      }
+    }
+    completed = true;
+    return {
+      gradQ: createTensor(gradQBuf, 'f32', [...q.shape], 'attn_grad_q'),
+      gradK: createTensor(gradKBuf, 'f32', [...k.shape], 'attn_grad_k'),
+      gradV: createTensor(gradVBuf, 'f32', [...v.shape], 'attn_grad_v'),
+    };
+  } finally {
+    if (!completed) {
+      releaseBuffer(gradQBuf);
+      releaseBuffer(gradKBuf);
+      releaseBuffer(gradVBuf);
+    }
+    if (!recorder) {
+      for (const tensor of ownedInputTensors) {
+        releaseTensorBuffer(tensor);
+      }
+    } else {
+      for (const tensor of ownedRecorderInputTensors) {
+        releaseTensorBuffer(tensor);
+      }
     }
   }
-
-  return {
-    gradQ: createTensor(gradQBuf, 'f32', [...q.shape], 'attn_grad_q'),
-    gradK: createTensor(gradKBuf, 'f32', [...k.shape], 'attn_grad_k'),
-    gradV: createTensor(gradVBuf, 'f32', [...v.shape], 'attn_grad_v'),
-  };
 }
 
 export async function runAttentionBackward(
@@ -256,11 +348,7 @@ export async function runAttentionBackward(
   if (!device) {
     throw new Error('runAttentionBackward requires a GPU device');
   }
-
-  const recorder = new CommandRecorder(device, 'attention_backward');
-  const result = await runAttentionBackwardCore(q, k, v, softmax, gradOutput, options, recorder);
-  recorder.submit();
-  return result;
+  return runAttentionBackwardCore(q, k, v, softmax, gradOutput, options);
 }
 
 export async function recordAttentionBackward(

package/src/gpu/kernels/backward/conv2d_backward.js

@@ -4,6 +4,19 @@ import { createPipeline, createUniformBufferWithView } from '../utils.js';
 import { dispatch, recordDispatch } from '../dispatch.js';
 import { getDevice } from '../../device.js';
 
+function destroyAfterSubmit(device, buffer) {
+  if (!buffer) {
+    return;
+  }
+  device.queue.onSubmittedWorkDone()
+    .then(() => {
+      buffer.destroy();
+    })
+    .catch(() => {
+      buffer.destroy();
+    });
+}
+
 export async function runConv2DBackward(input, weight, gradOutput, options = {}) {
   const { inChannels, outChannels, height, width, outHeight, outWidth, kernelH, kernelW, stride, pad, computeGradInput = true, computeGradWeight = true } = options;
 
@@ -67,7 +80,7 @@ export async function runConv2DBackward(input, weight, gradOutput, options = {})
     gradWeight = createTensor(outputBuf, 'f32', [outChannels, inChannels, kernelH, kernelW], 'conv2d_grad_weight');
   }
 
-  uniformBuffer.destroy();
+  destroyAfterSubmit(device, uniformBuffer);
   return { gradInput, gradWeight };
 }
 

package/src/gpu/kernels/bias_add.wgsl

@@ -14,6 +14,10 @@ struct Uniforms {
     dim: u32,
     data_offset: u32,  // byte offset into data buffer (divide by 4 for F32)
     bias_offset: u32,  // byte offset into bias buffer (divide by 4 for F32)
+    token_stride: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 override WORKGROUP_SIZE: u32 = 256u;
@@ -24,17 +28,15 @@ override WORKGROUP_SIZE: u32 = 256u;
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let total = u.num_tokens * u.dim;
-    if (idx >= total) {
+    let d = gid.x;
+    let token = gid.z * max(u.token_stride, 1u) + gid.y;
+    if (token >= u.num_tokens || d >= u.dim) {
         return;
     }
 
     // Convert byte offsets to F32 indices
     let data_base = u.data_offset / 4u;
     let bias_base = u.bias_offset / 4u;
-
-    let d = idx % u.dim;
+    let idx = token * u.dim + d;
     data[data_base + idx] = data[data_base + idx] + bias[bias_base + d];
 }
-

package/src/gpu/kernels/bias_add_f16.wgsl

@@ -18,6 +18,10 @@ struct Uniforms {
     dim: u32,
     data_offset: u32,  // byte offset into data buffer (divide by 2 for F16)
     bias_offset: u32,  // byte offset into bias buffer (divide by 2 for F16)
+    token_stride: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
 }
 
 override WORKGROUP_SIZE: u32 = 256u;
@@ -28,17 +32,16 @@ override WORKGROUP_SIZE: u32 = 256u;
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
-    let total = u.num_tokens * u.dim;
-    if (idx >= total) {
+    let d = gid.x;
+    let token = gid.z * max(u.token_stride, 1u) + gid.y;
+    if (token >= u.num_tokens || d >= u.dim) {
         return;
     }
 
     // Convert byte offsets to F16 indices
     let data_base = u.data_offset / 2u;
     let bias_base = u.bias_offset / 2u;
-
-    let d = idx % u.dim;
+    let idx = token * u.dim + d;
     let out = f32(data[data_base + idx]) + f32(bias[bias_base + d]);
     data[data_base + idx] = f16(out);
 }