@simulatte/doppler 0.1.5 → 0.1.7

package/src/inference/pipelines/text/logits/gpu.js

@@ -1,7 +1,7 @@
 
 
 import { getDevice, getKernelCapabilities } from '../../../../gpu/device.js';
-import { acquireBuffer, releaseBuffer, readBuffer } from '../../../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../../../memory/buffer-pool.js';
 import { runMatmul, runRMSNorm } from '../../../../gpu/kernel-selector.js';
 import { recordMatmul } from '../../../../gpu/kernels/matmul.js';
 import { recordRMSNorm } from '../../../../gpu/kernels/rmsnorm.js';
@@ -13,6 +13,7 @@ import { getRuntimeConfig } from '../../../../config/runtime.js';
 import { selectRuleValue } from '../../../../rules/rule-registry.js';
 import { runProbes } from '../probes.js';
 import { f16BufferToF32 } from './cpu.js';
+import { readBufferWithCleanup } from './utils.js';
 
 function shouldForceStableF32Logits(config, inputDtype) {
   // Small Gemma-family checkpoints can overflow in pure F16 logits path after RMSNorm offset.
@@ -187,14 +188,18 @@ export async function computeChunkedLogitsGPU(
     }
 
     const logitsBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: logitsTensor.dtype });
-    const chunkLogitsData = await readBuffer(logitsTensor.buffer, numTokens * rowCount * logitsBytes);
+    const chunkLogitsData = await readBufferWithCleanup(
+      logitsTensor.buffer,
+      numTokens * rowCount * logitsBytes,
+      () => {
+        releaseBuffer(logitsTensor.buffer);
+        releaseBuffer(weightBuffer.buffer);
+      }
+    );
     const chunkLogits = logitsTensor.dtype === 'f16'
       ? f16BufferToF32(chunkLogitsData)
       : new Float32Array(chunkLogitsData);
     writeChunkLogits(logits, chunkLogits, numTokens, vocabSize, rowOffset, rowCount);
-
-    releaseBuffer(logitsTensor.buffer);
-    releaseBuffer(weightBuffer.buffer);
   }
 
   return logits;

package/src/inference/pipelines/text/logits/index.js

@@ -7,7 +7,7 @@ export { rmsNormCPU, matmulCPU, applySoftcapping, f16ToF32, f16BufferToF32 } fro
 export { computeLogitsGPU, recordLogitsGPU, computeChunkedLogitsGPU, resolveCpuWeightDims, resolveLmHeadChunkRows, extractLmHeadChunk, writeChunkLogits } from './gpu.js';
 
 // Re-export utilities
-export { extractLastPositionLogits, finalizeLogits } from './utils.js';
+export { extractLastPositionLogits, finalizeLogits, readBufferWithCleanup } from './utils.js';
 
 // Imports for computeLogits orchestrator
 import { getDevice } from '../../../../gpu/device.js';
@@ -20,7 +20,7 @@ import { log, trace, isTraceEnabled } from '../../../../debug/index.js';
 import { runProbes } from '../probes.js';
 import { rmsNormCPU, matmulCPU, f16BufferToF32 } from './cpu.js';
 import { resolveCpuWeightDims, computeChunkedLogitsGPU } from './gpu.js';
-import { finalizeLogits } from './utils.js';
+import { finalizeLogits, readBufferWithCleanup } from './utils.js';
 import { getRuntimeConfig } from '../../../../config/runtime.js';
 import { selectRuleValue } from '../../../../rules/rule-registry.js';
 
@@ -288,15 +288,14 @@ export async function computeLogits(
   // 4. Read back logits
   const logitsBytes = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: logitsTensor.dtype });
   const logitsReadSize = matmulRows * matmulVocabSize * logitsBytes;
