@simulatte/doppler 0.1.6 → 0.1.7

package/src/gpu/kernels/repeat_channels.js

@@ -1,4 +1,4 @@
-import { acquireBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
 import { createTensor, dtypeBytes } from '../tensor.js';
 import { unifiedKernelWrapper } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -32,23 +32,31 @@ async function _repeatChannels(target, input, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = outChannels * height * width * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'repeat_channels_output');
+  const ownedOutput = outputBuffer ? null : output;
 
-  await unifiedKernelWrapper(
-    'repeat_channels',
-    target,
-    variant,
-    [input, output],
-    {
-      in_channels: inChannels,
-      height,
-      width,
-      repeats,
-      _pad0: 0,
-    },
-    [Math.ceil((height * width) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
-  );
-
-  return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
+  try {
+    await unifiedKernelWrapper(
+      'repeat_channels',
+      target,
+      variant,
+      [input, output],
+      {
+        in_channels: inChannels,
+        height,
+        width,
+        repeats,
+        _pad0: 0,
+      },
+      [Math.ceil((height * width) / WORKGROUP_SIZES.DEFAULT), outChannels, 1]
+    );
+
+    return createTensor(output, input.dtype, [outChannels, height, width], 'repeat_channels_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runRepeatChannels(input, options = {}) {

package/src/gpu/kernels/residual.js

@@ -82,6 +82,7 @@ function planResidualDispatch(target, size, elementsPerWorkgroup) {
 async function _residualAdd(target, a, b, size, options = {}) {
   const recorder = target && typeof target.beginComputePass === 'function' ? target : null;
   const { useVec4 = true, outputBuffer = null } = options;
+  const ownsOutput = outputBuffer == null;
 
   const { a: aAligned, b: bAligned, temps } = await alignResidualInputs(a, b, recorder);
   const outputDtype = inferOutputDtype(aAligned, bAligned);
@@ -97,15 +98,22 @@ async function _residualAdd(target, a, b, size, options = {}) {
     useVec4 ? VEC4_ELEMENTS_PER_WG : WORKGROUP_SIZES.DEFAULT
   );
 
-  await unifiedKernelWrapper(
-    'residual', target, variant,
-    [aAligned, bAligned, output],
-    { size, scale: 1, _pad1: dispatchPlan.dispatchStride, _pad2: 0 },
-    dispatchPlan.workgroups
-  );
-
-  cleanupTemps(temps, recorder);
-  return createTensor(output, outputDtype, [size], 'residual_output');
+  try {
+    await unifiedKernelWrapper(
+      'residual', target, variant,
+      [aAligned, bAligned, output],
+      { size, scale: 1, _pad1: dispatchPlan.dispatchStride, _pad2: 0 },
+      dispatchPlan.workgroups
+    );
+    return createTensor(output, outputDtype, [size], 'residual_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  } finally {
+    cleanupTemps(temps, recorder);
+  }
 }
 
 async function _biasAdd(target, data, bias, numTokens, dim, options = {}) {
@@ -126,24 +134,26 @@ async function _biasAdd(target, data, bias, numTokens, dim, options = {}) {
     Math.ceil(numTokens / tokenStride),
   ];
 
-  await unifiedKernelWrapper(
-    'bias_add', target, variant,
-    [data, biasAligned],
-    {
-      num_tokens: numTokens,
-      dim,
-      data_offset: dataOffset,
-      bias_offset: biasOffset,
-      token_stride: tokenStride,
-      _pad0: 0,
-      _pad1: 0,
-      _pad2: 0,
-    },
-    workgroups
-  );
-
-  cleanupTemps(temps, recorder);
-  return createTensor(data.buffer, data.dtype, [numTokens, dim], 'bias_add_output');
+  try {
+    await unifiedKernelWrapper(
+      'bias_add', target, variant,
+      [data, biasAligned],
+      {
+        num_tokens: numTokens,
+        dim,
+        data_offset: dataOffset,
+        bias_offset: biasOffset,
+        token_stride: tokenStride,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      workgroups
+    );
+    return createTensor(data.buffer, data.dtype, [numTokens, dim], 'bias_add_output');
+  } finally {
+    cleanupTemps(temps, recorder);
+  }
 }
 
 export async function runResidualAdd(a, b, size, options = {}) {

package/src/gpu/kernels/rmsnorm.js

@@ -1,7 +1,7 @@
 
 
 import { getKernelCapabilities } 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 { getKernelThresholds, padToQ4KBlock } from '../../config/schema/index.js';
 import { selectRuleValue } from './rule-registry.js';
@@ -119,31 +119,39 @@ export async function runRMSNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'rmsnorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
   const dispatchPlan = planRMSNormDispatch(null, batchSize);
 
   // Shader layout always includes the residual binding; when unused, bind a harmless placeholder.
   const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
 
-  await unifiedKernelWrapper(
-    'rmsnorm',
-    null,
-    variant,
-    [input, normWeightBuffer, outputBuf, residualBuf],
-    {
-      hidden_size: inferredHiddenSize,
-      num_tokens: batchSize,
-      eps,
-      has_residual: residual ? 1 : 0,
-      token_stride: dispatchPlan.tokenStride,
-      _pad0: 0,
-      _pad1: 0,
-      _pad2: 0,
-    },
-    dispatchPlan.workgroups,
-    { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
-  );
-
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  try {
+    await unifiedKernelWrapper(
+      'rmsnorm',
+      null,
+      variant,
+      [input, normWeightBuffer, outputBuf, residualBuf],
+      {
+        hidden_size: inferredHiddenSize,
+        num_tokens: batchSize,
+        eps,
+        has_residual: residual ? 1 : 0,
+        token_stride: dispatchPlan.tokenStride,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      dispatchPlan.workgroups,
+      { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
+    );
+
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function recordRMSNorm(
@@ -165,28 +173,36 @@ export async function recordRMSNorm(
   const paddedHiddenSize = padToQ4KBlock(inferredHiddenSize);
   const outputSize = batchSize * paddedHiddenSize * bytesPerElement;
   const outputBuf = outputBuffer || acquireBuffer(outputSize, undefined, 'rmsnorm_output');
+  const ownedOutput = outputBuffer ? null : outputBuf;
   const dispatchPlan = planRMSNormDispatch(recorder, batchSize);
 
   const residualBuf = residual?.buffer || residual || input?.buffer || input || outputBuf;
 
-  await unifiedKernelWrapper(
-    'rmsnorm',
-    recorder,
-    variant,
-    [input, normWeightBuffer, outputBuf, residualBuf],
-    {
-      hidden_size: inferredHiddenSize,
-      num_tokens: batchSize,
-      eps,
-      has_residual: residual ? 1 : 0,
-      token_stride: dispatchPlan.tokenStride,
-      _pad0: 0,
-      _pad1: 0,
-      _pad2: 0,
-    },
-    dispatchPlan.workgroups,
-    { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
-  );
-
-  return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  try {
+    await unifiedKernelWrapper(
+      'rmsnorm',
+      recorder,
+      variant,
+      [input, normWeightBuffer, outputBuf, residualBuf],
+      {
+        hidden_size: inferredHiddenSize,
+        num_tokens: batchSize,
+        eps,
+        has_residual: residual ? 1 : 0,
+        token_stride: dispatchPlan.tokenStride,
+        _pad0: 0,
+        _pad1: 0,
+        _pad2: 0,
+      },
+      dispatchPlan.workgroups,
+      { RMS_NORM_OFFSET: rmsNormWeightOffset, WEIGHT_IS_F16: normWeightDtype === 'f16' }
+    );
+
+    return createTensor(outputBuf, input.dtype, [batchSize, inferredHiddenSize], 'rmsnorm_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/rope.js

@@ -27,6 +27,9 @@ async function _rope(target, input, freqsCos, freqsSin, seqLen, options = {}) {
   if (rotaryDim <= 0 || rotaryDim > headDim) {
     throw new Error(`RoPE rotaryDim must be in (0, headDim]; got ${rotaryDim} for headDim ${headDim}`);
   }
+  if (input.dtype === 'f16' && (rotaryDim !== headDim || interleaved)) {
+    throw new Error('RoPE f16 kernel requires rotaryDim === headDim and interleaved === false.');
+  }
 
   const caps = getKernelCapabilities();
   const useF16 = input.dtype === 'f16' && caps.hasF16;

package/src/gpu/kernels/sample.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice, getKernelCapabilities } from '../device.js';
-import { acquireBuffer, releaseBuffer } from '../../memory/buffer-pool.js';
+import { acquireBuffer, readBufferSlice, releaseBuffer } from '../../memory/buffer-pool.js';
 import { WORKGROUP_SIZES } from './constants.js';
 import { createPipeline, createUniformBufferWithView, getOrCreateBindGroupLayout } from './utils.js';
 import { allowReadback } from '../perf-guards.js';
@@ -156,18 +156,19 @@ function ensureOutputBufferSize(outputBuffer, minBytes, label) {
   }
 }
 
-function readTokenFromOutput(device, outputBuffer, outputIndex, label) {
-  const stagingBuffer = device.createBuffer({
-    label,
-    size: 4,
-    usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
-  });
-
-  const copyEncoder = device.createCommandEncoder({ label: `${label}_copy` });
-  copyEncoder.copyBufferToBuffer(outputBuffer, outputIndex * 4, stagingBuffer, 0, 4);
-  device.queue.submit([copyEncoder.finish()]);
+async function readTokenFromOutput(outputBuffer, outputIndex) {
+  return new Uint32Array(await readBufferSlice(outputBuffer, outputIndex * 4, 4))[0];
+}
 
-  return stagingBuffer;
+function cleanupRunResources(uniformBuffer, ownedBuffers) {
+  if (uniformBuffer) {
+    uniformBuffer.destroy();
+  }
+  for (const buffer of ownedBuffers) {
+    if (buffer) {
+      releaseBuffer(buffer);
+    }
+  }
 }
 
 async function executeArgmaxRun(logits, vocabSize, options) {
@@ -238,20 +239,14 @@ async function executeArgmaxRun(logits, vocabSize, options) {
 
   device.queue.submit([encoder.finish()]);
 
-  const stagingBuffer = readTokenFromOutput(device, outputBuffer, outputIndex, 'argmax_staging');
-  await stagingBuffer.mapAsync(GPUMapMode.READ);
-  const tokenId = new Uint32Array(stagingBuffer.getMappedRange())[0];
-  stagingBuffer.unmap();
-
-  stagingBuffer.destroy();
-  uniformBuffer.destroy();
-  releaseBuffer(tempLogits);
-  releaseBuffer(tempIndices);
-  if (ownsOutputBuffer) {
-    releaseBuffer(outputBuffer);
+  try {
+    return await readTokenFromOutput(outputBuffer, outputIndex);
+  } finally {
+    cleanupRunResources(
+      uniformBuffer,
+      [tempLogits, tempIndices, ownsOutputBuffer ? outputBuffer : null]
+    );
   }
-
-  return tokenId;
 }
 
 async function executeArgmaxRecord(recorder, logits, vocabSize, options) {
@@ -428,20 +423,14 @@ export async function runGPUSample(
 
   device.queue.submit([encoder.finish()]);
 
-  const stagingBuffer = readTokenFromOutput(device, outputBuffer, outputIndex, 'sample_staging');
-  await stagingBuffer.mapAsync(GPUMapMode.READ);
-  const tokenId = new Uint32Array(stagingBuffer.getMappedRange())[0];
-  stagingBuffer.unmap();
-
-  stagingBuffer.destroy();
-  uniformBuffer.destroy();
-  releaseBuffer(topkLogits);
-  releaseBuffer(topkIndices);
-  if (ownsOutputBuffer) {
-    releaseBuffer(outputBuffer);
+  try {
+    return await readTokenFromOutput(outputBuffer, outputIndex);
+  } finally {
+    cleanupRunResources(
+      uniformBuffer,
+      [topkLogits, topkIndices, ownsOutputBuffer ? outputBuffer : null]
+    );
   }
-
-  return tokenId;
 }
 
 

package/src/gpu/kernels/sana_linear_attention.js

@@ -64,6 +64,8 @@ async function _sanaLinearAttention(target, query, key, value, options = {}) {
     outputBuffer = null,
     summaryBuffer = null,
   } = options;
+  const ownsSummary = summaryBuffer == null;
+  const ownsOutput = outputBuffer == null;
 
   if (
     !Number.isFinite(numHeads) ||
@@ -98,18 +100,24 @@ async function _sanaLinearAttention(target, query, key, value, options = {}) {
     eps,
   };
 
-  await runSummary(target, query, key, value, temporarySummary, uniforms, variant);
-  await runApply(target, query, temporarySummary, output, uniforms, variant);
-
-  if (!summaryBuffer) {
-    if (recorder) {
-      recorder.trackTemporaryBuffer(temporarySummary);
-    } else {
-      releaseBuffer(temporarySummary);
+  try {
+    await runSummary(target, query, key, value, temporarySummary, uniforms, variant);
+    await runApply(target, query, temporarySummary, output, uniforms, variant);
+    return createTensor(output, query.dtype, [numTokens, hiddenSize], 'sana_linear_attention_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(output);
+    }
+    throw error;
+  } finally {
+    if (ownsSummary) {
+      if (recorder) {
+        recorder.trackTemporaryBuffer(temporarySummary);
+      } else {
+        releaseBuffer(temporarySummary);
+      }
     }
   }
-
-  return createTensor(output, query.dtype, [numTokens, hiddenSize], 'sana_linear_attention_output');
 }
 
 export async function runSanaLinearAttention(query, key, value, options = {}) {

package/src/gpu/kernels/scale.js

@@ -1,4 +1,4 @@
-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';
@@ -6,6 +6,7 @@ import { selectRuleValue } from './rule-registry.js';
 
 async function _scale(target, input, scale, options = {}) {
   const { count, outputBuffer = null, inplace = false } = options;
+  const ownsOutput = !inplace && outputBuffer == null;
 
   const bytesPerElement = dtypeBytes(input.dtype);
   const inferredCount = count ?? Math.floor(input.buffer.size / bytesPerElement);
@@ -16,16 +17,22 @@ async function _scale(target, input, scale, options = {}) {
 
   const bindings = inplace ? [outputBuf, outputBuf] : [input, outputBuf];
 
-  await unifiedKernelWrapper(
-    'scale',
-    target,
-    variant,
-    bindings,
-    { size: inferredCount, scale },
-    Math.ceil(inferredCount / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  return createTensor(outputBuf, input.dtype, [...input.shape], 'scale_output');
+  try {
+    await unifiedKernelWrapper(
+      'scale',
+      target,
+      variant,
+      bindings,
+      { size: inferredCount, scale },
+      Math.ceil(inferredCount / WORKGROUP_SIZES.DEFAULT)
+    );
+    return createTensor(outputBuf, input.dtype, [...input.shape], 'scale_output');
+  } catch (error) {
+    if (ownsOutput) {
+      releaseBuffer(outputBuf);
+    }
+    throw error;
+  }
 }
 
 export async function runScale(input, scale, options = {}) {

package/src/gpu/kernels/shader-cache.js

@@ -138,8 +138,10 @@ export async function compileShader(
     code: source,
   });
 
-  // Check for compilation errors
-  const compilationInfo = await module.getCompilationInfo();
+  // Check for compilation errors (getCompilationInfo not available in all WebGPU providers)
+  const compilationInfo = typeof module.getCompilationInfo === 'function'
+    ? await module.getCompilationInfo()
+    : { messages: [] };
   if (compilationInfo.messages.length > 0) {
     for (const msg of compilationInfo.messages) {
       if (msg.type === 'error') {