@simulatte/doppler 0.1.0

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

@@ -0,0 +1,236 @@
+/**
+ * Loader Types
+ *
+ * Type definitions for the DopplerLoader.
+ *
+ * @module loader/loader-types
+ */
+
+import type { WeightBuffer } from '../gpu/weight-buffer.js';
+import type { TensorRole } from '../config/schema/index.js';
+
+/**
+ * Tensor location in loaded model
+ */
+export interface TensorLocation {
+  shardIndex: number;
+  offset: number;
+  size: number;
+  shape: number[];
+  dtype: string;
+  role: TensorRole;
+  group?: string;
+  spans?: Array<{ shardIndex: number; offset: number; size: number }>;
+  /** Weight storage layout: 'column' means pre-transposed for faster matmul */
+  layout?: 'row' | 'column';
+  /** Original shape before transpose (if layout is 'column') */
+  originalShape?: number[];
+}
+
+/**
+ * Loaded layer weights
+ */
+export interface LayerWeights {
+  inputNorm: GPUBuffer | Float32Array | null;
+  qProj: GPUBuffer | WeightBuffer | Float32Array | null;
+  kProj: GPUBuffer | WeightBuffer | Float32Array | null;
+  vProj: GPUBuffer | WeightBuffer | Float32Array | null;
+  oProj: GPUBuffer | WeightBuffer | Float32Array | null;
+  qkvProj?: GPUBuffer | WeightBuffer | Float32Array | null;
+  qkvSizes?: [number, number, number] | null;
+  qkvDtype?: 'f16' | 'f32' | null;
+  linearInProjZ?: GPUBuffer | WeightBuffer | Float32Array | null;
+  linearInProjA?: GPUBuffer | WeightBuffer | Float32Array | null;
+  linearInProjB?: GPUBuffer | WeightBuffer | Float32Array | null;
+  linearConv1D?: GPUBuffer | Float32Array | null;
+  linearDtBias?: GPUBuffer | Float32Array | null;
+  linearALog?: GPUBuffer | Float32Array | null;
+  linearNorm?: GPUBuffer | Float32Array | null;
+  qNorm: GPUBuffer | Float32Array | null;
+  kNorm: GPUBuffer | Float32Array | null;
+  postAttentionNorm: GPUBuffer | Float32Array | null;
+  preFeedforwardNorm: GPUBuffer | Float32Array | null;
+  postFeedforwardNorm: GPUBuffer | Float32Array | null;
+  postNorm: GPUBuffer | Float32Array | null;
+  postAttnNorm: GPUBuffer | Float32Array | null;
+  convInProj?: GPUBuffer | WeightBuffer | Float32Array | null;
+  convKernel?: GPUBuffer | WeightBuffer | Float32Array | null;
+  convOutProj?: GPUBuffer | WeightBuffer | Float32Array | null;
+  ffnGate: GPUBuffer | WeightBuffer | Float32Array | null;
+  ffnUp: GPUBuffer | WeightBuffer | Float32Array | null;
+  ffnDown: GPUBuffer | WeightBuffer | Float32Array | null;
+  /** Fused gate+up projection [intermediateSize*2, hiddenSize] for 2-pass FFN */
+  ffnGateUp?: GPUBuffer | WeightBuffer | Float32Array | null;
+  // Aliases for pipeline compatibility
+  gate?: GPUBuffer | WeightBuffer | Float32Array | null;
+  up?: GPUBuffer | WeightBuffer | Float32Array | null;
+  down?: GPUBuffer | WeightBuffer | Float32Array | null;
+  /** Fused gate+up for pipeline compatibility */
+  gateUp?: GPUBuffer | WeightBuffer | Float32Array | null;
+  routerWeight?: GPUBuffer | import('../gpu/weight-buffer.js').WeightBuffer | Float32Array | null;
+  routerBias?: GPUBuffer | Float32Array | null;
+  attentionSinks?: GPUBuffer | Float32Array | null;
+}
+
+/**
+ * Loading progress information
+ */
+export interface LoadProgress {
+  stage: 'manifest' | 'shards' | 'layers' | 'gpu_transfer' | 'complete';
+  progress: number;
+  /** Current layer index */
+  layer?: number;
+  /** Total layers */
+  total?: number;
+  /** Current shard index */
+  shard?: number;
+  /** Total shards */
+  totalShards?: number;
+  /** Bytes loaded so far */
+  bytesLoaded?: number;
+  /** Total bytes to load */
+  totalBytes?: number;
+  /** Loading speed in bytes per second */
+  bytesPerSecond?: number;
+  /** Human-readable message */
+  message?: string;
+}
+
+/**
+ * Loading options
+ */
+export interface LoadOptions {
+  onProgress?: (progress: LoadProgress) => void;
+  verifyHashes: boolean;
+}
+
+/**
+ * Shard load priority.
+ */
+export type ShardLoadPriority = 'high' | 'low';
+
+/**
+ * Shard loading options.
+ */
+export interface ShardLoadOptions {
+  priority?: ShardLoadPriority;
+}
+
+/**
+ * Custom shard loader options
+ */
+export interface CustomShardLoaderOptions {
+  verify?: boolean;
+  loadShardRange?: CustomShardRangeLoader;
+  streamShardRange?: CustomShardStreamLoader;
+}
+
+/**
+ * Custom shard loader function
+ */
+export type CustomShardLoader = (
+  shardIndex: number
+) => Promise<ArrayBuffer | Uint8Array>;
+
+/**
+ * Custom shard range loader function
+ */
+export type CustomShardRangeLoader = (
+  shardIndex: number,
+  offset: number,
+  length?: number | null
+) => Promise<ArrayBuffer | Uint8Array>;
+
+/**
+ * Custom shard range stream options
+ */
+export interface CustomShardStreamOptions {
+  chunkBytes?: number;
+}
+
+/**
+ * Custom shard range stream loader function
+ */
+export type CustomShardStreamLoader = (
+  shardIndex: number,
+  offset?: number,
+  length?: number | null,
+  options?: CustomShardStreamOptions
+) => AsyncIterable<Uint8Array>;
+
+/**
+ * Loader statistics
+ */
+export interface LoaderStats {
+  modelId: string | null;
+  isLoaded: boolean;
+  isMoE: boolean;
+  isUnifiedMemory: boolean;
+  layersLoaded: number;
+  expertsLoaded: number;
+  gpuBuffers: number;
+}
+
+/**
+ * GPU kernel capabilities
+ */
+export interface KernelCapabilities {
+  hasF16: boolean;
+  hasSubgroups: boolean;
+}
+
+/**
+ * Q4K loading configuration.
+ */
+export interface Q4KConfig {
+  /** Use fused Q4K matmul kernels (keeps raw quantized weights) */
+  useFusedQ4K: boolean;
+  /** Q4K layout: 'row' = fused kernel (fast), 'col' = dequant fallback */
+  q4kLayout: 'row' | 'col' | null;
+  /** Keep weights as F32 (disable F16 downcasting) */
+  keepF32Weights: boolean;
+}
+
+/**
+ * Model config (flexible structure from manifest)
+ */
+export interface ModelConfig {
+  num_hidden_layers?: number;
+  blockCount?: number;
+  text_config?: { num_hidden_layers?: number };
+  n_layer?: number;
+  num_local_experts?: number;
+  num_experts?: number;
+  architectures?: string[];
+  model_type?: string;
+  [key: string]: unknown;
+}
+
+/**
+ * Shard source tracking
+ */
+export interface ShardSourceInfo {
+  source: 'RAM' | 'OPFS' | 'custom' | 'network' | 'indexeddb' | 'memory' | 'storage' | string;
+  elapsed: number;
+  mode?: 'full' | 'range' | 'stream';
+  path?:
+    | 'cache'
+    | 'custom-loader'
+    | 'custom-range'
+    | 'custom-stream'
+    | 'custom-loader-slice'
+    | 'custom-range-fallback'
+    | 'backend-full'
+    | 'backend-range'
+    | 'backend-stream';
+  fallback?:
+    | 'none'
+    | 'custom_range_unavailable'
+    | 'custom_range_not_supported'
+    | 'custom_stream_not_supported'
+    | 'custom_stream_not_supported_resume'
+    | 'custom_stream_interrupted'
+    | 'custom_stream_interrupted_resume'
+    | 'custom_stream_partial_resume'
+    | 'custom_range_partial_retry';
+}

