@simulatte/doppler 0.1.5 → 0.1.7

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';
@@ -52,6 +52,23 @@ function buildProfileLabel(options = {}) {
   return `matmul${roleLabel}${layerLabel}`;
 }
 
+function assertBindGroupBuffer(kernelName, variant, bindingIndex, bindingLabel, buffer, details = []) {
+  const isGpuBuffer = buffer && (
+    typeof GPUBuffer === 'undefined'
+      ? true
+      : buffer instanceof GPUBuffer
+  );
+  if (isGpuBuffer) {
+    return;
+  }
+  const detailText = details.filter(Boolean).join(', ');
+  throw new Error(
+    `[${kernelName}] variant="${variant}" binding ${bindingIndex} "${bindingLabel}" requires a GPUBuffer` +
+    (detailText ? ` (${detailText})` : '') +
+    '.'
+  );
+}
+
 function createMatmulBindGroupEntries(variant, uniformBuffer, matmulInput, bBuffer, outputBuffer, offsets, bindingSizes) {
   const isQ4KF16 = variant === 'q4_fused_multicol_f16'
     || variant === 'q4_fused_f16a'
@@ -59,6 +76,14 @@ function createMatmulBindGroupEntries(variant, uniformBuffer, matmulInput, bBuff
     || variant === 'q4_fused_multicol_f16a'
     || variant === 'q4_fused_batched_f16a';
 
+  assertBindGroupBuffer('matmul', variant, 0, 'uniforms', uniformBuffer);
+  assertBindGroupBuffer('matmul', variant, 1, 'input', matmulInput?.buffer, [
+    `inputLabel=${matmulInput?.label ?? 'unknown'}`,
+    `inputDtype=${matmulInput?.dtype ?? 'unknown'}`,
+  ]);
+  assertBindGroupBuffer('matmul', variant, 2, 'weights', bBuffer);
+  assertBindGroupBuffer('matmul', variant, isQ4KF16 ? 4 : 3, 'output', outputBuffer);
+
   const entries = [
     { binding: 0, resource: { buffer: uniformBuffer } },
     { binding: 1, resource: { buffer: matmulInput.buffer, offset: offsets.aOffset, size: bindingSizes.aBindingSize } },
@@ -85,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,
@@ -114,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.`);
   }
@@ -203,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})`);
@@ -214,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');
 }
 
 

package/src/gpu/kernels/modulate.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';
@@ -24,20 +24,28 @@ async function _modulate(target, input, mod, options = {}) {
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = numTokens * hiddenSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'modulate_output');
-
-  await unifiedKernelWrapper(
-    'modulate', target, variant,
-    [input, mod, output],
-    {
-      num_tokens: numTokens, hidden_size: hiddenSize,
-      scale_offset: scaleOffset, shift_offset: shiftOffset,
-      gate_offset: gateOffset, has_gate: hasGate ? 1 : 0,
-      add_one: addOne ? 1 : 0, _pad0: 0,
-    },
-    Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  return createTensor(output, input.dtype, [numTokens, hiddenSize], 'modulate_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'modulate', target, variant,
+      [input, mod, output],
+      {
+        num_tokens: numTokens, hidden_size: hiddenSize,
+        scale_offset: scaleOffset, shift_offset: shiftOffset,
+        gate_offset: gateOffset, has_gate: hasGate ? 1 : 0,
+        add_one: addOne ? 1 : 0, _pad0: 0,
+      },
+      Math.ceil((numTokens * hiddenSize) / WORKGROUP_SIZES.DEFAULT)
+    );
+
+    return createTensor(output, input.dtype, [numTokens, hiddenSize], 'modulate_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 export async function runModulate(input, mod, options = {}) {