@simulatte/doppler 0.1.5 → 0.1.7

package/src/inference/tokenizers/bundled.js

@@ -64,6 +64,68 @@ function resolveSpecialTokens(specialTokensRaw, fallbackTokens, vocab) {
   return resolved;
 }
 
+function resolveByteLevelPretokenizerConfig(preTokenizer) {
+  if (!preTokenizer || typeof preTokenizer !== 'object') {
+    return {
+      useByteLevel: false,
+      addPrefixSpace: null,
+    };
+  }
+
+  if (preTokenizer.type === 'ByteLevel') {
+    return {
+      useByteLevel: true,
+      addPrefixSpace: preTokenizer.add_prefix_space === true,
+    };
+  }
+
+  if (preTokenizer.type === 'Sequence' && Array.isArray(preTokenizer.pretokenizers)) {
+    for (const entry of preTokenizer.pretokenizers) {
+      const resolved = resolveByteLevelPretokenizerConfig(entry);
+      if (resolved.useByteLevel) {
+        return resolved;
+      }
+    }
+  }
+
+  return {
+    useByteLevel: false,
+    addPrefixSpace: null,
+  };
+}
+
+function registerAddedTokens(addedTokens, vocab, reverseVocab, patterns, specialTokenIds, derivedSpecialTokens = null) {
+  let maxId = -1;
+  for (const token of addedTokens) {
+    const content = token?.content;
+    const id = typeof token?.id === 'number' ? token.id : parseInt(token?.id, 10);
+    if (!Number.isFinite(id) || !content) continue;
+    if (!vocab.has(content)) {
+      vocab.set(content, id);
+      reverseVocab.set(id, content);
+    }
+    if (id > maxId) maxId = id;
+    if (content.length > 1) {
+      patterns.push({ content, id });
+    }
+    if (token.special) {
+      specialTokenIds.add(id);
+      if (derivedSpecialTokens) {
+        if (derivedSpecialTokens.bos == null && (content === '<bos>' || content === '<s>' || content.includes('bos'))) {
+          derivedSpecialTokens.bos = id;
+        } else if (derivedSpecialTokens.eos == null && (content === '<eos>' || content === '</s>' || content.includes('eos'))) {
+          derivedSpecialTokens.eos = id;
+        } else if (derivedSpecialTokens.pad == null && (content === '<pad>' || content.includes('pad'))) {
+          derivedSpecialTokens.pad = id;
+        } else if (derivedSpecialTokens.unk == null && (content === '<unk>' || content.includes('unk'))) {
+          derivedSpecialTokens.unk = id;
+        }
+      }
+    }
+  }
+  return maxId;
+}
+
 
 export class TransformersTokenizer extends BaseTokenizer {
 
@@ -156,6 +218,10 @@ export class BundledTokenizer extends BaseTokenizer {
   
   #byteDecoder = null;
 
+  #byteEncoder = null;
+
+  #useByteLevelEncoding = false;
+
   
   constructor(config = {}) {
     // BundledTokenizer gets vocabSize from load(), so defer validation
@@ -164,6 +230,25 @@ export class BundledTokenizer extends BaseTokenizer {
     });
   }
 
+  #resetState() {
+    this.#vocab.clear();
+    this.#reverseVocab.clear();
+    this.#merges = [];
+    this.#mergeRanks.clear();
+    this.#scores = [];
+    this.#tokenTypes = [];
+    this.#type = null;
+    this.#byteTokens.clear();
+    this.#specialTokenPatterns = [];
+    this.#specialTokenIds = new Set();
+    this.#addSpacePrefix = true;
+    this.#spacePrefixChar = '▁';
+    this.#byteDecoder = null;
+    this.#byteEncoder = null;
+    this.#useByteLevelEncoding = false;
+    this.vocabSize = 0;
+  }
+
   
   isSpecialToken(tokenId) {
     if (this.#specialTokenIds.size > 0) {
@@ -199,13 +284,25 @@ export class BundledTokenizer extends BaseTokenizer {
     }
 
     this.#byteDecoder = new Map();
+    this.#byteEncoder = new Map();
     for (let i = 0; i < base.length; i++) {
       this.#byteDecoder.set(String.fromCodePoint(chars[i]), base[i]);
+      this.#byteEncoder.set(base[i], String.fromCodePoint(chars[i]));
     }
   }
 
+  #encodeByteLevelText(text) {
+    const bytes = new TextEncoder().encode(text);
+    let out = '';
+    for (const byte of bytes) {
+      out += this.#byteEncoder?.get(byte) ?? String.fromCharCode(byte);
+    }
+    return out;
+  }
+
   
   load(tokenizerJson) {
+    this.#resetState();
     // Detect format: HuggingFace has model.vocab, bundled has top-level vocab
     const isHuggingFace = 'model' in tokenizerJson && tokenizerJson.model?.vocab !== undefined;
 
@@ -290,30 +387,16 @@ export class BundledTokenizer extends BaseTokenizer {
       eos: null,
       unk: null,
     };
-    for (const token of addedTokens) {
-      const content = token.content;
-      const id = typeof token.id === 'number' ? token.id : parseInt( (token.id), 10);
-      if (!Number.isFinite(id) || !content) continue;
-      if (!this.#vocab.has(content)) {
-        this.#vocab.set(content, id);
-        this.#reverseVocab.set(id, content);
-      }
-      if (id > maxId) maxId = id;
-      if (token.special) {
-        specialTokenIds.add(id);
-        if (content.length > 1) {
-          specialTokenPatterns.push({ content, id });
-        }
-        if (derivedSpecialTokens.bos == null && (content === '<bos>' || content === '<s>' || content.includes('bos'))) {
-          derivedSpecialTokens.bos = id;
-        } else if (derivedSpecialTokens.eos == null && (content === '<eos>' || content === '</s>' || content.includes('eos'))) {
-          derivedSpecialTokens.eos = id;
-        } else if (derivedSpecialTokens.pad == null && (content === '<pad>' || content.includes('pad'))) {
-          derivedSpecialTokens.pad = id;
-        } else if (derivedSpecialTokens.unk == null && (content === '<unk>' || content.includes('unk'))) {
-          derivedSpecialTokens.unk = id;
-        }
-      }
+    const addedMaxId = registerAddedTokens(
+      addedTokens,
+      this.#vocab,
+      this.#reverseVocab,
+      specialTokenPatterns,
+      specialTokenIds,
+      derivedSpecialTokens
+    );
+    if (addedMaxId > maxId) {
+      maxId = addedMaxId;
     }
 
     const specialTokensRaw = hf.special_tokens_map || hf.specialTokens || hf.special_tokens || null;
@@ -351,6 +434,7 @@ export class BundledTokenizer extends BaseTokenizer {
 
     // Handle behavior flags (use HF config if present, else runtime defaults)
     const runtimeDefaults = getRuntimeConfig().inference.tokenizer;
+    const byteLevelPretokenizer = resolveByteLevelPretokenizerConfig(hf.pre_tokenizer);
     const configuredAddBosToken = this.addBosToken;
     const configuredAddEosToken = this.addEosToken;
     this.addBosToken =
@@ -378,9 +462,16 @@ export class BundledTokenizer extends BaseTokenizer {
     // - runtime config addSpacePrefix (user override or null for auto-detect)
     const decoderPrepend = hf.decoder?.prepend_scheme === 'always' || hf.decoder?.add_prefix_space === true;
     const normalizerPrepend = hf.normalizer?.prepend_scheme === 'always' || hf.normalizer?.add_prefix_space === true;
+    this.#useByteLevelEncoding = byteLevelPretokenizer.useByteLevel;
     const runtimeSpacePrefix = runtimeDefaults.addSpacePrefix;
     // Use explicit runtime config if set (non-null), otherwise auto-detect from tokenizer.json
-    this.#addSpacePrefix = runtimeSpacePrefix ?? model.add_prefix_space ?? model.add_dummy_prefix ?? decoderPrepend ?? normalizerPrepend ?? false;
+    this.#addSpacePrefix = runtimeSpacePrefix
+      ?? byteLevelPretokenizer.addPrefixSpace
+      ?? model.add_prefix_space
+      ?? model.add_dummy_prefix
+      ?? decoderPrepend
+      ?? normalizerPrepend
+      ?? false;
     log.debug('Tokenizer', `addSpacePrefix=${this.#addSpacePrefix} (runtime=${runtimeSpacePrefix}, model=${model.add_prefix_space ?? model.add_dummy_prefix}, decoder=${decoderPrepend}, normalizer=${normalizerPrepend})`);
 
     // Detect space prefix style by checking which WORD tokens exist in vocab
@@ -469,11 +560,47 @@ export class BundledTokenizer extends BaseTokenizer {
       this.#tokenTypes = tokenizerJson.tokenTypes;
     }
 
+    let maxId = -1;
+    for (const id of this.#vocab.values()) {
+      if (Number.isFinite(id) && id > maxId) {
+        maxId = id;
+      }
+    }
+
+    const addedTokens = Array.isArray(tokenizerJson.added_tokens) ? tokenizerJson.added_tokens : [];
+    const tokenPatterns = [];
+    const specialTokenIds = new Set();
+    const derivedSpecialTokens = {
+      pad: null,
+      bos: null,
+      eos: null,
+      unk: null,
+    };
+    const addedMaxId = registerAddedTokens(
+      addedTokens,
+      this.#vocab,
+      this.#reverseVocab,
+      tokenPatterns,
+      specialTokenIds,
+      derivedSpecialTokens
+    );
+    if (addedMaxId > maxId) {
+      maxId = addedMaxId;
+    }
+
     // Set special tokens - support both camelCase and snake_case formats
     const specialTokensRaw =  (tokenizerJson.specialTokens ||  (tokenizerJson).special_tokens);
-    this.specialTokens = resolveSpecialTokens(specialTokensRaw, this.specialTokens, this.#vocab);
+    this.specialTokens = resolveSpecialTokens(
+      specialTokensRaw,
+      {
+        ...derivedSpecialTokens,
+        ...this.specialTokens,
+      },
+      this.#vocab
+    );
     log.debug('Tokenizer', `Special tokens: BOS=${this.specialTokens.bos}, EOS=${this.specialTokens.eos}`);
-    this.#specialTokenIds = new Set();
+    this.#specialTokenIds = specialTokenIds;
+    this.#specialTokenPatterns = tokenPatterns;
     const builtinSpecials = [
       this.specialTokens.pad,
       this.specialTokens.bos,
@@ -485,8 +612,13 @@ export class BundledTokenizer extends BaseTokenizer {
         this.#specialTokenIds.add(id);
       }
     }
+    this.#specialTokenPatterns.sort((a, b) => b.content.length - a.content.length);
+    if (maxId >= 0) {
+      this.vocabSize = Math.max(this.vocabSize, maxId + 1);
+    }
 
     const runtimeDefaults = getRuntimeConfig().inference.tokenizer;
+    const byteLevelPretokenizer = resolveByteLevelPretokenizerConfig(tokenizerJson.pre_tokenizer);
     const configuredAddBosToken = this.addBosToken;
     const configuredAddEosToken = this.addEosToken;
     this.addBosToken =
@@ -505,9 +637,11 @@ export class BundledTokenizer extends BaseTokenizer {
     if (this.addEosToken && this.specialTokens.eos == null) {
       throw new Error('[Tokenizer] addEosToken is enabled but eos token is missing.');
     }
+    this.#useByteLevelEncoding = byteLevelPretokenizer.useByteLevel;
     // NOTE: Default to FALSE - first word shouldn't get space prefix
     // Space prefixes are only for words that follow a space in original text
-    this.#addSpacePrefix = tokenizerJson.addSpacePrefix === true;
+    this.#addSpacePrefix = tokenizerJson.addSpacePrefix === true
+      || byteLevelPretokenizer.addPrefixSpace === true;
 
     // Detect space prefix style based on vocab tokens
     // GPT-style uses 'Ġ' (U+0120), SentencePiece uses '▁' (U+2581)
@@ -548,7 +682,8 @@ export class BundledTokenizer extends BaseTokenizer {
       ids.push(this.specialTokens.bos);
     }
 
-    // Split text around special tokens and tokenize each segment
+    // Split text around literal added tokens and special tokens, then tokenize
+    // the remaining plain-text segments normally.
     const segments = this.#splitOnSpecialTokens(text);
     for (const seg of segments) {
       if (seg.isSpecial && seg.id !== undefined) {
@@ -690,11 +825,19 @@ export class BundledTokenizer extends BaseTokenizer {
     if (text.length === 0) return [];
 
     let normalized = text;
-    if (this.#addSpacePrefix && !normalized.startsWith(' ')) {
-      normalized = ` ${normalized}`;
+    let prefixed;
+    if (this.#useByteLevelEncoding) {
+      if (this.#addSpacePrefix && !normalized.startsWith(' ')) {
+        normalized = ` ${normalized}`;
+      }
+      prefixed = this.#encodeByteLevelText(normalized);
+    } else {
+      if (this.#addSpacePrefix && !normalized.startsWith(' ')) {
+        normalized = ` ${normalized}`;
+      }
+      const sp = this.#spacePrefixChar;
+      prefixed = normalized.replace(/ /g, sp);
     }
-    const sp = this.#spacePrefixChar;
-    const prefixed = normalized.replace(/ /g, sp);
 
     if (this.#mergeRanks.size === 0) {
       return this.#encodeBPEGreedy(prefixed);

package/src/inference/tokenizers/sentencepiece.js

@@ -31,8 +31,18 @@ export class SentencePieceTokenizer extends BaseTokenizer {
     });
   }
 
+  #resetState() {
+    this.#modelData = null;
+    this.#pieces.clear();
+    this.#reverseVocab.clear();
+    this.#algorithm = 'unigram';
+    this.#byteTokens.clear();
+    this.vocabSize = 0;
+  }
+
   
   async load(modelData) {
+    this.#resetState();
     this.#modelData = modelData;
 
     try {
@@ -42,6 +52,8 @@ export class SentencePieceTokenizer extends BaseTokenizer {
     } catch (err) {
       const message = err instanceof Error ? err.message : String(err);
       log.warn('Tokenizer', `Failed to parse model, using byte fallback: ${message}`);
+      this.#resetState();
+      this.#modelData = modelData;
       this.#initByteFallback();
     }
   }

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