-  const logitsData = await readBuffer(logitsTensor.buffer, logitsReadSize);
-
-  // Cleanup
-  if (inputBufferOwned) releaseBuffer(inputBuffer);
-  releaseBuffer(normedTensor.buffer);
-  if (matmulInputOwned) releaseBuffer(matmulInputTensor.buffer);
-  releaseBuffer(logitsTensor.buffer);
-  if (!getNormWeightBuffer && !(finalNorm instanceof GPUBuffer)) releaseBuffer(normWeightBuffer);
-  if (lmHeadBufferOwned) releaseBuffer(lmHeadGPU);
+  const logitsData = await readBufferWithCleanup(logitsTensor.buffer, logitsReadSize, () => {
+    if (inputBufferOwned) releaseBuffer(inputBuffer);
+    releaseBuffer(normedTensor.buffer);
+    if (matmulInputOwned) releaseBuffer(matmulInputTensor.buffer);
+    releaseBuffer(logitsTensor.buffer);
+    if (!getNormWeightBuffer && !(finalNorm instanceof GPUBuffer)) releaseBuffer(normWeightBuffer);
+    if (lmHeadBufferOwned) releaseBuffer(lmHeadGPU);
+  });
 
   const rawLogits = logitsTensor.dtype === 'f16'
     ? f16BufferToF32(logitsData)

package/src/inference/pipelines/text/logits/utils.d.ts

@@ -25,6 +25,13 @@ export function extractLastPositionLogits(
   vocabSize: number
 ): Float32Array;
 