package/src/loader/loader-types.js

@@ -0,0 +1,4 @@
+
+
+// Types-only module - all types are defined in loader-types.d.ts
+export {};

package/src/loader/manifest-config.d.ts

@@ -0,0 +1,97 @@
+/**
+ * Manifest Config - Model configuration resolution from manifest.
+ *
+ * Pure functions for extracting configuration from manifests:
+ * - Norm weight offset detection (Gemma models)
+ * - Large weight handling configuration
+ * - Weight layout resolution
+ *
+ * @module loader/manifest-config
+ */
+
+import type { RDRRManifest } from '../formats/rdrr/index.js';
+import type { TensorLocation, KernelCapabilities } from './loader-types.js';
+import type { WeightLayout, WeightDtype } from '../gpu/weight-buffer.js';
+
+/**
+ * Check if model requires (1 + weight) offset for RMSNorm weights.
+ *
+ * GGUF files do NOT have the offset baked in - they store raw weights.
+ * RMSNorm applies the +1 offset at runtime when this flag is true.
+ *
+ * @param manifest - Model manifest
+ * @returns Whether norm weight offset is needed
+ */
+export declare function needsNormWeightOffset(manifest: RDRRManifest | null): boolean;
+
+export interface LargeWeightConfig {
+  enabled: boolean;
+  safetyRatio: number;
+  preferF16: boolean;
+}
+
+/**
+ * Get large weight handling configuration from runtime config.
+ */
+export declare function getLargeWeightConfig(): LargeWeightConfig;
+
+/**
+ * Get maximum bytes for a single GPU buffer binding.
+ *
+ * @returns Max bytes, or null if large weight handling is disabled
+ */
+export declare function getLargeWeightMaxBytes(): number | null;
+
+/**
+ * Estimate GPU memory required for a matmul weight after dequantization.
+ *
+ * @param location - Tensor location info
+ * @param gpuCapabilities - GPU capabilities
+ * @param keepF32Weights - Whether to keep F32 (skip F16 downcast)
+ * @returns Estimated bytes and output dtype, or null if cannot estimate
+ */
+export declare function estimateMatmulWeightBytes(
+  location: TensorLocation,
+  gpuCapabilities: KernelCapabilities | null,
+  keepF32Weights: boolean
+): { bytes: number; dtype: WeightDtype } | null;
+
+/**
+ * Resolve weight layout from tensor location.
+ *
+ * Column layout is used for:
+ * - Explicit layout='column' in tensor info
+ * - Embeddings with transposed shape (dim0 < dim1)
+ *
+ * @param location - Tensor location info
+ * @returns Weight layout ('row' or 'column')
+ */
+export declare function resolveWeightLayout(location: TensorLocation): WeightLayout;
+
+/**
+ * Check if a large weight should use CPU streaming instead of GPU buffer.
+ *
+ * Logs a warning if the weight exceeds GPU limits and provides guidance.
+ *
+ * @param name - Tensor name
+ * @param location - Tensor location info
+ * @param label - Human-readable label for logging (e.g., 'Embedding', 'LM head')
+ * @param gpuCapabilities - GPU capabilities
+ * @param keepF32Weights - Whether to keep F32
+ * @returns Whether to use streaming
+ */
+export declare function shouldStreamLargeWeight(
+  name: string,
+  location: TensorLocation,
+  label: string,
+  gpuCapabilities: KernelCapabilities | null,
+  keepF32Weights: boolean
+): boolean;
+
+/**
+ * Check if model uses Mixture of Experts architecture.
+ *
+ * @param manifest - Model manifest
+ * @returns Whether model is MoE
+ */
+export declare function isMoEModel(manifest: RDRRManifest | null): boolean;

