@simulatte/doppler 0.1.6 → 0.1.7

package/src/loader/doppler-loader.js

@@ -9,7 +9,7 @@ import {
   verifyIntegrity,
   loadManifestFromStore,
 } from '../storage/shard-manager.js';
-import { parseManifest } from '../formats/rdrr/index.js';
+import { clearManifest, parseManifest, setManifest as setCurrentManifest } from '../formats/rdrr/index.js';
 import { initDevice, getDevice, getKernelCapabilities } from '../gpu/device.js';
 import { acquireBuffer, releaseBuffer, forceBufferPoolReclaim } from '../memory/buffer-pool.js';
 import { getExpertCache } from './experts/expert-cache.js';
@@ -50,6 +50,10 @@ function hasExpertGroups(manifest) {
   return Object.keys(manifest.groups).some((groupId) => groupId.includes('.expert.'));
 }
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
+
 // Re-export types for backward compatibility
 export {
   // Types are in .d.ts file
@@ -252,6 +256,7 @@ export class DopplerLoader {
   
   setManifest(manifest) {
     this.manifest = manifest;
+    setCurrentManifest(manifest);
     const moeConfig = manifest.moeConfig;
     this.isMoE = moeConfig != null && (moeConfig.numExperts ?? 0) > 1;
     if (!this.isMoE && hasExpertGroups(manifest)) {
@@ -259,6 +264,7 @@ export class DopplerLoader {
         `Manifest "${manifest.modelId ?? 'unknown'}" missing moeConfig for MoE model. Re-convert with moeConfig.`
       );
     }
+    this.shardCache.setManifest(this.manifest);
     this.shardCache.configureForModel(this.manifest, this.shardCache.hasCustomLoader);
     debugTrace.loader('Manifest set externally');
   }
@@ -679,7 +685,7 @@ export class DopplerLoader {
       const device = getDevice();
       if (!device) {
         log.warn('Loader', 'GPU device not available; falling back to CPU');
-        if (shardData instanceof GPUBuffer) {
+        if (isGpuBufferInstance(shardData)) {
           releaseBuffer(shardData);
           shardData = await this.#assembleShardData(location, name);
         }
@@ -708,7 +714,7 @@ export class DopplerLoader {
       return result.data;
     }
 
-    if (shardData instanceof GPUBuffer) {
+    if (isGpuBufferInstance(shardData)) {
       // Shouldn't happen (streaming is only used for toGPU), but keep this leak-proof.
       releaseBuffer(shardData);
       shardData = await this.#assembleShardData(location, name);
@@ -751,31 +757,40 @@ export class DopplerLoader {
     // queue.writeBuffer requires 4-byte aligned sizes; we pad the buffer.
     const alignedSize = Math.ceil(location.size / 4) * 4;
     const raw = acquireBuffer(alignedSize, undefined, `raw_${name}`);
+    let complete = false;
 
-    let dstOffset = 0;
-    const uploadChunk = (bytes) => {
-      device.queue.writeBuffer(raw, dstOffset, bytes, bytes.byteOffset, bytes.byteLength);
-      dstOffset += bytes.byteLength;
-    };
-    const streamRange = (idx, offset, length) => this.shardCache.streamRange(idx, offset, length, { chunkBytes });
+    try {
+      let dstOffset = 0;
+      const uploadChunk = (bytes) => {
+        device.queue.writeBuffer(raw, dstOffset, bytes, bytes.byteOffset, bytes.byteLength);
+        dstOffset += bytes.byteLength;
+      };
+      const streamRange = (idx, offset, length) => this.shardCache.streamRange(idx, offset, length, { chunkBytes });
 
-    if (location.spans) {
-      for (const span of location.spans) {
-        for await (const chunk of streamRange(span.shardIndex, span.offset, span.size)) {
+      if (location.spans) {
+        for (const span of location.spans) {
+          for await (const chunk of streamRange(span.shardIndex, span.offset, span.size)) {
+            uploadChunk(chunk);
+          }
+        }
+      } else {
+        for await (const chunk of streamRange(location.shardIndex, location.offset, location.size)) {
           uploadChunk(chunk);
         }
       }
-    } else {
-      for await (const chunk of streamRange(location.shardIndex, location.offset, location.size)) {
-        uploadChunk(chunk);
-      }
-    }
 
-    if (dstOffset < location.size) {
-      log.warn('Loader', `Stream upload short read for "${name}": got=${dstOffset}, expected=${location.size}`);
+      if (dstOffset !== location.size) {
+        throw new Error(
+          `Stream upload short read for "${name}": got=${dstOffset}, expected=${location.size}.`
+        );
+      }
+      complete = true;
+      return raw;
+    } finally {
+      if (!complete) {
+        releaseBuffer(raw);
+      }
     }
-
-    return raw;
   }
 
   
@@ -950,7 +965,7 @@ export class DopplerLoader {
       if (!value) return;
       const gpuBuffer = isWeightBuffer(value)
         ? value.buffer
-        : (value instanceof GPUBuffer ? value : null);
+        : (isGpuBufferInstance(value) ? value : null);
       if (!gpuBuffer) return;
       try {
         releaseBuffer(gpuBuffer);
@@ -990,6 +1005,7 @@ export class DopplerLoader {
     this.lmHead = null;
     this.finalNorm = null;
     this.manifest = null;
+    clearManifest();
     this.modelId = null;
     this.loadedShards.clear();
     this.isLoaded = false;

package/src/loader/dtype-utils.js

@@ -1,7 +1,4 @@
 
-
-import { getDevice } from '../gpu/device.js';
-import { isTraceEnabled, log, trace as debugTrace } from '../debug/index.js';
 import { selectRuleValue } from '../rules/rule-registry.js';
 import { tagBufferDtype } from '../gpu/weight-buffer.js';
 
@@ -26,46 +23,8 @@ export function f16ToF32(h) {
 
 
 export async function convertBF16ToF32GPU(srcBuffer, numElements, name) {
-  debugTrace.loader(`[BF16->F32] Importing cast.js...`);
-  const castModule = await import('../gpu/kernels/cast.js');
-  debugTrace.loader(`[BF16->F32] castModule keys:`, Object.keys(castModule));
-  const { runBF16ToF32 } = castModule;
-  debugTrace.loader(`[BF16->F32] runBF16ToF32 type: ${typeof runBF16ToF32}`);
+  const { runBF16ToF32 } = await import('../gpu/kernels/cast.js');
   const resultTensor = await runBF16ToF32(srcBuffer, [numElements], name);
-  debugTrace.loader(`[BF16->F32] runBF16ToF32 returned, result.size=${resultTensor.buffer?.size}`);
-
-  // Debug: Verify conversion produced non-zero values
-  const shouldCheckEmbed = isTraceEnabled('loader') &&
-    name.includes('embed') &&
-    name.includes('embed_tokens');
-  if (shouldCheckEmbed) {
-    try {
-      debugTrace.loader(`[BF16->F32] Checking embed buffer for non-zeros...`);
-      const device = getDevice();
-      const sampleSize = Math.min(1024, resultTensor.buffer.size);
-      debugTrace.loader(`[BF16->F32] Creating staging buffer size=${sampleSize}`);
-      const stagingBuffer = device.createBuffer({
-        size: sampleSize,
-        usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
-      });
-      debugTrace.loader(`[BF16->F32] Copying to staging buffer...`);
-      const encoder = device.createCommandEncoder();
-      encoder.copyBufferToBuffer(resultTensor.buffer, 0, stagingBuffer, 0, sampleSize);
-      device.queue.submit([encoder.finish()]);
-      debugTrace.loader(`[BF16->F32] Mapping staging buffer...`);
-      await stagingBuffer.mapAsync(GPUMapMode.READ);
-      debugTrace.loader(`[BF16->F32] Reading data...`);
-      const data = new Float32Array(stagingBuffer.getMappedRange().slice(0));
-      stagingBuffer.unmap();
-      stagingBuffer.destroy();
-      const nonZero = Array.from(data).filter(x => x !== 0);
-      const nanCount = data.filter(x => !Number.isFinite(x)).length;
-      debugTrace.loader(`[BF16->F32] nonZero=${nonZero.length}/${data.length}, nan=${nanCount}, sample=[${nonZero.slice(0, 5).map(x => x.toFixed(4)).join(', ')}]`);
-    } catch (err) {
-      log.error('Loader', 'BF16->F32 embed buffer check error:',  (err).message);
-    }
-  }
-
   return resultTensor.buffer;
 }
 
@@ -84,11 +43,11 @@ function normalizeBufferDtype(locationDtype, outputDtype) {
   if (explicit) {
     return explicit;
   }
-  const location = typeof locationDtype === 'string' ? locationDtype.toLowerCase() : null;
+  const location = typeof locationDtype === 'string' ? locationDtype.toUpperCase() : null;
   if (!location) {
     return null;
   }
-  return selectRuleValue('loader', 'weights', 'floatLocationDtype', { locationDtype: locationDtype });
+  return selectRuleValue('loader', 'weights', 'floatLocationDtype', { locationDtype: location });
 }
 
 export function applyBufferLayout(buffer, location, outputDtype = null) {

package/src/loader/embedding-loader.js

@@ -23,6 +23,10 @@ import { releaseBuffer } from '../memory/buffer-pool.js';
 const EMBEDDING_ROLE = 'embedding';
 const EMBEDDING_GROUP = 'embed';
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
+
 // ============================================================================
 // Main Function
 // ============================================================================
@@ -59,7 +63,7 @@ export async function loadEmbeddings(ctx) {
     }
 
     // Handle valid tensor types
-    if (tensor instanceof GPUBuffer || isWeightBuffer(tensor) || tensor instanceof Float32Array) {
+    if (isGpuBufferInstance(tensor) || isWeightBuffer(tensor) || tensor instanceof Float32Array) {
       const result = await processEmbeddingTensor(ctx, tensor, name, loc, shouldStream);
       if (result) {
         return result;
@@ -107,7 +111,7 @@ async function processEmbeddingTensor(ctx, tensor, name, loc, shouldStream) {
   }
 
   // Raw GPUBuffer - wrap with dtype/layout metadata
-  if (promoted instanceof GPUBuffer && loc?.shape && loc.shape.length === 2) {
+  if (isGpuBufferInstance(promoted) && loc?.shape && loc.shape.length === 2) {
     const layout = ctx.resolveWeightLayout(loc);
     
     const dtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
@@ -140,7 +144,7 @@ async function maybePromoteEmbeddingsToF32(ctx, current, name, loc) {
     return wrapped;
   }
 
-  if (!(current instanceof GPUBuffer)) return current;
+  if (!isGpuBufferInstance(current)) return current;
 
   const sourceDtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
     locationDtype: loc?.dtype,

package/src/loader/experts/expert-cache.js

@@ -3,6 +3,11 @@
 import { releaseBuffer } from '../../memory/buffer-pool.js';
 import { log, trace } from '../../debug/index.js';
 import { getRuntimeConfig } from '../../config/runtime.js';
+import { isWeightBuffer } from '../../gpu/weight-buffer.js';
+
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
 
 
 
@@ -256,12 +261,14 @@ export class ExpertCache {
     ];
 
     for (const buf of buffers) {
-      if (buf instanceof GPUBuffer) {
-        try {
-          releaseBuffer(buf);
-        } catch (e) {
-          // Buffer may already be released
-        }
+      const gpuBuffer = isWeightBuffer(buf)
+        ? buf.buffer
+        : (isGpuBufferInstance(buf) ? buf : null);
+      if (!gpuBuffer) continue;
+      try {
+        releaseBuffer(gpuBuffer);
+      } catch (e) {
+        // Buffer may already be released
       }
     }
   }

package/src/loader/experts/expert-loader.js

@@ -18,7 +18,7 @@ import { releaseBuffer } from '../../memory/buffer-pool.js';
 
 export async function preloadShardsForExpert(ctx, layerIdx, expertIdx, options) {
   // Get required shards from manifest mapping
-  const shardIndices = getShardsForExpert(layerIdx, expertIdx);
+  const shardIndices = getShardsForExpert(layerIdx, expertIdx, ctx.manifest);
   if (shardIndices.length === 0) {
     // No mapping available, fall back to loading all shards on demand
     return;
@@ -69,6 +69,10 @@ export function predictNextLayerExperts(currentExperts) {
   return currentExperts;
 }
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
+
 // ============================================================================
 // Expert Loading
 // ============================================================================
@@ -95,7 +99,7 @@ export async function loadExpert(ctx, layerIdx, expertIdx) {
   await preloadShardsForExpert(ctx, layerIdx, expertIdx);
 
   // Get tensor names from manifest if available (for logging/debugging)
-  const tensorNames = getTensorsForExpert(layerIdx, expertIdx);
+  const tensorNames = getTensorsForExpert(layerIdx, expertIdx, ctx.manifest);
   if (tensorNames.length > 0) {
     debugTrace.loader(`Expert ${layerIdx}_${expertIdx} tensors: ${tensorNames.length}`);
   }
@@ -260,7 +264,7 @@ function getGpuBuffer(value) {
   if (isWeightBuffer(value)) {
     return value.buffer;
   }
-  if (value instanceof GPUBuffer) {
+  if (isGpuBufferInstance(value)) {
     return value;
   }
   return null;
@@ -342,7 +346,7 @@ async function downcastExpertWeights(ctx, weights) {
     if (!buf) continue;
 
     // Only downcast GPUBuffer or WeightBuffer (not Float32Array)
-    if (!(buf instanceof GPUBuffer) && !isWeightBuffer(buf)) {
+    if (!isGpuBufferInstance(buf) && !isWeightBuffer(buf)) {
       continue;
     }
 
@@ -369,13 +373,13 @@ function calculateExpertSize(weights) {
     const buf = weights[k];
     if (isWeightBuffer(buf)) {
       sizeBytes += buf.buffer.size;
-    } else if (buf instanceof GPUBuffer) {
+    } else if (isGpuBufferInstance(buf)) {
       sizeBytes += buf.size;
     }
   }
 
   // Use manifest-provided expert size if available, otherwise use calculated
-  const manifestBytes = getExpertBytes();
+  const manifestBytes = getExpertBytes(ctx.manifest);
   if (manifestBytes > 0) {
     sizeBytes = manifestBytes;
   }

package/src/loader/final-weights-loader.js

@@ -20,6 +20,10 @@ const HEAD_GROUP = 'head';
 const FINAL_NORM_ROLE = 'norm';
 const LM_HEAD_ROLE = 'lm_head';
 
+function isGpuBufferInstance(value) {
+  return typeof GPUBuffer !== 'undefined' && value instanceof GPUBuffer;
+}
+
 function isLikelyFinalNormName(name) {
   const lower = String(name || '').toLowerCase();
   if (!lower) return false;
@@ -148,7 +152,7 @@ async function loadLmHead(ctx) {
       );
     }
 
-    if (tensor && (tensor instanceof GPUBuffer || isWeightBuffer(tensor) || tensor instanceof Float32Array)) {
+    if (tensor && (isGpuBufferInstance(tensor) || isWeightBuffer(tensor) || tensor instanceof Float32Array)) {
       lmHeadName = name;
       lmHeadLoc = loc;
       lmHead = processLmHeadTensor(ctx, tensor, name, loc, shouldStream);
@@ -189,7 +193,7 @@ function processLmHeadTensor(ctx, tensor, name, loc, shouldStream) {
   }
 
   // Raw GPUBuffer - wrap with dtype/layout metadata
-  if (tensor instanceof GPUBuffer && loc.shape && loc.shape.length === 2) {
+  if (isGpuBufferInstance(tensor) && loc.shape && loc.shape.length === 2) {
     const layout = ctx.resolveWeightLayout(loc);
     
     const dtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
@@ -209,7 +213,7 @@ async function maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc) {
   const tiedToEmbeddings =
     lmHead === ctx.embeddings ||
     (isWeightBuffer(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead.buffer === ctx.embeddings.buffer) ||
-    (lmHead instanceof GPUBuffer && isWeightBuffer(ctx.embeddings) && lmHead === ctx.embeddings.buffer);
+    (isGpuBufferInstance(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead === ctx.embeddings.buffer);
 
   if (tiedToEmbeddings) {
     return lmHead;
@@ -234,7 +238,7 @@ async function maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc) {
 
   // Get buffer for downcast
   const buffer = isWeightBuffer(lmHead) ? lmHead.buffer : lmHead;
-  if (!(buffer instanceof GPUBuffer)) {
+  if (!isGpuBufferInstance(buffer)) {
     return lmHead;
   }
 

package/src/loader/layer-loader.js

@@ -224,7 +224,8 @@ function createTryLoad(ctx, prefixes) {
     for (const prefix of prefixes) {
       for (const suffix of suffixes) {
         const tensor = await ctx.loadTensor(`${prefix}.${suffix}`, true, true);
-        if (tensor && (tensor instanceof GPUBuffer || tensor instanceof Float32Array || isWeightBuffer(tensor))) {
+        const isGpuBuffer = typeof GPUBuffer !== 'undefined' && tensor instanceof GPUBuffer;
+        if (tensor && (isGpuBuffer || tensor instanceof Float32Array || isWeightBuffer(tensor))) {
           return tensor;
         }
       }

package/src/loader/loader-state.js

@@ -122,14 +122,14 @@ export class LoaderState {
 
   static getGPUBuffer(weight) {
     if (!weight) return null;
-    if (weight instanceof GPUBuffer) return weight;
+    if (typeof GPUBuffer !== 'undefined' && weight instanceof GPUBuffer) return weight;
     if (isWeightBuffer(weight)) return weight.buffer;
     return null;
   }
 
   static isGPUBacked(weight) {
     if (!weight) return false;
-    if (weight instanceof GPUBuffer) return true;
+    if (typeof GPUBuffer !== 'undefined' && weight instanceof GPUBuffer) return true;
     if (isWeightBuffer(weight)) return true;
     if (isCpuWeightBuffer(weight)) return false;
     if (weight instanceof Float32Array) return false;

package/src/loader/memory-monitor.js

@@ -105,6 +105,10 @@ export class MemoryMonitor {
 
 
   start(getState) {
+    if (this.#interval) {
+      clearInterval(this.#interval);
+      this.#interval = null;
+    }
     this.#startTime = performance.now();
     this.#snapshots = [];
     this.#log('start', getState());
@@ -209,6 +213,10 @@ export class MemoryTimeSeries {
 
 
   start() {
+    if (this.#interval) {
+      clearInterval(this.#interval);
+      this.#interval = null;
+    }
     this.#startTime = performance.now();
     this.#samples = [];
     this.#capture('start');

package/src/loader/multi-model-loader.d.ts

@@ -22,6 +22,20 @@ export declare class MultiModelLoader {
   baseWeights: WeightLoadResult | null;
   adapters: Map<string, LoRAAdapter>;
 
+  _loadBaseWeights(
+    manifest: Manifest,
+    options: { storageContext?: { loadShard?: (index: number) => Promise<ArrayBuffer | Uint8Array> } },
+    runtimeConfig: unknown
+  ): Promise<WeightLoadResult>;
+
+  _resolveAdapterSource(source: AdapterSource): Promise<LoRAAdapter>;
+
+  _createPipeline(): InferencePipeline;
+
+  _getBaseLoader(): { unload(): Promise<void> };
+
+  unload(): Promise<void>;
+
   loadBase(
     manifest: Manifest,
     options?: { storageContext?: { loadShard?: (index: number) => Promise<ArrayBuffer | Uint8Array> } }

package/src/loader/multi-model-loader.js

@@ -17,37 +17,68 @@ export class MultiModelLoader {
   
   adapters = new Map();
 
-  
-  async loadBase(manifest, options = {}) {
-    // Get runtime model overrides to merge with manifest inference config
-    const runtimeConfig = getRuntimeConfig();
+  #pipelines = new Set();
+
+  async _loadBaseWeights(manifest, options, runtimeConfig) {
     const modelOverrides =  (runtimeConfig.inference.modelOverrides);
     const config = parseModelConfig(manifest, modelOverrides);
-    this.baseManifest = manifest;
-    this.baseWeights = await loadWeights(manifest, config, {
+    return loadWeights(manifest, config, {
       storageContext: options.storageContext,
       keepF32Weights: runtimeConfig.inference.compute.keepF32Weights === true,
     });
-    return this.baseWeights;
   }
 
-  
-  async loadAdapter(name, source) {
-    
-    let adapter;
-
+  async _resolveAdapterSource(source) {
     if (typeof source === 'string') {
-      adapter = await loadLoRAFromUrl(source);
-    } else if (this.#isRDRRManifest(source)) {
+      return loadLoRAFromUrl(source);
+    }
+    if (this.#isRDRRManifest(source)) {
       const loader = getDopplerLoader();
       await loader.init();
-      adapter = await loader.loadLoRAWeights(source);
-    } else if (this.#isLoRAManifest(source)) {
-      adapter = await loadLoRAFromManifest(source);
-    } else {
-      adapter = source;
+      return loader.loadLoRAWeights(source);
+    }
+    if (this.#isLoRAManifest(source)) {
+      return loadLoRAFromManifest(source);
+    }
+    return source;
+  }
+
+  _createPipeline() {
+    return new InferencePipeline();
+  }
+
+  _getBaseLoader() {
+    return getDopplerLoader();
+  }
+
+  async unload() {
+    const pipelines = Array.from(this.#pipelines);
+    this.#pipelines.clear();
+    await Promise.all(pipelines.map(async (pipeline) => pipeline.unload()));
+
+    if (this.baseWeights) {
+      const loader = this._getBaseLoader();
+      await loader.unload();
     }
 
+    this.baseManifest = null;
+    this.baseWeights = null;
+    this.adapters.clear();
+  }
+
+  async loadBase(manifest, options = {}) {
+    await this.unload();
+
+    const runtimeConfig = getRuntimeConfig();
+    const weights = await this._loadBaseWeights(manifest, options, runtimeConfig);
+    this.baseManifest = manifest;
+    this.baseWeights = weights;
+    return weights;
+  }
+
+  async loadAdapter(name, source) {
+    const adapter = await this._resolveAdapterSource(source);
+
     const adapterName = name || adapter.name;
     this.adapters.set(adapterName, adapter);
     return adapter;
@@ -68,11 +99,26 @@ export class MultiModelLoader {
     if (!this.baseManifest || !this.baseWeights) {
       throw new Error('Base model not loaded');
     }
-    const pipeline = new InferencePipeline();
-    await pipeline.initialize(contexts);
-    pipeline.setPreloadedWeights(this.baseWeights);
-    await pipeline.loadModel(this.baseManifest);
-    return pipeline;
+    const pipeline = this._createPipeline();
+    const unloadPipeline = pipeline.unload.bind(pipeline);
+    pipeline.unload = async () => {
+      try {
+        await unloadPipeline();
+      } finally {
+        this.#pipelines.delete(pipeline);
+      }
+    };
+
+    try {
+      await pipeline.initialize(contexts);
+      pipeline.setPreloadedWeights(this.baseWeights);
+      await pipeline.loadModel(this.baseManifest);
+      this.#pipelines.add(pipeline);
+      return pipeline;
+    } catch (error) {
+      await pipeline.unload().catch(() => {});
+      throw error;
+    }
   }