@simulatte/doppler 0.1.8 → 0.2.0

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

@@ -1,268 +0,0 @@
-
-
-import {
-  createWeightBuffer,
-  createCpuWeightBuffer,
-  isWeightBuffer,
-  isCpuWeightBuffer,
-} from '../gpu/weight-buffer.js';
-import { maybeDowncastToF16 } from './weight-downcast.js';
-import { getTensorNamesByRole } from './tensors/tensor-role.js';
-import { log, trace as debugTrace } from '../debug/index.js';
-import { selectRuleValue } from '../rules/rule-registry.js';
-
-// ============================================================================
-// Constants
-// ============================================================================
-
-
-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;
-  if (/layers?[._]\d+/.test(lower)) return false;
-  if (lower.includes('input_layernorm')) return false;
-  if (lower.includes('post_attention_layernorm')) return false;
-  if (lower.includes('pre_feedforward_layernorm')) return false;
-  if (lower.includes('post_feedforward_layernorm')) return false;
-
-  return (
-    lower === 'norm.weight' ||
-    lower.includes('model.norm.weight') ||
-    lower.includes('language_model.norm.weight') ||
-    lower.includes('model.language_model.norm.weight') ||
-    lower.includes('embedding_norm.weight') ||
-    lower.includes('model.embedding_norm.weight') ||
-    lower.includes('final_layernorm.weight') ||
-    lower.includes('final_layer_norm.weight') ||
-    lower.includes('norm_f.weight')
-  );
-}
-
-// ============================================================================
-// Main Function
-// ============================================================================
-
-
-export async function loadFinalWeights(ctx) {
-  let normOffsetDebugLogged = ctx.normOffsetDebugLogged;
-
-  // Load final norm
-  const { finalNorm, debugLogged: normDebugLogged } = await loadFinalNorm(ctx);
-  if (normDebugLogged) {
-    normOffsetDebugLogged = true;
-  }
-
-  // Load LM head
-  const lmHead = await loadLmHead(ctx);
-
-  return {
-    finalNorm,
-    lmHead,
-    normOffsetDebugLogged,
-  };
-}
-
-// ============================================================================
-// Final Norm Loading
-// ============================================================================
-
-
-async function loadFinalNorm(ctx) {
-  
-  let finalNorm = null;
-  let debugLogged = false;
-
-  let finalNormNames = getTensorNamesByRole(ctx.tensorLocations, FINAL_NORM_ROLE, HEAD_GROUP);
-  if (finalNormNames.length === 0) {
-    const legacyCandidates = getTensorNamesByRole(ctx.tensorLocations, FINAL_NORM_ROLE).filter(
-      (name) => isLikelyFinalNormName(name)
-    );
-    if (legacyCandidates.length > 0) {
-      finalNormNames = legacyCandidates;
-      log.warn(
-        'Loader',
-        '[FinalNorm] Falling back to role-only final norm selection because tensor groups are missing in manifest.'
-      );
-    }
-  }
-  if (finalNormNames.length === 0) {
-    throw new Error(
-      `[Loader] Final norm not found. Expected tensor with role="${FINAL_NORM_ROLE}" and group="${HEAD_GROUP}".`
-    );
-  }
-
-  for (const name of finalNormNames) {
-    const location = ctx.tensorLocations.get(name);
-    if (location) {
-      finalNorm =  (await ctx.loadTensor(name, true, true));
-      break;
-    }
-  }
-
-  if (finalNorm && ctx.needsNormWeightOffset() && !ctx.normOffsetDebugLogged) {
-    debugTrace.loader('Final norm uses RMSNorm weight offset (applied at runtime)');
-    debugLogged = true;
-  }
-
-  if (!finalNorm) {
-    throw new Error(
-      `[Loader] Final norm not found. Tried: ${finalNormNames.join(', ')}`
-    );
-  }
-
-  return { finalNorm, debugLogged };
-}
-
-// ============================================================================
-// LM Head Loading
-// ============================================================================
-
-
-async function loadLmHead(ctx) {
-  
-  let lmHead = null;
-  
-  let lmHeadName = null;
-  
-  let lmHeadLoc;
-
-  const lmHeadNames = getTensorNamesByRole(ctx.tensorLocations, LM_HEAD_ROLE, HEAD_GROUP);
-  if (lmHeadNames.length === 0 && !ctx.tieWordEmbeddings) {
-    throw new Error(
-      `[Loader] LM head not found. Expected tensor with role="${LM_HEAD_ROLE}" and group="${HEAD_GROUP}".`
-    );
-  }
-
-  for (const name of lmHeadNames) {
-    const loc = ctx.tensorLocations.get(name);
-    if (!loc) continue;
-
-    const shouldStream = ctx.shouldStreamLargeWeight(name, loc, 'LM head');
-    const tensor = await ctx.loadTensor(name, !shouldStream, true);
-
-    if (shouldStream && tensor && !(tensor instanceof Float32Array)) {
-      throw new Error(
-        `[Loader] LM head "${name}" too large for GPU and cannot be loaded on CPU (dtype=${loc.dtype}).`
-      );
-    }
-
-    if (tensor && (isGpuBufferInstance(tensor) || isWeightBuffer(tensor) || tensor instanceof Float32Array)) {
-      lmHeadName = name;
-      lmHeadLoc = loc;
-      lmHead = processLmHeadTensor(ctx, tensor, name, loc, shouldStream);
-      break;
-    }
-  }
-
-  // Use tied embeddings as fallback
-  if (!lmHead && ctx.embeddings && ctx.tieWordEmbeddings) {
-    debugTrace.loader('Using tied embeddings as LM head (manifest.tieWordEmbeddings=true)');
-    lmHead = ctx.embeddings;
-  } else if (!lmHead) {
-    throw new Error(
-      `[Loader] LM head not found. Tried: ${lmHeadNames.join(', ')}`
-    );
-  }
-
-  // Downcast LM head to F16 if applicable
-  if (lmHead && !isCpuWeightBuffer(lmHead)) {
-    lmHead = await maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc);
-  }
-
-  return lmHead;
-}
-
-
-function processLmHeadTensor(ctx, tensor, name, loc, shouldStream) {
-  // Float32Array streaming path
-  if (tensor instanceof Float32Array && shouldStream) {
-    const layout = ctx.resolveWeightLayout(loc);
-    
-    const dtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
-      locationDtype: loc.dtype,
-    });
-    const result = createCpuWeightBuffer(tensor, dtype, layout, loc.shape, name);
-    log.warn('Loader', `LM head stored on CPU for chunked matmul (layout=${layout})`);
-    return result;
-  }
-
-  // Raw GPUBuffer - wrap with dtype/layout metadata
-  if (isGpuBufferInstance(tensor) && loc.shape && loc.shape.length === 2) {
-    const layout = ctx.resolveWeightLayout(loc);
-    
-    const dtype = selectRuleValue('loader', 'weights', 'floatLocationDtype', {
-      locationDtype: loc.dtype,
-    });
-    const wrapped = createWeightBuffer(tensor, dtype, layout, loc.shape, name);
-    log.info('Loader', `Wrapped lm_head as WeightBuffer (layout=${layout}, dtype=${dtype})`);
-    return wrapped;
-  }
-
-  return tensor;
-}
-
-
-async function maybeDowncastLmHead(ctx, lmHead, lmHeadName, lmHeadLoc) {
-  // Check if tied to embeddings (skip downcast to avoid double-processing)
-  const tiedToEmbeddings =
-    lmHead === ctx.embeddings ||
-    (isWeightBuffer(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead.buffer === ctx.embeddings.buffer) ||
-    (isGpuBufferInstance(lmHead) && isWeightBuffer(ctx.embeddings) && lmHead === ctx.embeddings.buffer);
-
-  if (tiedToEmbeddings) {
-    return lmHead;
-  }
-
-  // Can't downcast Float32Array
-  if (lmHead instanceof Float32Array) {
-    return lmHead;
-  }
-
-  // Get current dtype
-  const dtype = isWeightBuffer(lmHead)
-    ? lmHead.dtype
-    : selectRuleValue('loader', 'weights', 'floatLocationDtype', {
-      locationDtype: lmHeadLoc?.dtype,
-    });
-
-  // Skip if not F32
-  if (dtype !== 'f32') {
-    return lmHead;
-  }
-
-  // Get buffer for downcast
-  const buffer = isWeightBuffer(lmHead) ? lmHead.buffer : lmHead;
-  if (!isGpuBufferInstance(buffer)) {
-    return lmHead;
-  }
-
-  const elems = buffer.size / 4;
-
-  // Attempt downcast
-  const result = await maybeDowncastToF16(lmHead, {
-    label: lmHeadName ?? 'lm_head',
-    keepF32: ctx.keepF32Weights,
-    dtype,
-    shape: isWeightBuffer(lmHead)
-      ? Array.from(lmHead.shape)
-      : (lmHeadLoc?.shape ?? [elems]),
-    layout: isWeightBuffer(lmHead)
-      ? lmHead.layout
-      : (lmHeadLoc ? ctx.resolveWeightLayout(lmHeadLoc) : 'row'),
-  });
-
-  if (result?.wasDowncast && result.newBuffer) {
-    ctx.gpuBuffers.add(result.newBuffer);
-    return  (result.buffer);
-  }
-
-  return lmHead;
-}

package/src/loader/index.d.ts

@@ -1,150 +0,0 @@
-// ============================================================================
-// Core Loader
-// ============================================================================
-
-export {
-  DopplerLoader,
-  getDopplerLoader,
-  createDopplerLoader,
-} from './doppler-loader.js';
-
-export type {
-  TensorLocation,
-  LayerWeights,
-  LoadProgress,
-  LoadOptions,
-  CustomShardLoader,
-  CustomShardLoaderOptions,
-  LoaderStats,
-  KernelCapabilities,
-  Q4KConfig,
-  ModelConfig,
-  ShardLoadPriority,
-  ShardLoadOptions,
-  CustomShardRangeLoader,
-  CustomShardStreamLoader,
-  CustomShardStreamOptions,
-  ShardSourceInfo,
-} from './loader-types.js';
-
-export { MultiModelLoader } from './multi-model-loader.js';
-export type { AdapterSource } from './multi-model-loader.js';
-
-export { LoaderState, createLoaderState } from './loader-state.js';
-export type { WeightType, LoaderStateSnapshot } from './loader-state.js';
-
-// ============================================================================
-// Manifest & Shard Resolution
-// ============================================================================
-
-export {
-  needsNormWeightOffset,
-  getLargeWeightConfig,
-  getLargeWeightMaxBytes,
-  estimateMatmulWeightBytes,
-  resolveWeightLayout,
-  shouldStreamLargeWeight,
-  isMoEModel,
-} from './manifest-config.js';
-export type { LargeWeightConfig } from './manifest-config.js';
-
-export { buildTensorLocations } from './shard-resolver.js';
-export type { BuildTensorLocationsOptions } from './shard-resolver.js';
-
-export { ShardCache, createShardCache } from './shard-cache.js';
-export type { ShardCacheConfig } from './shard-cache.js';
-
-// ============================================================================
-// Expert Loading & Caching
-// ============================================================================
-
-export {
-  preloadShardsForExpert,
-  prefetchExperts,
-  predictNextLayerExperts,
-  loadExpert,
-} from './experts/expert-loader.js';
-
-export {
-  ExpertCache,
-  getExpertCache,
-  createExpertCache,
-} from './experts/expert-cache.js';
-export type { CacheStats } from './experts/expert-cache.js';
-
-// ============================================================================
-// Tensor Loading
-// ============================================================================
-
-export {
-  isPackedQ4K,
-  shouldUseFusedQ4K,
-  getQ4KOutputDtype,
-  getWeightLayout,
-  convertBF16ToF32CPU,
-  convertF16ToF32CPU,
-  loadQ4KFused,
-  loadQ4KDequant,
-  loadQ6K,
-  loadBF16,
-  loadFloat,
-  loadTensorToGPU,
-  loadTensorToCPU,
-} from './tensors/tensor-loader.js';
-export type { TensorLoadConfig, TensorLoadResult } from './tensors/tensor-loader.js';
-
-export { assembleShardData } from './tensors/tensor-reader.js';
-
-export { getTensorNamesByRole } from './tensors/tensor-role.js';
-
-// ============================================================================
-// Weight Stage Loaders
-// ============================================================================
-
-export { loadEmbeddings } from './embedding-loader.js';
-
-export { loadLayer } from './layer-loader.js';
-
-export { loadFinalWeights } from './final-weights-loader.js';
-
-// ============================================================================
-// Utilities
-// ============================================================================
-
-export {
-  f16ToF32,
-  convertBF16ToF32GPU,
-  shouldDequantizeToF16,
-  applyBufferLayout,
-} from './dtype-utils.js';
-
-export {
-  maybeDowncastToF16,
-  batchDowncastWeights,
-} from './weight-downcast.js';
-
-export {
-  captureMemorySnapshot,
-  formatMemoryStats,
-  MemoryMonitor,
-  MemoryTimeSeries,
-} from './memory-monitor.js';
-export type { MemorySnapshot } from './memory-monitor.js';
-
-export {
-  QK_K,
-  QK4_K_BLOCK_SIZE,
-  K_SCALE_SIZE,
-  Q4K_BLOCK_BYTES,
-  Q6K_BLOCK_BYTES,
-  Q8_0_BLOCK_BYTES,
-  Q8_0_BLOCK_SIZE,
-  padToQ4KBlock,
-  q4kBlockCount,
-} from './quantization-constants.js';
-
-// ============================================================================
-// Type-only re-exports
-// ============================================================================
-
-export type { ExpertWeights } from './weights.js';

package/src/loader/index.js

@@ -1,124 +0,0 @@
-// ============================================================================
-// Core Loader
-// ============================================================================
-
-export {
-  DopplerLoader,
-  getDopplerLoader,
-  createDopplerLoader,
-} from './doppler-loader.js';
-
-export { MultiModelLoader } from './multi-model-loader.js';
-
-export { LoaderState, createLoaderState } from './loader-state.js';
-
-// ============================================================================
-// Manifest & Shard Resolution
-// ============================================================================
-
-export {
-  needsNormWeightOffset,
-  getLargeWeightConfig,
-  getLargeWeightMaxBytes,
-  estimateMatmulWeightBytes,
-  resolveWeightLayout,
-  shouldStreamLargeWeight,
-  isMoEModel,
-} from './manifest-config.js';
-
-export { buildTensorLocations } from './shard-resolver.js';
-
-export { ShardCache, createShardCache } from './shard-cache.js';
-
-// ============================================================================
-// Expert Loading & Caching
-// ============================================================================
-
-export {
-  preloadShardsForExpert,
-  prefetchExperts,
-  predictNextLayerExperts,
-  loadExpert,
-} from './experts/expert-loader.js';
-
-export {
-  ExpertCache,
-  getExpertCache,
-  createExpertCache,
-} from './experts/expert-cache.js';
-
-// ============================================================================
-// Tensor Loading
-// ============================================================================
-
-export {
-  isPackedQ4K,
-  shouldUseFusedQ4K,
-  getQ4KOutputDtype,
-  getWeightLayout,
-  convertBF16ToF32CPU,
-  convertF16ToF32CPU,
-  loadQ4KFused,
-  loadQ4KDequant,
-  loadQ6K,
-  loadBF16,
-  loadFloat,
-  loadTensorToGPU,
-  loadTensorToCPU,
-} from './tensors/tensor-loader.js';
-
-export { assembleShardData } from './tensors/tensor-reader.js';
-
-export { getTensorNamesByRole } from './tensors/tensor-role.js';
-
-// ============================================================================
-// Weight Stage Loaders
-// ============================================================================
-
-export { loadEmbeddings } from './embedding-loader.js';
-
-export { loadLayer } from './layer-loader.js';
-
-export { loadFinalWeights } from './final-weights-loader.js';
-
-// ============================================================================
-// Utilities
-// ============================================================================
-
-export {
-  f16ToF32,
-  convertBF16ToF32GPU,
-  shouldDequantizeToF16,
-  applyBufferLayout,
-} from './dtype-utils.js';
-
-export {
-  maybeDowncastToF16,
-  batchDowncastWeights,
-} from './weight-downcast.js';
-
-export {
-  captureMemorySnapshot,
-  formatMemoryStats,
-  MemoryMonitor,
-  MemoryTimeSeries,
-} from './memory-monitor.js';
-
-export {
-  QK_K,
-  QK4_K_BLOCK_SIZE,
-  K_SCALE_SIZE,
-  Q4K_BLOCK_BYTES,
-  Q6K_BLOCK_BYTES,
-  Q8_0_BLOCK_BYTES,
-  Q8_0_BLOCK_SIZE,
-  padToQ4KBlock,
-  q4kBlockCount,
-} from './quantization-constants.js';
-
-// ============================================================================
-// Type-only modules (re-exported for barrel completeness)
-// ============================================================================
-
-export {} from './loader-types.js';
-export {} from './weights.js';

package/src/loader/layer-loader.d.ts

@@ -1,63 +0,0 @@
-/**
- * Layer Loader - Load transformer layer weights.
- *
- * Handles loading of all weights for a single transformer layer:
- * - Attention weights (Q, K, V, O projections)
- * - Norm weights (input, post-attention, FFN norms)
- * - FFN weights (gate, up, down projections)
- * - MoE router weights
- *
- * @module loader/layer-loader
- */
-
-import type { WeightBuffer } from '../gpu/weight-buffer.js';
-import type { TensorLocation, LayerWeights } from './loader-types.js';
-
-/** Tensor loading function signature */
-export type TensorLoader = (
-  name: string,
-  toGPU?: boolean,
-  silent?: boolean
-) => Promise<GPUBuffer | WeightBuffer | Float32Array | Uint8Array | null>;
-
-/**
- * Context required for layer loading.
- */
-export interface LayerLoaderContext {
-  /** Tensor locations map */
-  tensorLocations: Map<string, TensorLocation>;
-  /** Load a tensor by name */
-  loadTensor: TensorLoader;
-  /** Check if model needs norm weight offset */
-  needsNormWeightOffset: () => boolean;
-  /** GPU buffers to track for cleanup */
-  gpuBuffers: Set<GPUBuffer>;
-  /** Keep F32 weights (skip downcast) */
-  keepF32Weights: boolean;
-  /** Whether model is MoE */
-  isMoE: boolean;
-  /** Check if layer is an expert layer */
-  isExpertLayer: (layerIdx: number) => boolean;
-  /** Optional attention dimensions for fused QKV sizing */
-  numHeads?: number;
-  numKVHeads?: number;
-  headDim?: number;
-  hiddenSize?: number;
-  /** Optional linear-attention dimensions for hybrid models (e.g. Qwen3.5) */
-  linearNumKeyHeads?: number;
-  linearNumValueHeads?: number;
-  linearKeyHeadDim?: number;
-  linearValueHeadDim?: number;
-}
-
-/**
- * Load all weights for a single transformer layer.
- *
- * @param ctx - Layer loader context
- * @param layerIdx - Layer index
- * @returns Loaded layer weights
- */
-export declare function loadLayer(
-  ctx: LayerLoaderContext,
-  layerIdx: number
-): Promise<LayerWeights>;