package/src/loader/manifest-config.js

@@ -0,0 +1,132 @@
+
+import { getDevice, getKernelCapabilities } from '../gpu/device.js';
+import { getRuntimeConfig } from '../config/runtime.js';
+import { DTYPE_SIZES } from '../config/schema/index.js';
+import { shouldDequantizeToF16 } from './dtype-utils.js';
+import { formatBytes } from '../storage/quota.js';
+import { log, trace as debugTrace } from '../debug/index.js';
+import { selectRuleValue } from '../rules/rule-registry.js';
+
+// ============================================================================
+// Norm Weight Offset Detection
+// ============================================================================
+
+export function needsNormWeightOffset(manifest) {
+  if (!manifest) {
+    debugTrace.loader('_needsNormWeightOffset: no manifest');
+    return false;
+  }
+
+  const inferenceFlag = manifest.inference?.normalization?.rmsNormWeightOffset;
+  if (inferenceFlag == null) {
+    const modelId = manifest.modelId ?? 'unknown';
+    throw new Error(
+      `Manifest "${modelId}" is missing inference.normalization.rmsNormWeightOffset. ` +
+      'Re-convert the model with a complete manifest.inference config.'
+    );
+  }
+
+  if (inferenceFlag) {
+    debugTrace.loader('RMSNorm weight offset enabled (manifest.inference.normalization.rmsNormWeightOffset=true)');
+  }
+  return inferenceFlag;
+}
+
+// ============================================================================
+// Large Weight Handling
+// ============================================================================
+
+export function getLargeWeightConfig() {
+  const config = getRuntimeConfig().inference.largeWeights;
+  if (!config) {
+    throw new Error('runtime.inference.largeWeights is required');
+  }
+  return config;
+}
+
+export function getLargeWeightMaxBytes() {
+  const config = getLargeWeightConfig();
+  if (!config.enabled) return null;
+
+  const device = getDevice();
+  if (!device) return null;
+
+  const safety = Math.min(Math.max(config.safetyRatio, 0.1), 1);
+  const maxBinding = Math.min(
+    device.limits.maxStorageBufferBindingSize,
+    device.limits.maxBufferSize
+  );
+  return Math.floor(maxBinding * safety);
+}
+
+export function estimateMatmulWeightBytes(location, gpuCapabilities, keepF32Weights) {
+  if (!location.shape || location.shape.length === 0) return null;
+
+  const numElements = location.shape.reduce((a, b) => a * b, 1);
+  if (!Number.isFinite(numElements) || numElements <= 0) return null;
+
+  const caps = gpuCapabilities || getKernelCapabilities();
+  const hasF16 = caps?.hasF16 ?? false;
+  const isMatmulWeight = shouldDequantizeToF16(location);
+
+  const dtype = selectRuleValue('loader', 'weights', 'matmulWeightDtype', {
+    locationDtype: location.dtype,
+    hasF16,
+    isMatmulWeight,
+    keepF32Weights: Boolean(keepF32Weights),
+  });
+
+  const bytesPerElement = DTYPE_SIZES[selectRuleValue('shared', 'dtype', 'f16OrF32FromDtype', { dtype })];
+  return { bytes: numElements * bytesPerElement, dtype };
+}
+
+export function resolveWeightLayout(location) {
+  const isEmbedding = location.role === 'embedding' || location.role === 'lm_head';
+  const useColumnWise = isEmbedding && location.shape?.length === 2
+    ? location.shape[0] < location.shape[1]
+    : false;
+
+  return selectRuleValue('loader', 'weights', 'weightLayout', {
+    layout: location.layout ?? null,
+    useColumnWise,
+  });
+}
+
+export function shouldStreamLargeWeight(name, location, label, gpuCapabilities, keepF32Weights) {
+  const maxBytes = getLargeWeightMaxBytes();
+  if (!maxBytes) return false;
+
+  const estimate = estimateMatmulWeightBytes(location, gpuCapabilities, keepF32Weights);
+  if (!estimate) return false;
+
+  if (estimate.bytes <= maxBytes) return false;
+
+  // Check if dtype can be streamed (only float types)
+  const canStream = location.dtype === 'F16' || location.dtype === 'F32' || location.dtype === 'BF16';
+  if (!canStream) {
+    log.warn(
+      'Loader',
+      `${label} weight "${name}" (${formatBytes(estimate.bytes)}) exceeds GPU binding limit (${formatBytes(maxBytes)}) ` +
+      `but dtype ${location.dtype} cannot be streamed. Regenerate with F16/F32 weights.`
+    );
+    return false;
+  }
+
+  log.warn(
+    'Loader',
+    `${label} weight "${name}" (${formatBytes(estimate.bytes)}) exceeds GPU binding limit (${formatBytes(maxBytes)}). ` +
+    'Using CPU-backed streaming.'
+  );
+  return true;
+}
+
+// ============================================================================
+// MoE Detection
+// ============================================================================
+
+export function isMoEModel(manifest) {
+  if (!manifest) return false;
+
+  // Explicit MoE config
+  return (manifest.moeConfig?.numExperts ?? 0) > 1;
+}

