@simulatte/doppler 0.1.6 → 0.1.8

package/src/gpu/kernels/silu.js

@@ -1,13 +1,26 @@
 
 
 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 { WORKGROUP_SIZES } from './constants.js';
 import { dispatch, recordDispatch } from './dispatch.js';
 import { getPipelineFast, createUniformBufferWithView } from './utils.js';
 import { selectRuleValue } from './rule-registry.js';
 
+function destroyAfterSubmit(device, buffer) {
+  if (!buffer) {
+    return;
+  }
+  device.queue.onSubmittedWorkDone()
+    .then(() => {
+      buffer.destroy();
+    })
+    .catch(() => {
+      buffer.destroy();
+    });
+}
+
 function canUseF16(input) {
   return input.dtype === 'f16';
 }
@@ -47,6 +60,12 @@ function createSiLUBindGroupEntries(uniformBuffer, input, output, gate) {
   ];
 }
 
+function cleanupRunResources(uniformBuffer, ownedOutput) {
+  if (ownedOutput) {
+    releaseBuffer(ownedOutput);
+  }
+}
+
 function planSiLUDispatch(device, size, useVec4) {
   const maxPerDim = Number.isFinite(device?.limits?.maxComputeWorkgroupsPerDimension)
     ? device.limits.maxComputeWorkgroupsPerDimension
@@ -97,6 +116,7 @@ export async function runSiLU(
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
   const dispatchPlan = planSiLUDispatch(device, inferredSize, useVec4);
 
   // Create uniform buffer
@@ -116,17 +136,21 @@ export async function runSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    dispatch(device, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -148,6 +172,7 @@ export async function runSwiGLURowsplitBias(
   const bytesPerElement = dtypeBytes(input.dtype);
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'swiglu_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -164,23 +189,27 @@ export async function runSwiGLURowsplitBias(
   );
 
   // Create bind group
-  const bindGroup = device.createBindGroup({
-    label: 'swiglu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: bias.buffer } },
-      { binding: 3, resource: { buffer: output } },
-    ],
-  });
-
-  const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
-  dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'swiglu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: bias.buffer } },
+        { binding: 3, resource: { buffer: output } },
+      ],
+    });
+
+    const workgroups = Math.ceil((numTokens * dim) / WORKGROUP_SIZES.DEFAULT);
+    dispatch(device, pipeline, bindGroup, workgroups, 'swiglu');
+    return createTensor(output, input.dtype, [numTokens, dim], 'swiglu_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    destroyAfterSubmit(device, uniformBuffer);
+  }
 }
 
 
@@ -202,6 +231,7 @@ export async function runSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Create uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -218,24 +248,28 @@ export async function runSiLURowSplit(
   );
 
   // Bind group: provide a dummy gate buffer to satisfy the fixed layout
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
-  dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  uniformBuffer.destroy();
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    dispatch(device, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    cleanupRunResources(null, ownedOutput);
+    throw error;
+  } finally {
+    uniformBuffer.destroy();
+  }
 }
 
 
