@simulatte/doppler 0.1.6 → 0.1.7

package/src/inference/kv-cache/base.js

@@ -3,6 +3,7 @@
 import { getDevice } from '../../gpu/device.js';
 import { allowReadback } from '../../gpu/perf-guards.js';
 import { log } from '../../debug/index.js';
+import { readBuffer } from '../../memory/buffer-pool.js';
 import {
   isContiguousLayer,
   isPagedLayer,
@@ -815,8 +816,52 @@ export class KVCache {
     }
   }
 
+  _destroyGpuBuffer(buffer) {
+    if (!buffer) return;
+    try {
+      buffer.destroy();
+    } catch {
+      // Ignore already-destroyed buffers during rollback.
+    }
+  }
+
+  _snapshotLayerGpuState(layer) {
+    return {
+      keysGPU: layer.keysGPU ?? null,
+      valuesGPU: layer.valuesGPU ?? null,
+      pageTableGPU: layer.pageTableGPU ?? null,
+    };
+  }
+
+  _rollbackMigratedLayers(snapshots) {
+    for (let l = 0; l < this.numLayers; l++) {
+      const layer = this.layers[l];
+      const snapshot = snapshots[l];
+      if (!snapshot) continue;
+
+      if (layer.keysGPU && layer.keysGPU !== snapshot.keysGPU) {
+        this._destroyGpuBuffer(layer.keysGPU);
+      }
+      if (layer.valuesGPU && layer.valuesGPU !== snapshot.valuesGPU) {
+        this._destroyGpuBuffer(layer.valuesGPU);
+      }
+      if (layer.pageTableGPU && layer.pageTableGPU !== snapshot.pageTableGPU) {
+        this._destroyGpuBuffer(layer.pageTableGPU);
+      }
+
+      layer.keysGPU = snapshot.keysGPU;
+      layer.valuesGPU = snapshot.valuesGPU;
+      if ('pageTableGPU' in layer) {
+        layer.pageTableGPU = snapshot.pageTableGPU;
+      }
+    }
+  }
+
   
   _migrateToGPU(device) {
+    const snapshots = this.layers.map((layer) => this._snapshotLayerGpuState(layer));
+
+    try {
     if (this.layout === 'paged') {
       log.info('KVCache', `Migrating ${this.currentSeqLen} positions to GPU (paged)...`);
       const numPages = Math.ceil(this.maxSeqLen / this.pageSize);
@@ -826,56 +871,66 @@ export class KVCache {
 
       for (let l = 0; l < this.numLayers; l++) {
         const layer =  (this.layers[l]);
-
-        if (!layer.keysGPU) {
-          layer.keysGPU = device.createBuffer({
+        let keysGPU = null;
+        let valuesGPU = null;
+        let pageTableGPU = null;
+        try {
+          keysGPU = device.createBuffer({
             label: `kv_cache_keys_paged_layer_${l}`,
             size: bytesPerLayer,
             usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
           });
-        }
-        if (!layer.valuesGPU) {
-          layer.valuesGPU = device.createBuffer({
+          valuesGPU = device.createBuffer({
             label: `kv_cache_values_paged_layer_${l}`,
             size: bytesPerLayer,
             usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
           });
-        }
 
-        if (!layer.pageTable) {
-          layer.pageTable = new Uint32Array(numPages);
-          for (let i = 0; i < numPages; i++) {
-            layer.pageTable[i] = i;
+          if (!layer.pageTable) {
+            layer.pageTable = new Uint32Array(numPages);
+            for (let i = 0; i < numPages; i++) {
+              layer.pageTable[i] = i;
+            }
           }
-        }
-        if (!layer.pageTableGPU) {
-          layer.pageTableGPU = device.createBuffer({
+          pageTableGPU = device.createBuffer({
             label: `kv_cache_page_table_layer_${l}`,
             size: pageTableBytes,
             usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
           });
-        }
-        device.queue.writeBuffer(layer.pageTableGPU, 0, layer.pageTable);
-
-        const allocatedPages = layer.allocatedPages ?? 0;
-        if (allocatedPages > 0) {
-          const pageElems = this.pageSize * this.kvSize;
-          const pageBytes = pageElems * this.bytesPerElem;
-          for (let p = 0; p < allocatedPages; p++) {
-            const keyPage = layer.keyPages?.[p];
-            const valuePage = layer.valuePages?.[p];
-            if (!keyPage || !valuePage) continue;
-            const byteOffset = p * pageBytes;
-            if (this.kvDtype === 'f16') {
-              const keysF16 = f32ToF16Array(keyPage);
-              const valuesF16 = f32ToF16Array(valuePage);
-              device.queue.writeBuffer(layer.keysGPU, byteOffset, keysF16);
-              device.queue.writeBuffer(layer.valuesGPU, byteOffset, valuesF16);
-            } else {
-              device.queue.writeBuffer(layer.keysGPU, byteOffset, keyPage);
-              device.queue.writeBuffer(layer.valuesGPU, byteOffset, valuePage);
+          device.queue.writeBuffer(pageTableGPU, 0, layer.pageTable);
+
+          const allocatedPages = layer.allocatedPages ?? 0;
+          if (allocatedPages > 0) {
+            const pageElems = this.pageSize * this.kvSize;
+            const pageBytes = pageElems * this.bytesPerElem;
+            for (let p = 0; p < allocatedPages; p++) {
+              const keyPage = layer.keyPages?.[p];
+              const valuePage = layer.valuePages?.[p];
+              if (!keyPage || !valuePage) continue;
+              const byteOffset = p * pageBytes;
+              if (this.kvDtype === 'f16') {
+                const keysF16 = f32ToF16Array(keyPage);
+                const valuesF16 = f32ToF16Array(valuePage);
+                device.queue.writeBuffer(keysGPU, byteOffset, keysF16);
+                device.queue.writeBuffer(valuesGPU, byteOffset, valuesF16);
+              } else {
+                device.queue.writeBuffer(keysGPU, byteOffset, keyPage);
+                device.queue.writeBuffer(valuesGPU, byteOffset, valuePage);
+              }
             }
           }
+
+          this._destroyGpuBuffer(layer.keysGPU);
+          this._destroyGpuBuffer(layer.valuesGPU);
+          this._destroyGpuBuffer(layer.pageTableGPU);
+          layer.keysGPU = keysGPU;
+          layer.valuesGPU = valuesGPU;
+          layer.pageTableGPU = pageTableGPU;
+        } catch (error) {
+          this._destroyGpuBuffer(keysGPU);
+          this._destroyGpuBuffer(valuesGPU);
+          this._destroyGpuBuffer(pageTableGPU);
+          throw error;
         }
       }
 
@@ -890,53 +945,64 @@ export class KVCache {
 
     for (let l = 0; l < this.numLayers; l++) {
       const layer =  (this.layers[l]);
-
-      // Create GPU buffers if they don't exist
-      if (!layer.keysGPU) {
-        layer.keysGPU = device.createBuffer({
+      let keysGPU = null;
+      let valuesGPU = null;
+      try {
+        keysGPU = device.createBuffer({
           label: `kv_cache_keys_layer_${l}`,
           size: bytesPerLayer,
           usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
         });
-      }
-      if (!layer.valuesGPU) {
-        layer.valuesGPU = device.createBuffer({
+        valuesGPU = device.createBuffer({
           label: `kv_cache_values_layer_${l}`,
           size: bytesPerLayer,
           usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
         });
-      }
 
-      // Upload existing CPU data to GPU
-      const usedElems = layer.seqLen * this.kvSize;
-      const usedSize = usedElems * this.bytesPerElem;
-      if (usedSize > 0) {
-        if (this.kvDtype === 'f16') {
-          const keysF16 = f32ToF16Array(layer.keys.subarray(0, usedElems));
-          const valuesF16 = f32ToF16Array(layer.values.subarray(0, usedElems));
-          device.queue.writeBuffer(layer.keysGPU, 0, keysF16);
-          device.queue.writeBuffer(layer.valuesGPU, 0, valuesF16);
-        } else {
-          device.queue.writeBuffer(
-            layer.keysGPU,
-            0,
-            layer.keys.buffer,
-            layer.keys.byteOffset,
-            usedSize
-          );
-          device.queue.writeBuffer(
-            layer.valuesGPU,
-            0,
-            layer.values.buffer,
-            layer.values.byteOffset,
-            usedSize
-          );
+        // Upload existing CPU data to GPU
+        const usedElems = layer.seqLen * this.kvSize;
+        const usedSize = usedElems * this.bytesPerElem;
+        if (usedSize > 0) {
+          if (this.kvDtype === 'f16') {
+            const keysF16 = f32ToF16Array(layer.keys.subarray(0, usedElems));
+            const valuesF16 = f32ToF16Array(layer.values.subarray(0, usedElems));
+            device.queue.writeBuffer(keysGPU, 0, keysF16);
+            device.queue.writeBuffer(valuesGPU, 0, valuesF16);
+          } else {
+            device.queue.writeBuffer(
+              keysGPU,
+              0,
+              layer.keys.buffer,
+              layer.keys.byteOffset,
+              usedSize
+            );
+            device.queue.writeBuffer(
+              valuesGPU,
+              0,
+              layer.values.buffer,
+              layer.values.byteOffset,
+              usedSize
+            );
+          }
         }
+
+        this._destroyGpuBuffer(layer.keysGPU);
+        this._destroyGpuBuffer(layer.valuesGPU);
+        layer.keysGPU = keysGPU;
+        layer.valuesGPU = valuesGPU;
+      } catch (error) {
+        this._destroyGpuBuffer(keysGPU);
+        this._destroyGpuBuffer(valuesGPU);
+        throw error;
       }
     }
 
     this.useGPU = true;
     log.info('KVCache', 'Migration complete');
+    } catch (error) {
+      this._rollbackMigratedLayers(snapshots);
+      throw error;
+    }
   }
 
   
@@ -962,44 +1028,24 @@ export class KVCache {
         layer.values = new Float32Array(sizePerLayer);
       }
 
-      // Create staging buffers for readback
-      const keysStaging = device.createBuffer({
-        size: usedSize,
-        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
-      });
-      const valuesStaging = device.createBuffer({
-        size: usedSize,
-        usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
-      });
-
-      // Copy from GPU cache to staging
-      const encoder = device.createCommandEncoder({ label: 'kv_cache_sync' });
-      encoder.copyBufferToBuffer(layer.keysGPU, 0, keysStaging, 0, usedSize);
-      encoder.copyBufferToBuffer(layer.valuesGPU, 0, valuesStaging, 0, usedSize);
-      device.queue.submit([encoder.finish()]);
-
-      // Map and copy to CPU arrays
-      await keysStaging.mapAsync(GPUMapMode.READ);
-      await valuesStaging.mapAsync(GPUMapMode.READ);
+      const [keysBytes, valuesBytes] = await Promise.all([
+        readBuffer(layer.keysGPU, usedSize),
+        readBuffer(layer.valuesGPU, usedSize),
+      ]);
 
       if (this.kvDtype === 'f16') {
-        const keysRaw = new Uint16Array(keysStaging.getMappedRange().slice(0));
-        const valuesRaw = new Uint16Array(valuesStaging.getMappedRange().slice(0));
+        const keysRaw = new Uint16Array(keysBytes);
+        const valuesRaw = new Uint16Array(valuesBytes);
         const keysData = f16ToF32Array(keysRaw);
         const valuesData = f16ToF32Array(valuesRaw);
         layer.keys.set(keysData);
         layer.values.set(valuesData);
       } else {
-        const keysData = new Float32Array(keysStaging.getMappedRange().slice(0));
-        const valuesData = new Float32Array(valuesStaging.getMappedRange().slice(0));
+        const keysData = new Float32Array(keysBytes);
+        const valuesData = new Float32Array(valuesBytes);
         layer.keys.set(keysData);
         layer.values.set(valuesData);
       }
-
-      keysStaging.unmap();
-      valuesStaging.unmap();
-      keysStaging.destroy();
-      valuesStaging.destroy();
     }
   }
 

package/src/inference/kv-cache/tiered.js

@@ -60,7 +60,7 @@ export class TieredKVCache {
       : (tiering.mode === 'int4' ? 'int4' : 'none');
     this.compression = tiering.compression ?? { mode: defaultCompressionMode, blockSize: 1 };
     
-    this.gating = tiering.gating ?? { mode: 'auto', minAluBwRatio: 0.0 };
+    this.gating = tiering.gating ?? { mode: 'force_off', minAluBwRatio: 0.0 };
     
     this.currentSeqLen = 0;
     
@@ -145,8 +145,10 @@ export class TieredKVCache {
     if (gating?.mode === 'force_off') return 'none';
     if (gating?.mode === 'force_on') return requested;
     if (gating?.mode === 'auto' && gating.minAluBwRatio > 0) {
-      const ratio = 1.0;
-      if (ratio < gating.minAluBwRatio) return 'none';
+      throw new Error(
+        'TieredKVCache auto compression gating requires an explicit measured ALU/BW ratio. ' +
+        'Use gating.mode="force_on"/"force_off" or set minAluBwRatio to 0.'
+      );
     }
     return requested;
   }

package/src/inference/moe-router.js

@@ -8,6 +8,9 @@ import { createTensor } from '../gpu/tensor.js';
 import { f16ToF32Array } from './kv-cache/types.js';
 import { selectRuleValue } from '../rules/rule-registry.js';
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
 
 
 
@@ -84,6 +87,12 @@ export class MoERouter {
 
   
   loadWeights(weights, bias = null) {
+    if (this._gateBiasGPU) {
+      this._gateBiasGPU.destroy();
+    }
+    if (this._gateWeightGPU) {
+      this._gateWeightGPU.destroy();
+    }
     this.gateWeight = weights;
     this.gateBias = bias;
     // Clear cached GPU uploads when swapping router parameters (e.g., per-layer routers).
@@ -91,13 +100,27 @@ export class MoERouter {
     this._gateWeightGPU = null;
   }
 
+  destroy() {
+    if (isGpuBufferInstance(this._gateBiasGPU)) {
+      this._gateBiasGPU.destroy();
+    }
+    if (isGpuBufferInstance(this._gateWeightGPU)) {
+      this._gateWeightGPU.destroy();
+    }
+    this._gateBiasGPU = null;
+    this._gateWeightGPU = null;
+    this.gateWeight = null;
+    this.gateBias = null;
+    this._biasAddPipelines.clear();
+  }
+
   
   computeRouterLogitsCPU(hiddenStates, numTokens) {
     if (!this.gateWeight) {
       throw new Error('Router gate weights not loaded');
     }
 
-    if (this.gateWeight instanceof GPUBuffer || isWeightBuffer(this.gateWeight)) {
+    if (isGpuBufferInstance(this.gateWeight) || isWeightBuffer(this.gateWeight)) {
       throw new Error('Gate weights are on GPU, use computeRouterLogitsGPU instead');
     }
 
@@ -140,13 +163,18 @@ export class MoERouter {
     if (!gateWeightBuffer) {
       throw new Error('Router gate weights not loaded');
     }
-    if (!isWeightBuffer(gateWeightBuffer) && !(gateWeightBuffer instanceof GPUBuffer)) {
+    if (!isWeightBuffer(gateWeightBuffer) && !isGpuBufferInstance(gateWeightBuffer)) {
       const uploaded = device.createBuffer({
         label: 'moe_gate_weight',
         size: gateWeightBuffer.byteLength,
         usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
       });
-      device.queue.writeBuffer(uploaded, 0,  (gateWeightBuffer));
+      try {
+        device.queue.writeBuffer(uploaded, 0, gateWeightBuffer);
+      } catch (error) {
+        uploaded.destroy();
+        throw error;
+      }
       this._gateWeightGPU = uploaded;
       this.gateWeight = uploaded;
       gateWeightBuffer = uploaded;
@@ -186,7 +214,7 @@ export class MoERouter {
 
   
   async _getGateBiasBuffer(device) {
-    if (this.gateBias instanceof GPUBuffer) return this.gateBias;
+    if (isGpuBufferInstance(this.gateBias)) return this.gateBias;
     if (this._gateBiasGPU) return this._gateBiasGPU;
 
     if (!(this.gateBias instanceof Float32Array)) {
@@ -198,7 +226,12 @@ export class MoERouter {
       size: this.gateBias.byteLength,
       usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
     });
-    device.queue.writeBuffer(buf, 0,  (this.gateBias));
+    try {
+      device.queue.writeBuffer(buf, 0, this.gateBias);
+    } catch (error) {
+      buf.destroy();
+      throw error;
+    }
     this._gateBiasGPU = buf;
     return buf;
   }
@@ -206,7 +239,7 @@ export class MoERouter {
   
   _inferBiasDtype(bias) {
     if (bias instanceof Float32Array) return 'f32';
-    if (bias instanceof GPUBuffer) {
+    if (isGpuBufferInstance(bias)) {
       const bytesPerElement = Math.round(bias.size / this.numExperts);
       return selectRuleValue('inference', 'dtype', 'f16OrF32FromBytes', { bytesPerElement });
     }
@@ -276,65 +309,64 @@ export class MoERouter {
       size: 16,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
     });
-    device.queue.writeBuffer(uniformBuffer, 0, uniformData);
-
-    const bindGroup = device.createBindGroup({
-      layout: pipeline.getBindGroupLayout(0),
-      entries: [
-        { binding: 0, resource: { buffer: uniformBuffer } },
-        { binding: 1, resource: { buffer: logits } },
-        { binding: 2, resource: { buffer: bias } },
-      ],
-    });
-
-    const encoder = device.createCommandEncoder({ label: 'moe_router_bias_add_encoder' });
-    const pass = encoder.beginComputePass({ label: 'moe_router_bias_add_pass' });
-    pass.setPipeline(pipeline);
-    pass.setBindGroup(0, bindGroup);
-    const total = numTokens * this.numExperts;
-    pass.dispatchWorkgroups(Math.ceil(total / 256));
-    pass.end();
-    device.queue.submit([encoder.finish()]);
+    try {
+      device.queue.writeBuffer(uniformBuffer, 0, uniformData);
+
+      const bindGroup = device.createBindGroup({
+        layout: pipeline.getBindGroupLayout(0),
+        entries: [
+          { binding: 0, resource: { buffer: uniformBuffer } },
+          { binding: 1, resource: { buffer: logits } },
+          { binding: 2, resource: { buffer: bias } },
+        ],
+      });
 
-    uniformBuffer.destroy();
+      const encoder = device.createCommandEncoder({ label: 'moe_router_bias_add_encoder' });
+      const pass = encoder.beginComputePass({ label: 'moe_router_bias_add_pass' });
+      pass.setPipeline(pipeline);
+      pass.setBindGroup(0, bindGroup);
+      const total = numTokens * this.numExperts;
+      pass.dispatchWorkgroups(Math.ceil(total / 256));
+      pass.end();
+      device.queue.submit([encoder.finish()]);
+    } finally {
+      uniformBuffer.destroy();
+    }
   }
 
   
   async routeGPU(hiddenStates, numTokens) {
     // Compute router logits on GPU
     const logitsBuffer = await this.computeRouterLogitsGPU(hiddenStates, numTokens);
-
-    // Read back logits to CPU for top-k selection
-    // (GPU top-k is complex and not always faster for small numExperts)
-    const logitsData = await readBuffer(logitsBuffer);
-    const logits = this.lastLogitsDtype === 'f16'
-      ? f16ToF32Array(new Uint16Array(logitsData))
-      : new Float32Array(logitsData);
-
-    
-    const selections = [];
-    this.activeExperts.clear();
-
-    for (let t = 0; t < numTokens; t++) {
-      const tokenLogits = logits.subarray(
-        t * this.numExperts,
-        (t + 1) * this.numExperts
-      );
-
-      const selection = this.selectExpertsForToken(tokenLogits);
-      selections.push(selection);
-
-      for (const idx of selection.indices) {
-        this.activeExperts.add(idx);
-        this.loadBalanceStats.expertCounts[idx]++;
+    try {
+      const logitsData = await readBuffer(logitsBuffer);
+      const logits = this.lastLogitsDtype === 'f16'
+        ? f16ToF32Array(new Uint16Array(logitsData))
+        : new Float32Array(logitsData);
+
+      const selections = [];
+      this.activeExperts.clear();
+
+      for (let t = 0; t < numTokens; t++) {
+        const tokenLogits = logits.subarray(
+          t * this.numExperts,
+          (t + 1) * this.numExperts
+        );
+
+        const selection = this.selectExpertsForToken(tokenLogits);
+        selections.push(selection);
+
+        for (const idx of selection.indices) {
+          this.activeExperts.add(idx);
+          this.loadBalanceStats.expertCounts[idx]++;
+        }
+        this.loadBalanceStats.totalTokens++;
       }
-      this.loadBalanceStats.totalTokens++;
-    }
-
-    // Clean up logits buffer
-    releaseBuffer(logitsBuffer);
 
-    return selections;
+      return selections;
+    } finally {
+      releaseBuffer(logitsBuffer);
+    }
   }
 
   

package/src/inference/multi-model-network.js

@@ -4,9 +4,10 @@ import { ExpertRouter } from './expert-router.js';
 import { MultiModelRecorder } from '../gpu/multi-model-recorder.js';
 import { applyRepetitionPenalty, sample, getTopK } from './pipelines/text/sampling.js';
 import { finalizeLogits, extractLastPositionLogits } from './pipelines/text/logits/index.js';
+import { readBufferWithCleanup } from './pipelines/text/logits/utils.js';
 import { isStopToken } from './pipelines/text/init.js';
 import { mergeMultipleLogits } from '../gpu/kernels/logit-merge.js';
-import { releaseBuffer, readBuffer } from '../memory/buffer-pool.js';
+import { releaseBuffer } from '../memory/buffer-pool.js';
 
 const MIN_AGREEMENT_WEIGHT = 1e-4;
 
@@ -478,8 +479,9 @@ export class MultiModelNetwork {
       if (canMergeOnGpu) {
         const buffers = voterResults.map((result) => result.logitsBuffer);
         const mergedBuffer = await mergeMultipleLogits(buffers, rawVocabSize, normalizedWeights, 1.0);
-        const mergedData = await readBuffer(mergedBuffer, rawVocabSize * 4);
-        releaseBuffer(mergedBuffer);
+        const mergedData = await readBufferWithCleanup(mergedBuffer, rawVocabSize * 4, () => {
+          releaseBuffer(mergedBuffer);
+        });
         const rawMerged = new Float32Array(mergedData);
         const finalized = await finalizeLogits(
           rawMerged,

package/src/inference/network-evolution.d.ts

@@ -35,12 +35,21 @@ export interface EvolutionConfig {
   generations?: number;
   eliteCount?: number;
   mutationRate?: number;
+  random: () => number;
   evaluate: (genome: NetworkGenome) => Promise<number>;
   randomGenome: () => NetworkGenome;
 }
 
-export declare const mutateGenome: (genome: NetworkGenome, mutationRate?: number) => NetworkGenome;
+export declare const mutateGenome: (
+  genome: NetworkGenome,
+  mutationRate?: number,
+  random?: (() => number) | null
+) => NetworkGenome;
 
-export declare const crossoverGenome: (a: NetworkGenome, b: NetworkGenome) => NetworkGenome;
+export declare const crossoverGenome: (
+  a: NetworkGenome,
+  b: NetworkGenome,
+  random?: (() => number) | null
+) => NetworkGenome;
 
 export declare function evolveNetwork(config: EvolutionConfig): Promise<NetworkGenome>;