package/src/loader/memory-monitor.d.ts

@@ -0,0 +1,112 @@
+/**
+ * Memory Monitor - Memory statistics and logging during model loading.
+ *
+ * Provides utilities for tracking memory usage across JS heap, GPU buffers,
+ * and shard cache during the model loading process.
+ *
+ * @module loader/memory-monitor
+ */
+
+// ============================================================================
+// Types
+// ============================================================================
+
+export interface MemorySnapshot {
+  /** JS heap used (Chrome only) */
+  jsHeapUsed?: number;
+  /** JS heap total (Chrome only) */
+  jsHeapTotal?: number;
+  /** JS heap limit (Chrome only) */
+  jsHeapLimit?: number;
+  /** GPU buffer pool stats */
+  gpu?: {
+    currentBytes: number;
+    activeBuffers: number;
+    pooledBuffers: number;
+    peakBytes: number;
+  };
+  /** Shard cache stats */
+  shardCache?: {
+    totalBytes: number;
+    shardCount: number;
+  };
+  /** Loaded model state */
+  modelState?: {
+    layerCount: number;
+    gpuBufferCount: number;
+  };
+}
+
+export interface MemoryMonitorState {
+  startTime: number;
+  interval: ReturnType<typeof setInterval> | null;
+}
+
+// ============================================================================
+// Memory Snapshot
+// ============================================================================
+
+/**
+ * Capture current memory statistics.
+ *
+ * @returns Memory snapshot with available stats
+ */
+export declare function captureMemorySnapshot(): MemorySnapshot;
+
+/**
+ * Format memory snapshot for logging.
+ *
+ * @param phase - Loading phase label
+ * @param elapsed - Elapsed seconds since start
+ * @param snapshot - Memory snapshot
+ * @param shardCacheBytes - Shard cache total bytes
+ * @param shardCount - Number of cached shards
+ * @param layerCount - Number of loaded layers
+ * @param gpuBufferCount - Number of GPU buffers
+ * @returns Formatted log string
+ */
+export declare function formatMemoryStats(
+  phase: string,
+  elapsed: number,
+  snapshot: MemorySnapshot,
+  shardCacheBytes: number,
+  shardCount: number,
+  layerCount: number,
+  gpuBufferCount: number
+): string;
+
+// ============================================================================
+// Memory Monitor Class
+// ============================================================================
+
+/**
+ * Memory monitor for tracking loading progress.
+ *
+ * Manages periodic memory logging during model loading.
+ */
+export declare class MemoryMonitor {
+  constructor(logIntervalMs?: number);
+
+  /**
+   * Start memory monitoring.
+   *
+   * @param getState - Function to get current loader state for logging
+   */
+  start(getState: () => { shardCacheBytes: number; shardCount: number; layerCount: number; gpuBufferCount: number }): void;
+
+  /**
+   * Stop memory monitoring.
+   *
+   * @param phase - Final phase label ('complete' or 'failed')
+   * @param getState - Function to get current loader state
+   */
+  stop(
+    phase: 'complete' | 'failed',
+    getState: () => { shardCacheBytes: number; shardCount: number; layerCount: number; gpuBufferCount: number }
+  ): void;
+
+  /**
+   * Get elapsed time since monitoring started.
+   */
+  getElapsed(): number;
+}