@@ -258,6 +292,7 @@ export async function recordSiLURowSplit(
 
   const outputSize = numTokens * dim * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_rowsplit_output');
+  const ownedOutput = outputBuffer ? null : output;
 
   // Uniform buffer
   const uniformBuffer = createUniformBufferWithView(
@@ -272,22 +307,28 @@ export async function recordSiLURowSplit(
     recorder
   );
 
-  const gateBuffer = input.buffer;
-  const bindGroup = device.createBindGroup({
-    label: 'silu_rowsplit_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: input.buffer } },
-      { binding: 2, resource: { buffer: output } },
-      { binding: 3, resource: { buffer: gateBuffer } },
-    ],
-  });
-
-  const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
-  recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
-
-  return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  try {
+    const gateBuffer = input.buffer;
+    const bindGroup = device.createBindGroup({
+      label: 'silu_rowsplit_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: input.buffer } },
+        { binding: 2, resource: { buffer: output } },
+        { binding: 3, resource: { buffer: gateBuffer } },
+      ],
+    });
+
+    const workgroups = [Math.ceil(dim / WORKGROUP_SIZES.DEFAULT), numTokens, 1];
+    recordDispatch(recorder, pipeline, bindGroup, workgroups, 'silu_rowsplit');
+    return createTensor(output, input.dtype, [numTokens, dim], 'silu_rowsplit_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }
 
 
@@ -328,6 +369,7 @@ export async function recordSiLU(
   const inferredSize = size || (input.buffer.size / bytesPerElement);
   const outputSize = inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'silu_output');
+  const ownedOutput = outputBuffer ? null : output;
   const dispatchPlan = planSiLUDispatch(device, inferredSize, false);
 
   // Uniform buffer
@@ -346,13 +388,19 @@ export async function recordSiLU(
   // Create bind group using helper
   const entries = createSiLUBindGroupEntries(uniformBuffer, input, output, gate);
 
-  const bindGroup = device.createBindGroup({
-    label: 'silu_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries,
-  });
-
-  recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
-
-  return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  try {
+    const bindGroup = device.createBindGroup({
+      label: 'silu_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries,
+    });
+
+    recordDispatch(recorder, pipeline, bindGroup, dispatchPlan.workgroups, 'silu');
+    return createTensor(output, input.dtype, [inferredSize], 'silu_output');
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 }

package/src/gpu/kernels/softmax.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 { unifiedKernelWrapper } from './utils.js';
 import { createPipeline, createUniformBufferWithView, createBindGroupWithValidation } from './utils.js';
@@ -20,23 +20,34 @@ function selectSoftmaxVariant(innerSize) {
 
 async function _softmax(target, input, axis, options = {}) {
   const { batchSize = 1, size, seqLen, temperature = 1.0, outputBuffer = null } = options;
+  if (input.dtype !== 'f32') {
+    throw new Error(`Softmax requires f32 input, got ${input.dtype}.`);
+  }
 
-  const bytesPerElement = input.dtype === 'f16' ? 2 : 4;
+  const bytesPerElement = 4;
   const inferredSize = size || seqLen || (input.buffer.size / (batchSize * bytesPerElement));
   const variant = selectSoftmaxVariant(inferredSize);
   trace.kernels(`Softmax: size=${inferredSize}, variant=${variant}`);
 
   const outputSize = batchSize * inferredSize * bytesPerElement;
   const output = outputBuffer || acquireBuffer(outputSize, undefined, 'softmax_output');
+  const ownedOutput = outputBuffer ? null : output;
+
+  try {
+    await unifiedKernelWrapper(
+      'softmax', target, variant,
+      [input, output],
+      { inner_size: inferredSize, outer_size: batchSize, temperature },
+      batchSize
+    );
+  } catch (error) {
+    if (ownedOutput) {
+      releaseBuffer(ownedOutput);
+    }
+    throw error;
+  }
 
-  await unifiedKernelWrapper(
-    'softmax', target, variant,
-    [input, output],
-    { inner_size: inferredSize, outer_size: batchSize, temperature },
-    batchSize
-  );
-
-  return createTensor(output, input.dtype, [batchSize, inferredSize], 'softmax_output');
+  return createTensor(output, 'f32', [batchSize, inferredSize], 'softmax_output');
 }
 
 export async function runSoftmax(input, axis, options = {}) {
@@ -76,6 +87,7 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
 
   const indices = acquireBuffer(indicesSize, undefined, 'softmax_topk_indices');
   const weights = acquireBuffer(weightsSize, undefined, 'softmax_topk_weights');
+  let completed = false;
 
   const uniformBuffer = createUniformBufferWithView(
     'softmax_topk_uniforms', 16,
@@ -88,19 +100,26 @@ export async function runSoftmaxTopK(logits, numTokens, numExperts, topK, option
     null, device
   );
 
-  const bindGroup = await createBindGroupWithValidation(device, {
-    label: 'softmax_topk_bind_group',
-    layout: pipeline.getBindGroupLayout(0),
-    entries: [
-      { binding: 0, resource: { buffer: uniformBuffer } },
-      { binding: 1, resource: { buffer: logits } },
-      { binding: 2, resource: { buffer: indices } },
-      { binding: 3, resource: { buffer: weights } },
-    ],
-  }, `topk:${variant}`);
-
-  dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
-  uniformBuffer.destroy();
-
-  return { indices, weights };
+  try {
+    const bindGroup = await createBindGroupWithValidation(device, {
+      label: 'softmax_topk_bind_group',
+      layout: pipeline.getBindGroupLayout(0),
+      entries: [
+        { binding: 0, resource: { buffer: uniformBuffer } },
+        { binding: 1, resource: { buffer: logits } },
+        { binding: 2, resource: { buffer: indices } },
+        { binding: 3, resource: { buffer: weights } },
+      ],
+    }, `topk:${variant}`);
+
+    dispatchKernel(null, pipeline, bindGroup, numTokens, 'softmax_topk');
+    completed = true;
+    return { indices, weights };
+  } finally {
+    uniformBuffer.destroy();
+    if (!completed) {
+      releaseBuffer(indices);
+      releaseBuffer(weights);
+    }
+  }
 }

package/src/gpu/kernels/split_qg.d.ts

@@ -0,0 +1,50 @@
+/**
+ * Split Q and Gate Kernel
+ *
+ * De-interleaves Q and Gate projections from q_proj output for attentionOutputGate models.
+ * Models like Qwen 3.5 store q_proj weights in per-head interleaved layout:
+ *   rows [h*headDim*2 : h*headDim*2+headDim] = Q for head h
+ *   rows [h*headDim*2+headDim : (h+1)*headDim*2] = Gate for head h
+ * This kernel separates the full matmul output into contiguous Q and Gate tensors.
+ */
+
+import type { Tensor } from '../tensor.js';
+import type { CommandRecorder } from '../command-recorder.js';
+
+/** Split Q and Gate options */
+export interface SplitQGOptions {
+  numTokens: number;
+  numHeads: number;
+  headDim: number;
+  /** Pre-allocated Q output tensor */
+  qTensor?: Tensor | null;
+  /** Pre-allocated Gate output tensor */
+  gTensor?: Tensor | null;
+}
+
+/** Split Q and Gate result */
+export interface SplitQGResult {
+  Q: Tensor;
+  G: Tensor;
+}
+
+/**
+ * De-interleave Q and Gate from q_proj output.
+ *
+ * @param qgTensor - Full q_proj output [numTokens, numHeads * headDim * 2] (interleaved)
+ * @param options - Split configuration
+ * @returns Separate Q and Gate tensors, each [numTokens, numHeads * headDim]
+ */
+export declare function runSplitQG(
+  qgTensor: Tensor,
+  options: SplitQGOptions
+): Promise<SplitQGResult>;
+
+/**
+ * Record split Q and Gate (batched, no submit).
+ */
+export declare function recordSplitQG(
+  recorder: CommandRecorder,
+  qgTensor: Tensor,
+  options: SplitQGOptions
+): Promise<SplitQGResult>;

package/src/gpu/kernels/split_qg.js

@@ -0,0 +1,46 @@
+
+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';
+import { selectRuleValue } from './rule-registry.js';
+
+async function _splitQG(target, qgTensor, options) {
+  const { numTokens, numHeads, headDim, qTensor = null, gTensor = null } = options;
+  const ownsQ = qTensor == null;
+  const ownsG = gTensor == null;
+
+  const outputDtype = qgTensor.dtype;
+  const pipelineVariant = selectRuleValue('splitQg', 'variant', { outputDtype });
+  const bytesPerElement = dtypeBytes(outputDtype);
+  const qSize = numHeads * headDim;
+
+  const qBuffer = qTensor?.buffer || acquireBuffer(numTokens * qSize * bytesPerElement, undefined, 'Q');
+  const gBuffer = gTensor?.buffer || acquireBuffer(numTokens * qSize * bytesPerElement, undefined, 'Q_gate');
+
+  try {
+    await unifiedKernelWrapper(
+      'split_qg', target, pipelineVariant,
+      [qgTensor, qBuffer, gBuffer],
+      { num_tokens: numTokens, num_heads: numHeads, head_dim: headDim, _pad: 0 },
+      Math.ceil((numTokens * qSize) / WORKGROUP_SIZES.DEFAULT)
+    );
+
+    const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
+    const G = gTensor || createTensor(gBuffer, outputDtype, [numTokens, qSize], 'Q_gate');
+
+    return { Q, G };
+  } catch (error) {
+    if (ownsQ) releaseBuffer(qBuffer);
+    if (ownsG) releaseBuffer(gBuffer);
+    throw error;
+  }
+}
+
+export async function runSplitQG(qgTensor, options) {
+  return _splitQG(null, qgTensor, options);
+}
+
+export async function recordSplitQG(recorder, qgTensor, options) {
+  return _splitQG(recorder, qgTensor, options);
+}

package/src/gpu/kernels/split_qg.wgsl

@@ -0,0 +1,58 @@
+// split_qg.wgsl
+
+/**
+ * De-interleave Q and Gate projections from q_proj output for attentionOutputGate models.
+ *
+ * Models like Qwen 3.5 store q_proj weights with interleaved head layout:
+ *   rows [h*headDim*2 : h*headDim*2+headDim]     = Q for head h
+ *   rows [h*headDim*2+headDim : (h+1)*headDim*2] = Gate for head h
+ *
+ * A single full matmul over all 2*qSize rows produces interleaved output:
+ *   input[token, h*headDim*2 : h*headDim*2+headDim]     = Q head h
+ *   input[token, h*headDim*2+headDim : (h+1)*headDim*2] = Gate head h
+ *
+ * This kernel separates them into contiguous Q and G outputs:
+ *   Q[token, h*headDim + dim] = input[token, h*headDim*2 + dim]
+ *   G[token, h*headDim + dim] = input[token, h*headDim*2 + headDim + dim]
+ *
+ * Input layout  (row-major): [numTokens, numHeads * headDim * 2]
+ * Output Q layout (row-major): [numTokens, numHeads * headDim]
+ * Output G layout (row-major): [numTokens, numHeads * headDim]
+ */
+
+struct Params {
+    num_tokens: u32,
+    num_heads: u32,
+    head_dim: u32,
+    _pad: u32,
+}
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+@group(0) @binding(0) var<uniform> params: Params;
+@group(0) @binding(1) var<storage, read> input: array<f32>;
+@group(0) @binding(2) var<storage, read_write> Q: array<f32>;
+@group(0) @binding(3) var<storage, read_write> G: array<f32>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let q_size = params.num_heads * params.head_dim;
+    let total_elements = params.num_tokens * q_size;
+
+    if (idx >= total_elements) {
+        return;
+    }
+
+    let token = idx / q_size;
+    let elem = idx % q_size;
+    let head = elem / params.head_dim;
+    let dim = elem % params.head_dim;
+
+    // Input is interleaved per head: [Q_h (headDim elems), G_h (headDim elems)]
+    let src_q = token * (q_size * 2u) + head * (params.head_dim * 2u) + dim;
+    let src_g = src_q + params.head_dim;
+
+    Q[idx] = input[src_q];
+    G[idx] = input[src_g];
+}

package/src/gpu/kernels/split_qg_f16.wgsl

@@ -0,0 +1,62 @@
+// AUTO-GENERATED from src/gpu/kernels/split_qg.wgsl.
+// Edit the source kernel and tools/configs/wgsl-variants.js, then run `npm run kernels:generate`.
+// split_qg_f16.wgsl
+
+/**
+ * De-interleave Q and Gate projections from q_proj output for attentionOutputGate models (f16).
+ *
+ * Models like Qwen 3.5 store q_proj weights with interleaved head layout:
+ *   rows [h*headDim*2 : h*headDim*2+headDim]     = Q for head h
+ *   rows [h*headDim*2+headDim : (h+1)*headDim*2] = Gate for head h
+ *
+ * A single full matmul over all 2*qSize rows produces interleaved output:
+ *   input[token, h*headDim*2 : h*headDim*2+headDim]     = Q head h
+ *   input[token, h*headDim*2+headDim : (h+1)*headDim*2] = Gate head h
+ *
+ * This kernel separates them into contiguous Q and G outputs:
+ *   Q[token, h*headDim + dim] = input[token, h*headDim*2 + dim]
+ *   G[token, h*headDim + dim] = input[token, h*headDim*2 + headDim + dim]
+ *
+ * Input layout  (row-major): [numTokens, numHeads * headDim * 2]
+ * Output Q layout (row-major): [numTokens, numHeads * headDim]
+ * Output G layout (row-major): [numTokens, numHeads * headDim]
+ */
+
+enable f16;
+
+struct Params {
+    num_tokens: u32,
+    num_heads: u32,
+    head_dim: u32,
+    _pad: u32,
+}
+
+override WORKGROUP_SIZE: u32 = 256u;
+
+@group(0) @binding(0) var<uniform> params: Params;
+@group(0) @binding(1) var<storage, read> input: array<f16>;
+@group(0) @binding(2) var<storage, read_write> Q: array<f16>;
+@group(0) @binding(3) var<storage, read_write> G: array<f16>;
+
+@compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    let q_size = params.num_heads * params.head_dim;
+    let total_elements = params.num_tokens * q_size;
+
+    if (idx >= total_elements) {
+        return;
+    }
+
+    let token = idx / q_size;
+    let elem = idx % q_size;
+    let head = elem / params.head_dim;
+    let dim = elem % params.head_dim;
+
+    // Input is interleaved per head: [Q_h (headDim elems), G_h (headDim elems)]
+    let src_q = token * (q_size * 2u) + head * (params.head_dim * 2u) + dim;
+    let src_g = src_q + params.head_dim;
+
+    Q[idx] = input[src_q];
+    G[idx] = input[src_g];
+}

package/src/gpu/kernels/split_qkv.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';
@@ -7,6 +7,9 @@ import { selectRuleValue } from './rule-registry.js';
 
 async function _splitQKV(target, qkvTensor, options) {
   const { numTokens, qSize, kSize, vSize, qTensor = null, kTensor = null, vTensor = null } = options;
+  const ownsQ = qTensor == null;
+  const ownsK = kTensor == null;
+  const ownsV = vTensor == null;
 
   const outputDtype = qkvTensor.dtype;
   const pipelineVariant = selectRuleValue('splitQkv', 'variant', { outputDtype });
@@ -18,18 +21,25 @@ async function _splitQKV(target, qkvTensor, options) {
 
   const totalElements = numTokens * (qSize + kSize + vSize);
 
-  await unifiedKernelWrapper(
-    'split_qkv', target, pipelineVariant,
-    [qkvTensor, qBuffer, kBuffer, vBuffer],
-    { num_tokens: numTokens, q_size: qSize, k_size: kSize, v_size: vSize },
-    Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT)
-  );
-
-  const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
-  const K = kTensor || createTensor(kBuffer, outputDtype, [numTokens, kSize], 'K');
-  const V = vTensor || createTensor(vBuffer, outputDtype, [numTokens, vSize], 'V');
-
-  return { Q, K, V };
+  try {
+    await unifiedKernelWrapper(
+      'split_qkv', target, pipelineVariant,
+      [qkvTensor, qBuffer, kBuffer, vBuffer],
+      { num_tokens: numTokens, q_size: qSize, k_size: kSize, v_size: vSize },
+      Math.ceil(totalElements / WORKGROUP_SIZES.DEFAULT)
+    );
+
+    const Q = qTensor || createTensor(qBuffer, outputDtype, [numTokens, qSize], 'Q');
+    const K = kTensor || createTensor(kBuffer, outputDtype, [numTokens, kSize], 'K');
+    const V = vTensor || createTensor(vBuffer, outputDtype, [numTokens, vSize], 'V');
+
+    return { Q, K, V };
+  } catch (error) {
+    if (ownsQ) releaseBuffer(qBuffer);
+    if (ownsK) releaseBuffer(kBuffer);
+    if (ownsV) releaseBuffer(vBuffer);
+    throw error;
+  }
 }
 
 export async function runSplitQKV(qkvTensor, options) {