@simulatte/doppler 0.1.5 → 0.1.7

package/src/inference/pipelines/text/moe-cpu-gptoss.js

@@ -30,81 +30,86 @@ export async function runGptOssExpertCPU(layerIdx, expertIdx, input, config, exp
   const inputBuffer = acquireBuffer(input.byteLength, undefined, 'moe_cpu_gptoss_input');
   device.queue.writeBuffer(inputBuffer, 0, input);
   const inputTensor = createTensor(inputBuffer, 'f32', [numTokens, hiddenSize], 'moe_cpu_gptoss_input');
+  let gateUpWeight = null;
+  let downWeight = null;
+  let gateUpOut = null;
+  let activated = null;
+  let downOut = null;
+  try {
+    gateUpWeight = await dequantizeMXFP4Expert(
+      weights.gateUpBlocks,
+      weights.gateUpScales,
+      expertIdx,
+      numExperts,
+      outDim,
+      gateUpGroups,
+      { outputDtype: 'f32', modelType: 'gpt-oss', groupSize: 32, dequantTileShape: 'scalar' }
+    );
 
-  const gateUpWeight = await dequantizeMXFP4Expert(
-    weights.gateUpBlocks,
-    weights.gateUpScales,
-    expertIdx,
-    numExperts,
-    outDim,
-    gateUpGroups,
-    { outputDtype: 'f32', modelType: 'gpt-oss', groupSize: 32, dequantTileShape: 'scalar' }
-  );
+    downWeight = await dequantizeMXFP4Expert(
+      weights.downBlocks,
+      weights.downScales,
+      expertIdx,
+      numExperts,
+      hiddenSize,
+      downGroups,
+      { outputDtype: 'f32', modelType: 'gpt-oss', groupSize: 32, dequantTileShape: 'scalar' }
+    );
 
-  const downWeight = await dequantizeMXFP4Expert(
-    weights.downBlocks,
-    weights.downScales,
-    expertIdx,
-    numExperts,
-    hiddenSize,
-    downGroups,
-    { outputDtype: 'f32', modelType: 'gpt-oss', groupSize: 32, dequantTileShape: 'scalar' }
-  );
+    gateUpOut = await runMatmul(
+      inputTensor,
+      gateUpWeight.buffer,
+      numTokens,
+      outDim,
+      hiddenSize,
+      {
+        transposeB: 'auto',
+        bDtype: 'f32',
+        outputDtype: 'f32',
+        role: 'moe_gate_up',
+        kernelPath,
+      }
+    );
 
-  const gateUpOut = await runMatmul(
-    inputTensor,
-    gateUpWeight.buffer,
-    numTokens,
-    outDim,
-    hiddenSize,
-    {
-      transposeB: 'auto',
-      bDtype: 'f32',
-      outputDtype: 'f32',
-      role: 'moe_gate_up',
-      kernelPath,
-    }
-  );
+    const biasOffset = expertIdx * outDim * 4;
+    const biasTensor = createTensor(weights.gateUpBias, 'f32', [numExperts * outDim], 'moe_cpu_gptoss_bias');
+    activated = await runSwiGLURowsplitBias(gateUpOut, biasTensor, numTokens, intermediateSize, {
+      biasOffset,
+      swigluLimit,
+    });
 
-  const biasOffset = expertIdx * outDim * 4;
-  const biasTensor = createTensor(weights.gateUpBias, 'f32', [numExperts * outDim], 'moe_cpu_gptoss_bias');
-  const activated = await runSwiGLURowsplitBias(gateUpOut, biasTensor, numTokens, intermediateSize, {
-    biasOffset,
-    swigluLimit,
-  });
+    downOut = await runMatmul(
+      activated,
+      downWeight.buffer,
+      numTokens,
+      hiddenSize,
+      intermediateSize,
+      {
+        transposeB: 'auto',
+        bDtype: 'f32',
+        outputDtype: 'f32',
+        role: 'moe_down',
+        kernelPath,
+      }
+    );
 
-  const downOut = await runMatmul(
-    activated,
-    downWeight.buffer,
-    numTokens,
-    hiddenSize,
-    intermediateSize,
-    {
-      transposeB: 'auto',
-      bDtype: 'f32',
-      outputDtype: 'f32',
-      role: 'moe_down',
-      kernelPath,
+    if (weights.downBias) {
+      const downBiasOffset = expertIdx * hiddenSize * 4;
+      const downBiasTensor = createTensor(weights.downBias, 'f32', [numExperts * hiddenSize], 'moe_cpu_gptoss_down_bias');
+      await runBiasAdd(downOut, downBiasTensor, numTokens, hiddenSize, {
+        dataOffset: 0,
+        biasOffset: downBiasOffset,
+      });
     }
-  );
 
-  if (weights.downBias) {
-    const downBiasOffset = expertIdx * hiddenSize * 4;
-    const downBiasTensor = createTensor(weights.downBias, 'f32', [numExperts * hiddenSize], 'moe_cpu_gptoss_down_bias');
-    await runBiasAdd(downOut, downBiasTensor, numTokens, hiddenSize, {
-      dataOffset: 0,
-      biasOffset: downBiasOffset,
-    });
+    const outputData = await readBuffer(downOut.buffer, input.byteLength);
+    return new Float32Array(outputData);
+  } finally {
+    releaseBuffer(inputBuffer);
+    if (gateUpWeight) releaseBuffer(gateUpWeight.buffer);
+    if (downWeight) releaseBuffer(downWeight.buffer);
+    if (gateUpOut) releaseBuffer(gateUpOut.buffer);
+    if (activated) releaseBuffer(activated.buffer);
+    if (downOut) releaseBuffer(downOut.buffer);
   }
-
-  const outputData = await readBuffer(downOut.buffer, input.byteLength);
-
-  releaseBuffer(inputBuffer);
-  releaseBuffer(gateUpWeight.buffer);
-  releaseBuffer(downWeight.buffer);
-  releaseBuffer(gateUpOut.buffer);
-  releaseBuffer(activated.buffer);
-  releaseBuffer(downOut.buffer);
-
-  return new Float32Array(outputData);
 }

package/src/inference/pipelines/text/moe-cpu.js

@@ -75,42 +75,46 @@ async function runExpertCPU(layerIdx, expertIdx, input, config, expertWeights) {
   const inputBuffer = acquireBuffer(input.byteLength, undefined, 'expert_input');
   device.queue.writeBuffer(inputBuffer, 0, input);
   const inputTensor = createTensor(inputBuffer, 'f32', [numTokens, hiddenSize], 'expert_input');
-
-  const gateOutput = await runMatmul(inputTensor, weights.gate, numTokens, intermediateSize, hiddenSize, {
-    transposeB: 'auto',
-    role: 'moe_gate',
-    kernelPath,
-  });
-
-  const upOutput = await runMatmul(inputTensor, weights.up, numTokens, intermediateSize, hiddenSize, {
-    transposeB: 'auto',
-    role: 'moe_up',
-    kernelPath,
-  });
-
-  const activationFn = {
-    gelu: runGeLU,
-    silu: runSiLU,
-  }[selectRuleValue('inference', 'ffn', 'activationOp', { hiddenActivation })];
-  const activatedOutput = await activationFn(upOutput, {
-    size: numTokens * intermediateSize,
-    gate: gateOutput,
-    swigluLimit,
-  });
-
-  const output = await runMatmul(activatedOutput, weights.down, numTokens, hiddenSize, intermediateSize, {
-    transposeB: 'auto',
-    role: 'moe_down',
-    kernelPath,
-  });
-
-  const outputData = await readBuffer(output.buffer, input.byteLength);
-
-  releaseBuffer(inputBuffer);
-  releaseBuffer(gateOutput.buffer);
-  releaseBuffer(upOutput.buffer);
-  releaseBuffer(activatedOutput.buffer);
-  releaseBuffer(output.buffer);
-
-  return new Float32Array(outputData);
+  let gateOutput = null;
+  let upOutput = null;
+  let activatedOutput = null;
+  let output = null;
+  try {
+    gateOutput = await runMatmul(inputTensor, weights.gate, numTokens, intermediateSize, hiddenSize, {
+      transposeB: 'auto',
+      role: 'moe_gate',
+      kernelPath,
+    });
+
+    upOutput = await runMatmul(inputTensor, weights.up, numTokens, intermediateSize, hiddenSize, {
+      transposeB: 'auto',
+      role: 'moe_up',
+      kernelPath,
+    });
+
+    const activationFn = {
+      gelu: runGeLU,
+      silu: runSiLU,
+    }[selectRuleValue('inference', 'ffn', 'activationOp', { hiddenActivation })];
+    activatedOutput = await activationFn(upOutput, {
+      size: numTokens * intermediateSize,
+      gate: gateOutput,
+      swigluLimit,
+    });
+
+    output = await runMatmul(activatedOutput, weights.down, numTokens, hiddenSize, intermediateSize, {
+      transposeB: 'auto',
+      role: 'moe_down',
+      kernelPath,
+    });
+
+    const outputData = await readBuffer(output.buffer, input.byteLength);
+    return new Float32Array(outputData);
+  } finally {
+    releaseBuffer(inputBuffer);
+    if (gateOutput) releaseBuffer(gateOutput.buffer);
+    if (upOutput) releaseBuffer(upOutput.buffer);
+    if (activatedOutput) releaseBuffer(activatedOutput.buffer);
+    if (output) releaseBuffer(output.buffer);
+  }
 }

package/src/inference/pipelines/text/moe-gpu.js

@@ -81,17 +81,27 @@ export async function moeFeedForwardGPU(
   };
 
   const inputTensor = createTensor(inputBuffer, activationDtype, [numTokens, hiddenSize], 'moe_input');
+  let logitsBuffer = null;
+  let indicesBuffer = null;
+  let weightsBuffer = null;
+  let gathered = null;
+  let tokenCounts = null;
+  let tokenMap = null;
+  let expertOutputs = null;
+  let tokenOffsets = null;
+  let outputTensor = null;
 
   const layerRouter = layerRouterWeights?.get(layerIdx) || null;
   if (layerRouter) {
     moeRouter.loadWeights(layerRouter.weight, layerRouter.bias || null);
   }
 
-  let stepStart = perfMark();
-  const logitsBuffer = await moeRouter.computeRouterLogitsGPU(inputTensor.buffer, numTokens, null, {
-    inputDtype: activationDtype,
-    outputDtype: activationDtype,
-  });
+  try {
+    let stepStart = perfMark();
+    logitsBuffer = await moeRouter.computeRouterLogitsGPU(inputTensor.buffer, numTokens, null, {
+      inputDtype: activationDtype,
+      outputDtype: activationDtype,
+    });
   const logitsDtype = moeRouter.lastLogitsDtype ?? activationDtype;
   perfLog(`MoE L${layerIdx} router`, stepStart, { numTokens, logitsDtype });
 
@@ -134,18 +144,18 @@ export async function moeFeedForwardGPU(
   }
 
   stepStart = perfMark();
-  const { indices: indicesBuffer, weights: weightsBuffer } = await runSoftmaxTopK(
-    logitsBuffer,
-    numTokens,
-    numExperts,
-    topK,
-    {
-      normalize: moeRouter.normalizeWeights,
-      inputDtype: logitsDtype,
-      weightsDtype: activationDtype,
-      modelType,
-    }
-  );
+    ({ indices: indicesBuffer, weights: weightsBuffer } = await runSoftmaxTopK(
+      logitsBuffer,
+      numTokens,
+      numExperts,
+      topK,
+      {
+        normalize: moeRouter.normalizeWeights,
+        inputDtype: logitsDtype,
+        weightsDtype: activationDtype,
+        modelType,
+      }
+    ));
   perfLog(`MoE L${layerIdx} topk`, stepStart, {
     topK,
     modelType,
@@ -195,7 +205,8 @@ export async function moeFeedForwardGPU(
     trace.buffers(`MoE L${layerIdx} topk_weights`, { minW, maxW, nanW, dtype: activationDtype });
   }
 
-  releaseBuffer(logitsBuffer);
+    releaseBuffer(logitsBuffer);
+    logitsBuffer = null;
 
   const bytesPerElement = selectRuleValue('shared', 'dtype', 'bytesFromDtype', { dtype: activationDtype });
   const bytesPerToken = hiddenSize * bytesPerElement;
@@ -207,47 +218,44 @@ export async function moeFeedForwardGPU(
     );
   }
 
-  let gathered;
-  let tokenCounts;
-  let tokenMap;
-
-  stepStart = perfMark();
-  ({ gathered, tokenCounts, tokenMap } = await runMoEGather(
-    inputTensor,
-    indicesBuffer,
-    numTokens,
-    hiddenSize,
-    numExperts,
-    topK,
-    { maxTokensPerExpert }
-  ));
+    stepStart = perfMark();
+    ({ gathered, tokenCounts, tokenMap } = await runMoEGather(
+      inputTensor,
+      indicesBuffer,
+      numTokens,
+      hiddenSize,
+      numExperts,
+      topK,
+      { maxTokensPerExpert }
+    ));
   perfLog(`MoE L${layerIdx} gather`, stepStart, { maxTokensPerExpert });
 
-  const expertOutputs = acquireBuffer(
-    numExperts * maxTokensPerExpert * hiddenSize * bytesPerElement,
-    undefined,
-    'moe_expert_outputs_gathered'
-  );
+    expertOutputs = acquireBuffer(
+      numExperts * maxTokensPerExpert * hiddenSize * bytesPerElement,
+      undefined,
+      'moe_expert_outputs_gathered'
+    );
 
   const zeroEncoder = device.createCommandEncoder({ label: 'zero_moe_expert_outputs' });
   zeroEncoder.clearBuffer(expertOutputs, 0, numExperts * maxTokensPerExpert * hiddenSize * bytesPerElement);
   device.queue.submit([zeroEncoder.finish()]);
 
-  stepStart = perfMark();
-  const tokenOffsets = await runMoEBuildTokenOffsets(
-    tokenCounts,
-    tokenMap,
-    numTokens,
-    numExperts,
-    topK,
-    maxTokensPerExpert
-  );
+    stepStart = perfMark();
+    tokenOffsets = await runMoEBuildTokenOffsets(
+      tokenCounts,
+      tokenMap,
+      numTokens,
+      numExperts,
+      topK,
+      maxTokensPerExpert
+    );
   perfLog(`MoE L${layerIdx} offsets_kernel`, stepStart, {
     totalSlots: numExperts * maxTokensPerExpert,
     routingSlots: numTokens * topK,
   });
 
-  releaseBuffer(tokenCounts);
+    releaseBuffer(tokenCounts);
+    tokenCounts = null;
 
   const expertStrideBytes = maxTokensPerExpert * bytesPerToken;
   const rowsPerExpert = maxTokensPerExpert;
@@ -322,49 +330,65 @@ export async function moeFeedForwardGPU(
     perfLog(`MoE L${layerIdx} expert_exec`, stepStart, { expertIdx, count });
   }
 
-  const expertOutputsTensor = createTensor(
-    expertOutputs,
-    activationDtype,
-    [numExperts, maxTokensPerExpert, hiddenSize],
-    'moe_expert_outputs'
-  );
-  stepStart = perfMark();
-  const outputTensor = await runScatterAddDynamic(
-    expertOutputsTensor,
-    indicesBuffer,
-    weightsBuffer,
-    tokenOffsets,
-    numTokens,
-    hiddenSize,
-    topK,
-    { weightsDtype: activationDtype }
-  );
+    const expertOutputsTensor = createTensor(
+      expertOutputs,
+      activationDtype,
+      [numExperts, maxTokensPerExpert, hiddenSize],
+      'moe_expert_outputs'
+    );
+    stepStart = perfMark();
+    outputTensor = await runScatterAddDynamic(
+      expertOutputsTensor,
+      indicesBuffer,
+      weightsBuffer,
+      tokenOffsets,
+      numTokens,
+      hiddenSize,
+      topK,
+      { weightsDtype: activationDtype }
+    );
   perfLog(`MoE L${layerIdx} scatter`, stepStart, { numTokens, hiddenSize });
 
-  releaseBuffer(gathered.buffer);
-  releaseBuffer(tokenMap);
-  releaseBuffer(expertOutputs);
-  releaseBuffer(tokenOffsets);
-  releaseBuffer(indicesBuffer);
-  releaseBuffer(weightsBuffer);
+    releaseBuffer(gathered.buffer);
+    gathered = null;
+    releaseBuffer(tokenMap);
+    tokenMap = null;
+    releaseBuffer(expertOutputs);
+    expertOutputs = null;
+    releaseBuffer(tokenOffsets);
+    tokenOffsets = null;
+    releaseBuffer(indicesBuffer);
+    indicesBuffer = null;
+    releaseBuffer(weightsBuffer);
+    weightsBuffer = null;
+
+    if (perfEnabled) {
+      const cacheStats = getDequantCacheStats();
+      trace.perf(`MoE L${layerIdx} done`, {
+        numTokens,
+        topK,
+        executedExperts: numExperts,
+        rowsPerExpert,
+        maxTokensPerExpert,
+        dequantCacheHits: cacheStats.hits,
+        dequantCacheMisses: cacheStats.misses,
+        expertCache: typeof expertLoader?.getExpertCacheStats === 'function'
+          ? expertLoader.getExpertCacheStats()
+          : null,
+      });
+    }
 
-  if (perfEnabled) {
-    const cacheStats = getDequantCacheStats();
-    trace.perf(`MoE L${layerIdx} done`, {
-      numTokens,
-      topK,
-      executedExperts: numExperts,
-      rowsPerExpert,
-      maxTokensPerExpert,
-      dequantCacheHits: cacheStats.hits,
-      dequantCacheMisses: cacheStats.misses,
-      expertCache: typeof expertLoader?.getExpertCacheStats === 'function'
-        ? expertLoader.getExpertCacheStats()
-        : null,
-    });
+    return outputTensor.buffer;
+  } finally {
+    if (logitsBuffer) releaseBuffer(logitsBuffer);
+    if (tokenCounts) releaseBuffer(tokenCounts);
+    if (gathered?.buffer) releaseBuffer(gathered.buffer);
+    if (tokenMap) releaseBuffer(tokenMap);
+    if (expertOutputs) releaseBuffer(expertOutputs);
+    if (tokenOffsets) releaseBuffer(tokenOffsets);
+    if (indicesBuffer) releaseBuffer(indicesBuffer);
+    if (weightsBuffer) releaseBuffer(weightsBuffer);
   }
-
-  return outputTensor.buffer;
 }
 
 function inferBufferDtype(buffer, expectedElements) {