+export function readBufferWithCleanup(
+  buffer: GPUBuffer,
+  byteLength: number,
+  cleanup?: (() => void) | null,
+  reader?: ((buffer: GPUBuffer, byteLength: number) => Promise<ArrayBuffer>) | null
+): Promise<ArrayBuffer>;
+
 /**
  * Finalize logits by applying padding and softcapping.
  *

package/src/inference/pipelines/text/logits/utils.js

@@ -1,5 +1,6 @@
 
 
+import { readBuffer } from '../../../../memory/buffer-pool.js';
 import { runProbes } from '../probes.js';
 import { applySoftcapping } from './cpu.js';
 
@@ -19,6 +20,14 @@ export function extractLastPositionLogits(
   return lastPosLogits;
 }
 
+export async function readBufferWithCleanup(buffer, byteLength, cleanup, reader = readBuffer) {
+  try {
+    return await reader(buffer, byteLength);
+  } finally {
+    cleanup?.();
+  }
+}
+
 
 export async function finalizeLogits(
   rawLogits,

package/src/inference/pipelines/text/lora-apply.js

@@ -17,42 +17,60 @@ export async function applyLoRA(input, baseOutput, lora, dims, getWeightBuffer,
 
   const aBuf = getWeightBuffer(lora.a, 'lora_a');
   const bBuf = getWeightBuffer(lora.b, 'lora_b');
-  const ownsA = !(lora.a instanceof GPUBuffer) && !isWeightBuffer(lora.a);
-  const ownsB = !(lora.b instanceof GPUBuffer) && !isWeightBuffer(lora.b);
-
-  const loraIntermediate = recorder
-    ? await recordMatmul(recorder, input, aBuf, M, rank, K, { transposeB: 'auto', role: 'lora_a', kernelPath })
-    : await runMatmul(input, aBuf, M, rank, K, { transposeB: 'auto', role: 'lora_a', kernelPath });
+  const ownsA = !(typeof GPUBuffer !== 'undefined' && lora.a instanceof GPUBuffer) && !isWeightBuffer(lora.a);
+  const ownsB = !(typeof GPUBuffer !== 'undefined' && lora.b instanceof GPUBuffer) && !isWeightBuffer(lora.b);
+  // Extract underlying GPUBuffer for WeightBuffers
+  const aBufGPU = isWeightBuffer(aBuf) ? aBuf.buffer : aBuf;
+  const bBufGPU = isWeightBuffer(bBuf) ? bBuf.buffer : bBuf;
+  let loraIntermediate = null;
+  let loraOutput = null;
+  let scaled = null;
+  try {
+    loraIntermediate = recorder
+      ? await recordMatmul(recorder, input, aBuf, M, rank, K, { transposeB: 'auto', role: 'lora_a', kernelPath })
+      : await runMatmul(input, aBuf, M, rank, K, { transposeB: 'auto', role: 'lora_a', kernelPath });
 
-  const loraOutput = recorder
-    ? await recordMatmul(recorder, loraIntermediate, bBuf, M, N, rank, { transposeB: 'auto', role: 'lora_b', kernelPath })
-    : await runMatmul(loraIntermediate, bBuf, M, N, rank, { transposeB: 'auto', role: 'lora_b', kernelPath });
+    loraOutput = recorder
+      ? await recordMatmul(recorder, loraIntermediate, bBuf, M, N, rank, { transposeB: 'auto', role: 'lora_b', kernelPath })
+      : await runMatmul(loraIntermediate, bBuf, M, N, rank, { transposeB: 'auto', role: 'lora_b', kernelPath });
 
-  const scaled = recorder
-    ? await recordScale(recorder, loraOutput, lora.scale, { outputBuffer: null })
-    : await runScale(loraOutput, lora.scale, { outputBuffer: null });
+    scaled = recorder
+      ? await recordScale(recorder, loraOutput, lora.scale, { outputBuffer: null })
+      : await runScale(loraOutput, lora.scale, { outputBuffer: null });
 
-  const combined = recorder
-    ? await recordResidualAdd(recorder, baseOutput, scaled, M * N)
-    : await runResidualAdd(baseOutput, scaled, M * N);
+    const combined = recorder
+      ? await recordResidualAdd(recorder, baseOutput, scaled, M * N)
+      : await runResidualAdd(baseOutput, scaled, M * N);
 
-  // Extract underlying GPUBuffer for WeightBuffers
-  const aBufGPU = isWeightBuffer(aBuf) ? aBuf.buffer : aBuf;
-  const bBufGPU = isWeightBuffer(bBuf) ? bBuf.buffer : bBuf;
+    if (recorder) {
+      recorder.trackTemporaryBuffer(loraIntermediate.buffer);
+      recorder.trackTemporaryBuffer(loraOutput.buffer);
+      recorder.trackTemporaryBuffer(scaled.buffer);
+      if (ownsA) recorder.trackTemporaryBuffer(aBufGPU);
+      if (ownsB) recorder.trackTemporaryBuffer(bBufGPU);
+    } else {
+      releaseBuffer(loraIntermediate.buffer);
+      releaseBuffer(loraOutput.buffer);
+      releaseBuffer(scaled.buffer);
+      if (ownsA) releaseBuffer(aBufGPU);
+      if (ownsB) releaseBuffer(bBufGPU);
+    }
 
-  if (recorder) {
-    recorder.trackTemporaryBuffer(loraIntermediate.buffer);
-    recorder.trackTemporaryBuffer(loraOutput.buffer);
-    recorder.trackTemporaryBuffer(scaled.buffer);
-    if (ownsA) recorder.trackTemporaryBuffer(aBufGPU);
-    if (ownsB) recorder.trackTemporaryBuffer(bBufGPU);
-  } else {
-    releaseBuffer(loraIntermediate.buffer);
-    releaseBuffer(loraOutput.buffer);
-    releaseBuffer(scaled.buffer);
-    if (ownsA) releaseBuffer(aBufGPU);
-    if (ownsB) releaseBuffer(bBufGPU);
+    return combined;
+  } catch (error) {
+    if (recorder) {
+      if (loraIntermediate) recorder.trackTemporaryBuffer(loraIntermediate.buffer);
+      if (loraOutput) recorder.trackTemporaryBuffer(loraOutput.buffer);
+      if (scaled) recorder.trackTemporaryBuffer(scaled.buffer);
+      if (ownsA) recorder.trackTemporaryBuffer(aBufGPU);
+      if (ownsB) recorder.trackTemporaryBuffer(bBufGPU);
+    } else {
+      if (loraIntermediate) releaseBuffer(loraIntermediate.buffer);
+      if (loraOutput) releaseBuffer(loraOutput.buffer);
+      if (scaled) releaseBuffer(scaled.buffer);
+      if (ownsA) releaseBuffer(aBufGPU);
+      if (ownsB) releaseBuffer(bBufGPU);
+    }
+    throw error;
   }
-
-  return combined;
 }