@simulatte/doppler 0.1.5 → 0.1.7

package/src/gpu/kernels/fused_matmul_residual.js

@@ -1,7 +1,7 @@
 
 
 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 { getBuffer } from '../weight-buffer.js';
 import { dispatch, recordDispatch } from './dispatch.js';
@@ -47,7 +47,12 @@ export async function runMatmulResidualFused(
   const pipelineVariant = resolveFusedResidualVariant(input, residual);
   const pipeline = await getPipelineFast('fused_matmul_residual', pipelineVariant);
 
-  const output = outputBuffer || acquireBuffer(N * dtypeBytes(outputDtype), undefined, 'matmul_residual_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(
+    N * dtypeBytes(outputDtype),
+    undefined,
+    'matmul_residual_output'
+  );
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer (same layout as matmul_gemv)
   const uniformBuffer = createUniformBufferWithView(
@@ -68,21 +73,28 @@ export async function runMatmulResidualFused(
   );
 
   // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_residual_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: output } },
-      { binding: 4, resource: { buffer: residual.buffer } },
-    ],
-  });
-
-  // One workgroup per output element
-  const workgroups = N;
-  dispatch(device, pipeline, bindGroup, workgroups, 'matmul_residual_fused');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'matmul_residual_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: output } },
+        { binding: 4, resource: { buffer: residual.buffer } },
+      ],
+    });
+
+    const workgroups = N;
+    dispatch(device, pipeline, bindGroup, workgroups, 'matmul_residual_fused');
+  } catch (error) {
+    uniformBuffer.destroy();
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   uniformBuffer.destroy();
 
@@ -112,7 +124,12 @@ export async function recordMatmulResidualFused(
   const pipelineVariant = resolveFusedResidualVariant(input, residual);
   const pipeline = await getPipelineFast('fused_matmul_residual', pipelineVariant);
 
-  const output = outputBuffer || acquireBuffer(N * dtypeBytes(outputDtype), undefined, 'matmul_residual_output');
+  const ownedOutput = outputBuffer ? null : acquireBuffer(
+    N * dtypeBytes(outputDtype),
+    undefined,
+    'matmul_residual_output'
+  );
+  const output = outputBuffer || ownedOutput;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -132,21 +149,27 @@ export async function recordMatmulResidualFused(
   );
 
   // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'matmul_residual_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: output } },
-      { binding: 4, resource: { buffer: residual.buffer } },
-    ],
-  });
-
-  // One workgroup per output element
-  const workgroups = N;
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'matmul_residual_fused');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'matmul_residual_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: output } },
+        { binding: 4, resource: { buffer: residual.buffer } },
+      ],
+    });
+
+    const workgroups = N;
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, 'matmul_residual_fused');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
   return createTensor(output, outputDtype, [1, N], 'matmul_residual_output');
 }

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

@@ -42,6 +42,7 @@ async function _groupedPointwiseConv2D(target, input, weight, bias, options = {}
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = outChannels * height * width * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'grouped_pointwise_conv2d_output');
+  const spatial = height * width;
 
   const weightBuffer = getBuffer(weight);
   let biasBuffer = getBuffer(bias);
@@ -54,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((outChannels * height * width) / WORKGROUP_SIZES.DEFAULT)
-  );
+  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/grouped_pointwise_conv2d.wgsl

@@ -19,17 +19,14 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let spatial = u.height * u.width;
-    let out_size = u.out_channels * spatial;
-    if (idx >= out_size) {
+    let spatial_idx = gid.x;
+    let out_channel = gid.y;
+    if (spatial_idx >= spatial || out_channel >= u.out_channels) {
         return;
     }
-
-    let out_channel = idx / spatial;
-    let rem = idx - out_channel * spatial;
-    let y = rem / u.width;
-    let x = rem - y * u.width;
+    let y = spatial_idx / u.width;
+    let x = spatial_idx - y * u.width;
 
     let in_per_group = u.in_channels / u.groups;
     let out_per_group = u.out_channels / u.groups;
@@ -43,5 +40,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
         sum = sum + input[input_idx] * weight[weight_idx];
     }
 
-    output[idx] = sum;
+    output[out_channel * spatial + spatial_idx] = sum;
 }

package/src/gpu/kernels/grouped_pointwise_conv2d_f16.wgsl

@@ -23,17 +23,14 @@ struct Uniforms {
 
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-    let idx = gid.x;
     let spatial = u.height * u.width;
-    let out_size = u.out_channels * spatial;
-    if (idx >= out_size) {
+    let spatial_idx = gid.x;
+    let out_channel = gid.y;
+    if (spatial_idx >= spatial || out_channel >= u.out_channels) {
         return;
     }
-
-    let out_channel = idx / spatial;
-    let rem = idx - out_channel * spatial;
-    let y = rem / u.width;
-    let x = rem - y * u.width;
+    let y = spatial_idx / u.width;
+    let x = spatial_idx - y * u.width;
 
     let in_per_group = u.in_channels / u.groups;
     let out_per_group = u.out_channels / u.groups;
@@ -47,5 +44,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
         sum = sum + f32(input[input_idx]) * f32(weight[weight_idx]);
     }
 
-    output[idx] = f16(sum);
+    output[out_channel * spatial + spatial_idx] = f16(sum);
 }

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();
+  }
 }