@simulatte/doppler 0.1.6 → 0.1.7

package/src/inference/pipelines/text/init.js

@@ -2,7 +2,7 @@
 
 import { parseModelConfig } from './config.js';
 import { getDevice, getDeviceLimits, getKernelCapabilities } from '../../../gpu/device.js';
-import { acquireBuffer } from '../../../memory/buffer-pool.js';
+import { acquireBuffer, releaseBuffer } from '../../../memory/buffer-pool.js';
 import { KVCache, SlidingWindowKVCache, TieredKVCache, BasisDecomposedPagedCache } from '../../kv-cache.js';
 import { Tokenizer } from '../../tokenizer.js';
 import { MoERouter } from '../../moe-router.js';
@@ -14,6 +14,10 @@ import { PAGED_LAYOUT_SEQ_LEN_THRESHOLD } from '../../../config/schema/index.js'
 import { isKernelPathFusedQ4K } from '../../../config/kernel-path-loader.js';
 import { createWeightBuffer, getWeightDtype, isWeightBuffer } from '../../../gpu/weight-buffer.js';
 import { selectRuleValue } from '../../../rules/rule-registry.js';
+import {
+  createSourceStorageContext,
+  getSourceRuntimeMetadata,
+} from '../../../tooling/source-runtime-bundle.js';
 
 function resolveErrorMessage(error) {
   if (error && typeof error === 'object' && typeof error.message === 'string') {
@@ -56,12 +60,61 @@ function normalizeBaseUrl(baseUrl) {
   return baseUrl.replace(/\/$/, '');
 }
 
+async function fetchBytes(url, offset = null, length = null) {
+  const headers = {};
+  if (Number.isFinite(offset) && Number.isFinite(length) && length > 0) {
+    const start = Math.max(0, Math.floor(offset));
+    const end = start + Math.max(0, Math.floor(length)) - 1;
+    headers.Range = `bytes=${start}-${end}`;
+  }
+  const response = await fetch(url, { headers });
+  if (!response.ok) {
+    throw new Error(`Failed to fetch ${url}: ${response.status}`);
+  }
+  return new Uint8Array(await response.arrayBuffer());
+}
+
 function createRemoteStorageContext(baseUrl, manifest) {
   const root = normalizeBaseUrl(baseUrl);
   if (!root || !isRDRRManifest(manifest)) {
     return null;
   }
 
+  const sourceRuntime = getSourceRuntimeMetadata(manifest);
+  if (sourceRuntime) {
+    const readRange = async (relativePath, offset, length) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) {
+        throw new Error('Direct-source artifact path is required.');
+      }
+      const url = `${root}/${filename}`;
+      return fetchBytes(url, offset, length);
+    };
+    const readText = async (relativePath) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) return null;
+      const response = await fetch(`${root}/${filename}`);
+      if (!response.ok) {
+        throw new Error(`Failed to fetch ${filename} from ${root}: ${response.status}`);
+      }
+      return response.text();
+    };
+    const readBinary = async (relativePath) => {
+      const filename = String(relativePath || '').replace(/^\/+/, '');
+      if (!filename) {
+        throw new Error('Direct-source binary asset path is required.');
+      }
+      return fetchBytes(`${root}/${filename}`);
+    };
+    return createSourceStorageContext({
+      manifest,
+      readRange,
+      readText,
+      readBinary,
+      verifyHashes: true,
+    });
+  }
+
   return {
     async loadShard(index) {
       const shard = manifest.shards[index];
@@ -69,11 +122,7 @@ function createRemoteStorageContext(baseUrl, manifest) {
       if (!filename) {
         throw new Error(`Manifest shard ${index} is missing filename.`);
       }
-      const response = await fetch(`${root}/${filename.replace(/^\/+/, '')}`);
-      if (!response.ok) {
-        throw new Error(`Failed to fetch shard ${index} from ${root}: ${response.status}`);
-      }
-      return new Uint8Array(await response.arrayBuffer());
+      return fetchBytes(`${root}/${filename.replace(/^\/+/, '')}`);
     },
   };
 }
