@simulatte/doppler 0.1.6 → 0.1.7

package/src/gpu/kernels/fused_matmul_rmsnorm.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice } from '../device.js';
-import { acquireBuffer, getBufferRequestedSize } from '../../memory/buffer-pool.js';
+import { acquireBuffer, getBufferRequestedSize, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { getBuffer } from '../weight-buffer.js';
 import { dispatch, recordDispatch } from './dispatch.js';
@@ -91,7 +91,8 @@ export async function runMatmulRMSNormFused(
   // Output buffer: [1, N] - size depends on dtype
   const bytesPerElement = dtype === 'f16' ? 2 : 4;
   const outputSize = N * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer (8 u32/f32 = 32 bytes, padded for alignment)
   const uniformBuffer = createUniformBufferWithView(
@@ -110,36 +111,44 @@ export async function runMatmulRMSNormFused(
   );
 
   // Create placeholder for residual if not provided
+  const ownsResidualBuffer = !residual;
   const residualBuffer = residual || device.createBuffer({
     label: 'matmul_rmsnorm_residual_placeholder',
     size: 4,
     usage: GPUBufferUsage.STORAGE,
   });
 
-  // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_rmsnorm_fused_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: weightBuffer } },
-      { binding: 3, resource: { buffer: normWeightBuffer } },
-      { binding: 4, resource: { buffer: output } },
-      { binding: 5, resource: { buffer: residualBuffer } },
-    ],
-  });
-
-  // Calculate workgroups
-  
-  const workgroups = 1;
-
-  const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
-  dispatch(device, pipeline, bindGroup, workgroups, dispatchLabel);
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'matmul_rmsnorm_fused_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: weightBuffer } },
+        { binding: 3, resource: { buffer: normWeightBuffer } },
+        { binding: 4, resource: { buffer: output } },
+        { binding: 5, resource: { buffer: residualBuffer } },
+      ],
+    });
+
+    const workgroups = 1;
+    const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+    dispatch(device, pipeline, bindGroup, workgroups, dispatchLabel);
+  } catch (error) {
+    uniformBuffer.destroy();
+    if (ownsResidualBuffer) {
+      residualBuffer.destroy();
+    }
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   // Cleanup
   uniformBuffer.destroy();
-  if (!residual) residualBuffer.destroy();
+  if (ownsResidualBuffer) residualBuffer.destroy();
 
   // Output dtype matches input dtype
   return createTensor(output, input.dtype, [1, N], 'matmul_rmsnorm_fused_output');
@@ -199,7 +208,8 @@ export async function recordMatmulRMSNormFused(
   // Output buffer - size depends on dtype
   const bytesPerElement = dtype === 'f16' ? 2 : 4;
   const outputSize = N * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'matmul_rmsnorm_fused_output');
+  const output = outputBuffer || ownedOutput;
 
   // Uniform buffer via recorder (8 u32/f32 = 32 bytes, padded for alignment)
   const uniformBuffer = createUniformBufferWithView(
@@ -217,35 +227,42 @@ export async function recordMatmulRMSNormFused(
   );
 
   // Placeholder for residual
+  const ownsResidualBuffer = !residual;
   const residualBuffer = residual || device.createBuffer({
     label: 'matmul_rmsnorm_residual_placeholder',
     size: 4,
     usage: GPUBufferUsage.STORAGE,
   });
 
-  // Bind group
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_rmsnorm_fused_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: weightBuffer } },
-      { binding: 3, resource: { buffer: normWeightBuffer } },
-      { binding: 4, resource: { buffer: output } },
-      { binding: 5, resource: { buffer: residualBuffer } },
-    ],
-  });
-
-  // Calculate workgroups
-  
-  const workgroups = 1;
-
-  const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, dispatchLabel);
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'matmul_rmsnorm_fused_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: weightBuffer } },
+        { binding: 3, resource: { buffer: normWeightBuffer } },
+        { binding: 4, resource: { buffer: output } },
+        { binding: 5, resource: { buffer: residualBuffer } },
+      ],
+    });
+
+    const workgroups = 1;
+    const dispatchLabel = label ? `matmul_rmsnorm_fused:${label}` : 'matmul_rmsnorm_fused';
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, dispatchLabel);
+  } catch (error) {
+    if (ownsResidualBuffer) {
+      residualBuffer.destroy();
+    }
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   // Track placeholder for cleanup
-  if (!residual) {
+  if (ownsResidualBuffer) {
     recorder.trackTemporaryBuffer(residualBuffer);
   }
 

package/src/gpu/kernels/gather.js

@@ -1,5 +1,5 @@
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { WORKGROUP_SIZES, VEC4_ELEMENTS_PER_WG } from './constants.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { trace } from '../../debug/index.js';
@@ -26,7 +26,6 @@ async function _gather(
   options = {}
 ) {
   const {
-    useVec4 = true,
     outputBuffer = null,
     embeddingDtype,
     outputDtype,
@@ -43,9 +42,22 @@ async function _gather(
   if (outputDtype == null) {
     throw new Error('[Gather] outputDtype is required.');
   }
+  if (embeddingDtype === 'f16' && !caps.hasF16) {
+    throw new Error('[Gather] embeddingDtype=f16 requires shader-f16 support.');
+  }
+  if (outputDtype === 'f16' && !caps.hasF16) {
+    throw new Error('[Gather] outputDtype=f16 requires shader-f16 support.');
+  }
 
-  const useF16Input = embeddingDtype === 'f16' && caps.hasF16;
-  const useF16Output = outputDtype === 'f16' && caps.hasF16;
+  const requestedVec4 = options.useVec4;
+  const wantsVec4 = requestedVec4 ?? true;
+  if (requestedVec4 === true && hiddenSize % 4 !== 0) {
+    throw new Error('[Gather] useVec4=true requires hiddenSize to be divisible by 4.');
+  }
+
+  const useF16Input = embeddingDtype === 'f16';
+  const useF16Output = outputDtype === 'f16';
+  const useVec4 = wantsVec4 && hiddenSize % 4 === 0;
 
   trace.embed(
     `Gather: numTokens=${numTokens}, hiddenSize=${hiddenSize}, vocabSize=${vocabSize}, ` +
@@ -64,6 +76,7 @@ async function _gather(
   const paddedHiddenSize = padToQ4KBlock(hiddenSize);
   const outputSize = numTokens * paddedHiddenSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gather_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   const uniforms = {
     num_tokens: numTokens,
@@ -82,16 +95,22 @@ async function _gather(
       ? Math.ceil((numTokens * hiddenSize) / VEC4_ELEMENTS_PER_WG)
       : Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT));
 
-  await unifiedKernelWrapper(
-    'gather',
-    target,
-    variant,
-    [indices, embeddings, output],
-    uniforms,
-    workgroups
-  );
-
-  return createTensor(output, actualDtype, [numTokens, hiddenSize], 'gather_output');
+  try {
+    await unifiedKernelWrapper(
+      'gather',
+      target,
+      variant,
+      [indices, embeddings, output],
+      uniforms,
+      workgroups
+    );
+    return createTensor(output, actualDtype, [numTokens, hiddenSize], 'gather_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runGather(
@@ -116,4 +135,3 @@ export async function recordGather(
 ) {
   return _gather(recorder, indices, embeddings, numTokens, hiddenSize, vocabSize, options);
 }
-

package/src/gpu/kernels/gelu.js

@@ -1,5 +1,5 @@
 
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { WORKGROUP_SIZES } from './constants.js';
 import { unifiedKernelWrapper } from './utils.js';
@@ -26,16 +26,24 @@ async function _gelu(target, input, options = {}) {
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'gelu_output');
   const gateBuffer = gate ?? input;
-
-  await unifiedKernelWrapper(
-    'gelu', target, variant,
-    [input, output, gateBuffer],
-    { size: inferredSize, rowsplit_dim: 0 },
-    Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT),
-    overrides
-  );
-
-  return createTensor(output, input.dtype, [inferredSize], 'gelu_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'gelu', target, variant,
+      [input, output, gateBuffer],
+      { size: inferredSize, rowsplit_dim: 0 },
+      Math.ceil(inferredSize / WORKGROUP_SIZES.DEFAULT),
+      overrides
+    );
+
+    return createTensor(output, input.dtype, [inferredSize], 'gelu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runGeLU(input, options = {}) {

package/src/gpu/kernels/grouped_pointwise_conv2d.js

@@ -55,33 +55,43 @@ async function _groupedPointwiseConv2D(target, input, weight, bias, options = {}
     device.queue.writeBuffer(biasBuffer, 0, new Uint8Array(paddedSize));
   }
 
-  await unifiedKernelWrapper(
-    'grouped_pointwise_conv2d',
-    target,
-    variant,
-    [input, weightBuffer, biasBuffer, output],
-    {
-      in_channels: inChannels,
-      out_channels: outChannels,
-      height,
-      width,
-      groups,
-      _pad0: 0,
-      _pad1: 0,
-      _pad2: 0,
-    },
-    [Math.ceil(spatial / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
-  );
+  try {
+    await unifiedKernelWrapper(
+      'grouped_pointwise_conv2d',
+      target,
+      variant,
+      [input, weightBuffer, biasBuffer, output],
+      {
+        in_channels: inChannels,
+        out_channels: outChannels,
+        height,
+        width,
+        groups,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      [Math.ceil(spatial / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
+
+    if (tempBias) {
+      if (recorder) {
+        recorder.trackTemporaryBuffer(tempBias);
+      } else {
+        releaseBuffer(tempBias);
+      }
+    }
 
-  if (tempBias) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(tempBias);
-    } else {
+    return createTensor(output, input.dtype, [outChannels, height, width], 'grouped_pointwise_conv2d_output');
+  } catch (error) {
+    if (tempBias) {
       releaseBuffer(tempBias);
     }
+    if (!outputBuffer) {
+      releaseBuffer(output);
+    }
+    throw error;
   }
-
-  return createTensor(output, input.dtype, [outChannels, height, width], 'grouped_pointwise_conv2d_output');
 }
 
 export async function runGroupedPointwiseConv2D(input, weight, bias, options = {}) {

package/src/gpu/kernels/groupnorm.js

@@ -17,6 +17,9 @@ function validateOptions(options) {
   if (!Number.isFinite(numGroups) || numGroups <= 0) {
     throw new Error('GroupNorm requires numGroups > 0.');
   }
+  if (channels % numGroups !== 0) {
+    throw new Error('GroupNorm requires channels to be divisible by numGroups.');
+  }
   if (!Number.isFinite(eps)) {
     throw new Error('GroupNorm requires eps.');
   }
@@ -44,34 +47,42 @@ async function _groupNorm(target, input, weight, bias, options = {}) {
 
   const statsSize = numGroups * 2 * 4;
   const statsBuffer = acquireBuffer(statsSize, undefined, 'groupnorm_stats');
-
-  await unifiedKernelWrapper(
-    'groupnorm_stats',
-    target,
-    statsVariant,
-    [input, statsBuffer],
-    uniforms,
-    numGroups
-  );
-
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = channels * height * width * bytesPerElement;
-  const output = outputBuffer || acquireBuffer(outputSize, undefined, 'groupnorm_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(outputSize, undefined, 'groupnorm_output');
+  const output = outputBuffer || ownedOutput;
 
-  const weightBuffer = getBuffer(weight);
-  const biasBuffer = getBuffer(bias);
+  try {
+    await unifiedKernelWrapper(
+      'groupnorm_stats',
+      target,
+      statsVariant,
+      [input, statsBuffer],
+      uniforms,
+      numGroups
+    );
 
-  const total = channels * height * width;
-  const workgroups = Math.ceil(total / WORKGROUP_SIZES.DEFAULT);
+    const weightBuffer = getBuffer(weight);
+    const biasBuffer = getBuffer(bias);
 
-  await unifiedKernelWrapper(
-    'groupnorm_apply',
-    target,
-    applyVariant,
-    [input, statsBuffer, weightBuffer, biasBuffer, output],
-    uniforms,
-    workgroups
-  );
+    const total = channels * height * width;
+    const workgroups = Math.ceil(total / WORKGROUP_SIZES.DEFAULT);
+
+    await unifiedKernelWrapper(
+      'groupnorm_apply',
+      target,
+      applyVariant,
+      [input, statsBuffer, weightBuffer, biasBuffer, output],
+      uniforms,
+      workgroups
+    );
+  } catch (error) {
+    releaseBuffer(statsBuffer);
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   if (recorder) {
     recorder.trackTemporaryBuffer(statsBuffer);

package/src/gpu/kernels/kv-quantize.js

@@ -78,8 +78,11 @@ export async function runKVQuantize(
   });
 
   const workgroups = [numKVHeads, numTokens, 1];
-  dispatch(device, pipeline, bindGroup, workgroups, 'kv_quantize');
-  uniformBuffer.destroy();
+  try {
+    dispatch(device, pipeline, bindGroup, workgroups, 'kv_quantize');
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 

package/src/gpu/kernels/layernorm.js

@@ -1,6 +1,6 @@
 
 import { getKernelCapabilities } from '../device.js';
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor } from '../tensor.js';
 import { padToQ4KBlock } from '../../config/schema/index.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -36,17 +36,25 @@ export async function runLayerNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'layernorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
 
-  await unifiedKernelWrapper(
-    'layernorm',
-    null,
-    variant,
-    [input, weight, bias, outputBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
-    batchSize
-  );
+  try {
+    await unifiedKernelWrapper(
+      'layernorm',
+      null,
+      variant,
+      [input, weight, bias, outputBuf],
+      { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
+      batchSize
+    );
 
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function recordLayerNorm(
@@ -66,15 +74,23 @@ export async function recordLayerNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'layernorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
 
-  await unifiedKernelWrapper(
-    'layernorm',
-    recorder,
-    variant,
-    [input, weight, bias, outputBuf],
-    { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
-    batchSize
-  );
+  try {
+    await unifiedKernelWrapper(
+      'layernorm',
+      recorder,
+      variant,
+      [input, weight, bias, outputBuf],
+      { hidden_size: inferredHiddenSize, num_tokens: batchSize, eps },
+      batchSize
+    );
 
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'layernorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/logit-merge.js

@@ -266,9 +266,11 @@ export class LogitMergeKernel {
     pass.end();
 
     this.#device.queue.submit([encoder.finish()]);
-
-    // Cleanup temporary buffer
-    paramsBuffer.destroy();
+    this.#device.queue.onSubmittedWorkDone()
+      .catch(() => {})
+      .finally(() => {
+        paramsBuffer.destroy();
+      });
 
     return mergedBuffer;
   }

package/src/gpu/kernels/matmul.js

@@ -1,4 +1,4 @@
-import { getDevice } from '../device.js';
+import { getDevice, getKernelCapabilities } from '../device.js';
 import { createTensor } from '../tensor.js';
 import { getBuffer, getLayout, getWeightDtype } from '../weight-buffer.js';
 import { log, trace, isTraceEnabled } from '../../debug/index.js';
@@ -110,6 +110,7 @@ async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
   const mode = isRecord ? 'record' : 'run';
   const opLabel = isRecord ? 'recordMatmul' : 'runMatmul';
   const device = recorder?.device || getDevice();
+  const capabilities = getKernelCapabilities();
 
   const {
     alpha = 1.0,
@@ -139,6 +140,13 @@ async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
   const bDtype = toMatmulDtype(weightDtype ?? options.bDtype);
   const requestedOutputDtype = options.outputDtype || A.dtype;
 
+  if (bDtype === 'f16' && capabilities?.hasF16 !== true) {
+    throw new Error(`[${opLabel}] f16 weights require shader-f16 support.`);
+  }
+  if (requestedOutputDtype === 'f16' && capabilities?.hasF16 !== true) {
+    throw new Error(`[${opLabel}] f16 output requires shader-f16 support.`);
+  }
+
   if (!isRecord && isTraceEnabled('kernels') && !weightDtype && !options.bDtype && M <= 2) {
     log.warn('Matmul', `runMatmul: B buffer dtype unknown! size=${bBuffer.size}, M=${M}, N=${N}, K=${K}. Assuming f32.`);
   }
@@ -228,6 +236,7 @@ async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
     N,
     outputBuffer
   );
+  const ownsOutput = outputBuffer == null;
 
   if (!Number.isFinite(outputSize) || outputSize <= 0) {
     throw new Error(`[${opLabel}] Invalid output size: ${outputSize} (M=${M}, N=${N})`);
@@ -239,50 +248,60 @@ async function executeMatmul(recorder, A, B, M, N, K, options = {}) {
   }
 
   const dispatchPlan = calculateMatmulDispatch(variant, useQ4KFused, useGemv, M, N, config);
-  const uniformBuffer = createMatmulUniformBuffer(
-    'matmul_uniforms',
-    M,
-    N,
-    K,
-    alpha,
-    useQ4KFused,
-    transposeB,
-    dispatchPlan.uniformWorkgroupsX,
-    recorder || null,
-    device
-  );
+  let uniformBuffer = null;
+  let completed = false;
+  try {
+    uniformBuffer = createMatmulUniformBuffer(
+      'matmul_uniforms',
+      M,
+      N,
+      K,
+      alpha,
+      useQ4KFused,
+      transposeB,
+      dispatchPlan.uniformWorkgroupsX,
+      recorder || null,
+      device
+    );
 
-  const entries = createMatmulBindGroupEntries(
-    variant,
-    uniformBuffer,
-    matmulInput,
-    bBuffer,
-    C,
-    { aOffset, bOffset, cOffset },
-    {
-      aBindingSize: bindingSizes.aBindingSize,
-      bBindingSize: bindingSizes.bBindingSize,
-      cBindingSize,
-    }
-  );
+    const entries = createMatmulBindGroupEntries(
+      variant,
+      uniformBuffer,
+      matmulInput,
+      bBuffer,
+      C,
+      { aOffset, bOffset, cOffset },
+      {
+        aBindingSize: bindingSizes.aBindingSize,
+        bBindingSize: bindingSizes.bBindingSize,
+        cBindingSize,
+      }
+    );
 
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
+    const bindGroup = device.createBindGroup({
+      label: 'matmul_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
 
-  if (isRecord) {
-    kernel.record(recorder, pipeline, bindGroup, dispatchPlan.workgroups, buildProfileLabel(options));
-  } else {
-    kernel.dispatch(pipeline, bindGroup, dispatchPlan.workgroups);
-    releaseUniformBuffer(uniformBuffer);
-    if (castedInput) {
+    if (isRecord) {
+      kernel.record(recorder, pipeline, bindGroup, dispatchPlan.workgroups, buildProfileLabel(options));
+    } else {
+      kernel.dispatch(pipeline, bindGroup, dispatchPlan.workgroups);
+    }
+    completed = true;
+    return createTensor(C, actualOutputDtype, [M, N], 'matmul_output');
+  } finally {
+    if (!isRecord && uniformBuffer) {
+      releaseUniformBuffer(uniformBuffer);
+    }
+    if (!isRecord && castedInput) {
       releaseBuffer(castedInput.buffer);
     }
+    if (!completed && ownsOutput) {
+      releaseBuffer(C);
+    }
   }
-
-  return createTensor(C, actualOutputDtype, [M, N], 'matmul_output');
 }