@@ -326,20 +375,29 @@ export async function initRoPEFrequencies(config, useGPU) {
   // Upload to GPU if available
   const device = getDevice();
   if (device && useGPU) {
-    const cosBuffer = acquireBuffer(globalFreqs.cos.byteLength, undefined, 'rope_cos');
-    const sinBuffer = acquireBuffer(globalFreqs.sin.byteLength, undefined, 'rope_sin');
-    device.queue.writeBuffer(cosBuffer, 0, globalFreqs.cos.buffer, globalFreqs.cos.byteOffset, globalFreqs.cos.byteLength);
-    device.queue.writeBuffer(sinBuffer, 0, globalFreqs.sin.buffer, globalFreqs.sin.byteOffset, globalFreqs.sin.byteLength);
-
-    
-    let localCosBuffer;
-    
-    let localSinBuffer;
-    if (localFreqs) {
-      localCosBuffer = acquireBuffer(localFreqs.cos.byteLength, undefined, 'rope_local_cos');
-      localSinBuffer = acquireBuffer(localFreqs.sin.byteLength, undefined, 'rope_local_sin');
-      device.queue.writeBuffer(localCosBuffer, 0, localFreqs.cos.buffer, localFreqs.cos.byteOffset, localFreqs.cos.byteLength);
-      device.queue.writeBuffer(localSinBuffer, 0, localFreqs.sin.buffer, localFreqs.sin.byteOffset, localFreqs.sin.byteLength);
+    let cosBuffer = null;
+    let sinBuffer = null;
+    let localCosBuffer = null;
+    let localSinBuffer = null;
+    try {
+      cosBuffer = acquireBuffer(globalFreqs.cos.byteLength, undefined, 'rope_cos');
+      sinBuffer = acquireBuffer(globalFreqs.sin.byteLength, undefined, 'rope_sin');
+      device.queue.writeBuffer(cosBuffer, 0, globalFreqs.cos.buffer, globalFreqs.cos.byteOffset, globalFreqs.cos.byteLength);
+      device.queue.writeBuffer(sinBuffer, 0, globalFreqs.sin.buffer, globalFreqs.sin.byteOffset, globalFreqs.sin.byteLength);
+
+      if (localFreqs) {
+        localCosBuffer = acquireBuffer(localFreqs.cos.byteLength, undefined, 'rope_local_cos');
+        localSinBuffer = acquireBuffer(localFreqs.sin.byteLength, undefined, 'rope_local_sin');
+        device.queue.writeBuffer(localCosBuffer, 0, localFreqs.cos.buffer, localFreqs.cos.byteOffset, localFreqs.cos.byteLength);
+        device.queue.writeBuffer(localSinBuffer, 0, localFreqs.sin.buffer, localFreqs.sin.byteOffset, localFreqs.sin.byteLength);
+      }
+    } catch (error) {
+      for (const buffer of [cosBuffer, sinBuffer, localCosBuffer, localSinBuffer]) {
+        if (buffer) {
+          releaseBuffer(buffer);
+        }
+      }
+      throw error;
     }
 
     log.debug(

package/src/inference/pipelines/text/kernel-path-auto-select.js

@@ -78,6 +78,7 @@ export function resolveCapabilityKernelPathRef(configuredKernelPathRef, kernelPa
 
   const normalizedPolicy = resolveKernelPathPolicy(kernelPathPolicy);
   const hasSubgroups = capabilities?.hasSubgroups === true;
+  const hasF16 = capabilities?.hasF16 === true;
   const normalizedSource = normalizeKernelPathSource(kernelPathSource);
   const allowCapabilityAutoSelection = normalizedPolicy.mode === 'capability-aware'
     && normalizedPolicy.sourceScope.includes(normalizedSource);
@@ -85,6 +86,7 @@ export function resolveCapabilityKernelPathRef(configuredKernelPathRef, kernelPa
   return selectRuleValue('inference', 'kernelPath', 'autoSelect', {
     kernelPathRef: configuredKernelPathRef,
     hasSubgroups,
+    hasF16,
     allowCapabilityAutoSelection,
   });
 }

package/src/inference/pipelines/text/kernel-trace.d.ts

@@ -12,6 +12,8 @@
  * Snapshot of a tensor's statistics (no full data, just stats).
  */
 export interface TensorSnapshot {
+  ok: boolean;
+  error: string | null;
   shape: number[];
   dtype: string;
   stats: {

package/src/inference/pipelines/text/kernel-trace.js

@@ -283,6 +283,9 @@ export async function traceStep(name, label, layer, outputBuffer, outputShape, o
   if (layer >= 0 && !kernelTrace.shouldTraceLayer(layer)) return;
 
   const output = await snapshotTensor(outputBuffer, outputShape);
+  if (!output.ok) {
+    throw new Error(`[TRACE] Failed to snapshot output for ${label}: ${output.error}`);
+  }
 
   // Snapshot inputs if provided (expensive - only do if tracing)
   
@@ -290,6 +293,9 @@ export async function traceStep(name, label, layer, outputBuffer, outputShape, o
   if (options?.inputs && options?.inputShapes) {
     for (let i = 0; i < options.inputs.length; i++) {
       const snap = await snapshotTensor(options.inputs[i], options.inputShapes[i]);
+      if (!snap.ok) {
+        throw new Error(`[TRACE] Failed to snapshot input ${i} for ${label}: ${snap.error}`);
+      }
       inputs.push(snap);
     }
   }

package/src/inference/pipelines/text/layer.js

@@ -2,7 +2,7 @@
 
 import { log, trace } from '../../../debug/index.js';
 import { getDevice } from '../../../gpu/device.js';
-import { releaseBuffer } from '../../../memory/buffer-pool.js';
+import { releaseBuffer, readBuffer } from '../../../memory/buffer-pool.js';
 import { allowReadback } from '../../../gpu/perf-guards.js';
 import { createTensor } from '../../../gpu/tensor.js';
 import {
@@ -228,6 +228,7 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       linearRuntime: context.linearAttentionRuntime ?? null,
       getWeightBuffer: (weight, label) => getWeightBuffer(weight, label),
       getNormWeightBuffer: (weight, label) => getNormWeightBuffer(weight, label, weightConfig, debugFlags),
+      debugProbes: context.debugProbes,
       recorder: recorder ?? null,
     });
   } else {
@@ -314,14 +315,7 @@ export async function processLayerGPU(layerIdx, inputBuffer, numTokens, isPrefil
       if (allowReadback(`layer.attn-out.${layerIdx}`)) {
         try {
           const sampleSize = Math.min(128, attnOutput.buffer.size);
-          const staging = device.createBuffer({ size: sampleSize, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ });
-          const enc = device.createCommandEncoder();
-          enc.copyBufferToBuffer(attnOutput.buffer, 0, staging, 0, sampleSize);
-          device.queue.submit([enc.finish()]);
-          await staging.mapAsync(GPUMapMode.READ);
-          const data = new Float32Array(staging.getMappedRange().slice(0));
-          staging.unmap();
-          staging.destroy();
+          const data = new Float32Array(await readBuffer(attnOutput.buffer, sampleSize));
           let maxAbs = 0;
           for (let i = 0; i < data.length; i++) {
             const abs = Math.abs(data[i]);

package/src/inference/pipelines/text/linear-attention.d.ts

@@ -3,6 +3,7 @@ import type { Tensor } from '../../../gpu/tensor.js';
 import type { WeightBuffer } from '../../../gpu/weight-buffer.js';
 import type { CommandRecorder } from '../../../gpu/command-recorder.js';
 import type { LinearNormMode } from '../../../config/schema/index.js';
+import type { ProbeConfigSchema } from '../../../config/schema/index.js';
 
 export interface LinearLayerRuntimeState {
   layerIdx: number;
@@ -67,6 +68,7 @@ export interface RunLinearAttentionLayerOptions {
     weight: GPUBuffer | Float32Array | ArrayBuffer,
     label: string
   ) => GPUBuffer;
+  debugProbes?: ProbeConfigSchema[] | null;
   recorder?: CommandRecorder | null;
 }
 
@@ -74,6 +76,14 @@ export declare function hasLinearAttentionLayers(layerTypes: unknown): boolean;
 
 export declare function createLinearAttentionRuntime(): LinearAttentionRuntime;
 
+export declare function inferLinearNormMode(
+  weight: { size?: number; dtype?: string } | GPUBuffer | WeightBuffer | ArrayBufferView | ArrayBuffer | null | undefined,
+  projectionLayout: {
+    headVDim: number;
+    valueDim: number;
+  }
+): LinearNormMode | null;
+
 export declare function resetLinearAttentionRuntime(
   runtime: LinearAttentionRuntime | null | undefined
 ): LinearAttentionRuntime;

package/src/inference/pipelines/text/linear-attention.js

@@ -4,6 +4,7 @@ import { readBuffer, releaseBuffer, uploadData, acquireBuffer } from '../../../m
 import { log } from '../../../debug/index.js';
 import { decodeReadback } from './debug-utils/index.js';
 import { runLinearAttentionCoreGPU } from '../../../gpu/kernels/linear-attention-core.js';
+import { runProbes } from './probes.js';
 
 const LINEAR_RUNTIME_SCHEMA_VERSION = 1;
 const QK_L2NORM_EPS = 1e-6;
@@ -173,9 +174,22 @@ function inferLinearNormModeFromWeight(weight, projectionLayout) {
   if (weight instanceof ArrayBuffer) {
     return classify(Math.trunc(weight.byteLength / Float32Array.BYTES_PER_ELEMENT));
   }
+  const explicitDtype = typeof weight?.dtype === 'string' ? weight.dtype.toLowerCase() : null;
+  const trackedDtype = isGpuBuffer(weight) ? String(getBufferDtype(weight) ?? '').toLowerCase() : '';
+  const bytesPerElement = bytesFromDtype(explicitDtype || trackedDtype || null);
+  const sizedElements = Number.isFinite(weight?.size)
+    ? Math.trunc(Number(weight.size) / bytesPerElement)
+    : null;
+  if (sizedElements && Number(weight.size) % bytesPerElement === 0) {
+    return classify(sizedElements);
+  }
   return null;
 }
 
+export function inferLinearNormMode(weight, projectionLayout) {
+  return inferLinearNormModeFromWeight(weight, projectionLayout);
+}
+
 function resolveLinearNormMode(configNormMode, normWeight, projectionLayout, layerIdx) {
   const configuredMode = normalizeLinearNormMode(configNormMode);
   const inferredMode = inferLinearNormModeFromWeight(normWeight, projectionLayout);
@@ -185,7 +199,15 @@ function resolveLinearNormMode(configNormMode, normWeight, projectionLayout, lay
       `but norm.weight shape implies "${inferredMode}".`
     );
   }
-  return configuredMode ?? inferredMode ?? 'shared';
+  if (configuredMode) {
+    return configuredMode;
+  }
+  if (inferredMode) {
+    return inferredMode;
+  }
+  throw new Error(
+    `linear_attention layer ${layerIdx} requires explicit linearNormMode or a norm.weight shape that resolves it.`
+  );
 }
 
 async function readWeightAsF32(weight, expectedElements, label) {
@@ -395,10 +417,17 @@ async function createLayerRuntimeState(
 
   let convKernelSize = toPositiveInt(config.linearConvKernelDim) ?? null;
   if (isWeightBuffer(convKernel) && Array.isArray(convKernel.shape) && convKernel.shape.length >= 3) {
-    convKernelSize = toPositiveInt(convKernel.shape[2]) ?? convKernelSize;
+    const shapeKernelSize = toPositiveInt(convKernel.shape[2]) ?? null;
+    if (convKernelSize != null && shapeKernelSize != null && convKernelSize !== shapeKernelSize) {
+      throw new Error(
+        `linear_attention layer ${layerIdx} declares linearConvKernelDim=${convKernelSize}, ` +
+        `but conv1d weight shape implies ${shapeKernelSize}.`
+      );
+    }
+    convKernelSize = shapeKernelSize ?? convKernelSize;
   }
   if (!convKernelSize) {
-    convKernelSize = 4;
+    throw new Error(`linear_attention layer ${layerIdx} requires linearConvKernelDim.`);
   }
 
   const convWeight = await readWeightAsF32(
@@ -435,6 +464,11 @@ async function createLayerRuntimeState(
   const recurrentState = new Float32Array(
     projectionLayout.numVHeads * projectionLayout.headKDim * projectionLayout.headVDim
   );
+  const rmsNormEps = Number(config.rmsNormEps);
+  if (!Number.isFinite(rmsNormEps) || rmsNormEps <= 0) {
+    throw new Error(`linear_attention layer ${layerIdx} requires a positive rmsNormEps.`);
+  }
+
   const layerState = {
     layerIdx,
     seqLen: currentSeqLen,
@@ -452,7 +486,7 @@ async function createLayerRuntimeState(
     vSize: projectionLayout.vSize,
     qRep: projectionLayout.qRep,
     normMode,
-    rmsNormEps: Number(config.rmsNormEps) || 1e-6,
+    rmsNormEps,
     convWeight,
     dtBias,
     aNegExp,
@@ -681,13 +715,13 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
     const normWeightBuffer = getNormWeightBuffer(layerWeights.inputNorm, `L${layerIdx}.linear_input_norm`);
     try {
       if (recorder) {
-        normedTensor = await recordRMSNorm(recorder, inputTensor, normWeightBuffer, Number(config.rmsNormEps) || 1e-6, {
+        normedTensor = await recordRMSNorm(recorder, inputTensor, normWeightBuffer, layerState.rmsNormEps, {
           batchSize: numTokens,
           hiddenSize,
           rmsNormWeightOffset: config.rmsNormWeightOffset,
         });
       } else {
-        normedTensor = await runRMSNorm(inputTensor, normWeightBuffer, Number(config.rmsNormEps) || 1e-6, {
+        normedTensor = await runRMSNorm(inputTensor, normWeightBuffer, layerState.rmsNormEps, {
           batchSize: numTokens,
           hiddenSize,
           rmsNormWeightOffset: config.rmsNormWeightOffset,
@@ -755,6 +789,38 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
   });
 
   try {
+    await runProbes('linear_qkv_proj', qkvTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.convDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: qkvTensor.dtype,
+    });
+    await runProbes('linear_z_proj', zTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.valueDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: zTensor.dtype,
+    });
+    await runProbes('linear_a_proj', aTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.numVHeads,
+      probes: options.debugProbes,
+      recorder,
+      dtype: aTensor.dtype,
+    });
+    await runProbes('linear_b_proj', bTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.numVHeads,
+      probes: options.debugProbes,
+      recorder,
+      dtype: bTensor.dtype,
+    });
     const coreTensor = await runLinearAttentionCoreGPU(
       qkvTensor,
       zTensor,
@@ -768,6 +834,14 @@ export async function runLinearAttentionLayer(inputTensor, layerWeights, options
         recorder,
       }
     );
+    await runProbes('linear_core_out', coreTensor.buffer, {
+      layerIdx,
+      numTokens,
+      hiddenSize: projectionLayout.valueDim,
+      probes: options.debugProbes,
+      recorder,
+      dtype: coreTensor.dtype,
+    });
     layerState.seqLen = currentSeqLen + numTokens;
     const outProjWeight = getWeightBuffer(layerWeights.oProj, `L${layerIdx}.linear_out_proj`);
     try {

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;
 }