edgeinfer 0.0.0-stage → 2.0.0

package/src/runtime-manager.ts

@@ -0,0 +1,338 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2026 Soumya Debnath
+
+import * as ort from 'onnxruntime-web';
+
+export interface VRAMProbeResult {
+  available: boolean;
+  /** Largest single storage buffer the device actually allocated. */
+  maxAllocationBytes: number;
+  testedAllocations: number[];
+  failedAllocationBytes: number | null;
+  adapterLimits: {
+    maxBufferSize: number;
+    maxStorageBufferBindingSize: number;
+  };
+  deviceCreated: boolean;
+  probe: 'empirical' | 'limits' | 'unavailable';
+}
+
+export interface VRAMProbeOptions {
+  startBytes?: number;
+  maxBytes?: number;
+  stepBytes?: number;
+}
+
+export interface GPUCapabilities {
+  provider: string;
+  hasWebGPU: boolean;
+  hasWebNN: boolean;
+  /**
+   * NOT an amount of video memory. This is the largest single GPU storage
+   * buffer allocated by the empirical probe, bounded by the adapter's
+   * `maxBufferSize`. Browser baselines are commonly 256 MiB regardless of the
+   * physical card, so treat this as a maximum-tensor-size budget, not as VRAM.
+   */
+  estimatedVRAM: number;
+  /** Raw adapter limits, kept separate from the empirical probe. */
+  adapterLimits: {
+    maxBufferSize: number;
+    maxStorageBufferBindingSize: number;
+  } | null;
+  /** Evidence for how `estimatedVRAM` was obtained. */
+  vramProbe: VRAMProbeResult | null;
+  recommendedQuantization: 'fp32' | 'fp16' | 'int8' | 'int4';
+}
+
+/** Max time to wait for navigator.gpu.requestAdapter() before giving up. */
+const ADAPTER_TIMEOUT_MS = 5000;
+
+/**
+ * RuntimeManager handles execution provider detection and ONNX session creation.
+ * Implements a smart fallback chain: WebGPU → WebNN → WASM
+ * with automatic quantization level recommendation based on available GPU memory.
+ */
+export class RuntimeManager {
+  private static cachedCapabilities: GPUCapabilities | null = null;
+
+  /**
+   * Detect GPU capabilities and recommend the best execution configuration.
+   */
+  static async detectCapabilities(): Promise<GPUCapabilities> {
+    if (this.cachedCapabilities) return this.cachedCapabilities;
+
+    const capabilities: GPUCapabilities = {
+      provider: 'wasm',
+      hasWebGPU: false,
+      hasWebNN: false,
+      estimatedVRAM: 0,
+      adapterLimits: null,
+      vramProbe: null,
+      recommendedQuantization: 'int8',
+    };
+
+    // Check WebGPU availability.
+    // NOTE: `'gpu' in navigator` is NOT sufficient to conclude WebGPU works.
+    // A browser can expose navigator.gpu and still return null from
+    // requestAdapter() (no GPU, driver blocklisted, headless). We therefore
+    // always request an adapter, and we bound the wait: requestAdapter() is
+    // specified to settle, but a hung driver must not stall the whole SDK.
+    if (typeof navigator !== 'undefined' && 'gpu' in navigator) {
+      try {
+        const gpu = (navigator as any).gpu;
+        const adapter = await RuntimeManager.withTimeout(
+          Promise.resolve(gpu.requestAdapter()),
+          ADAPTER_TIMEOUT_MS
+        );
+        if (adapter) {
+          capabilities.hasWebGPU = true;
+          capabilities.provider = 'webgpu';
+
+          const maxBufferSize = adapter.limits?.maxBufferSize || 0;
+          const maxStorageSize = adapter.limits?.maxStorageBufferBindingSize || 0;
+          capabilities.adapterLimits = {
+            maxBufferSize,
+            maxStorageBufferBindingSize: maxStorageSize,
+          };
+
+          // Probe what the device actually allocates instead of treating the
+          // adapter's baseline limits as physical VRAM.
+          const probe = await RuntimeManager.probeVRAM(adapter);
+          capabilities.vramProbe = probe;
+          capabilities.estimatedVRAM = probe.available
+            ? probe.maxAllocationBytes
+            : Math.max(maxBufferSize, maxStorageSize);
+
+          const vramGB = capabilities.estimatedVRAM / (1024 * 1024 * 1024);
+          if (vramGB >= 4) {
+            capabilities.recommendedQuantization = 'fp16';
+          } else if (vramGB >= 1) {
+            capabilities.recommendedQuantization = 'int8';
+          } else {
+            capabilities.recommendedQuantization = 'int4';
+          }
+        }
+      } catch {
+        // WebGPU not available or adapter request failed
+      }
+    }
+
+    // Check WebNN availability
+    if (typeof navigator !== 'undefined' && 'ml' in navigator) {
+      try {
+        const ml = (navigator as any).ml;
+        if (ml) {
+          capabilities.hasWebNN = true;
+          if (!capabilities.hasWebGPU) {
+            capabilities.provider = 'webnn';
+          }
+        }
+      } catch {
+        // WebNN not available
+      }
+    }
+
+    this.cachedCapabilities = capabilities;
+    return capabilities;
+  }
+
+  /**
+   * Get the best available execution provider with smart fallback.
+   * Priority: WebGPU → WebNN → WASM
+   */
+  static async getBestExecutionProvider(): Promise<string> {
+    const caps = await this.detectCapabilities();
+    return caps.provider;
+  }
+
+  /**
+   * Create an ONNX inference session with automatic provider selection.
+   * Falls back gracefully if the preferred provider fails.
+   *
+   * Returns BOTH the session and the provider that actually served it. The
+   * detected capability (`detectCapabilities().provider`) is only a preference:
+   * a provider can be detected as available and still fail to create a session
+   * (for example onnxruntime-web only registers the WebGPU execution provider
+   * in some builds/versions). Callers must report the returned `provider`, not
+   * the detected one, or `EdgeInfer.executionProvider` will lie.
+   */
+  static async createSession(
+    modelBuffer: ArrayBuffer,
+    providers?: string[]
+  ): Promise<{ session: ort.InferenceSession; provider: string }> {
+    const requestedProviders = providers || [await this.getBestExecutionProvider()];
+
+    // Try each provider in order, falling back on failure
+    for (const provider of requestedProviders) {
+      try {
+        const session = await ort.InferenceSession.create(
+          modelBuffer,
+          {
+            executionProviders: [provider],
+            graphOptimizationLevel: 'all',
+          }
+        );
+        return { session, provider };
+      } catch (err) {
+        console.warn(`EdgeInfer: Failed to create session with provider "${provider}", trying next...`, err);
+      }
+    }
+
+    // Final fallback: WASM (should always work)
+    console.warn('EdgeInfer: All preferred providers failed, falling back to WASM');
+    const session = await ort.InferenceSession.create(modelBuffer, {
+      executionProviders: ['wasm'],
+      graphOptimizationLevel: 'all',
+    });
+    return { session, provider: 'wasm' };
+  }
+
+  /**
+   * Empirically probe the largest single storage buffer the adapter's device
+   * will allocate. Allocates and destroys one buffer at a time; never relies on
+   * `maxBufferSize` as if it were physical VRAM.
+   */
+  static async probeVRAM(
+    adapter: any,
+    options: VRAMProbeOptions = {}
+  ): Promise<VRAMProbeResult> {
+    const limits = {
+      maxBufferSize: Number(adapter?.limits?.maxBufferSize) || 0,
+      maxStorageBufferBindingSize:
+        Number(adapter?.limits?.maxStorageBufferBindingSize) || 0,
+    };
+    const unavailable: VRAMProbeResult = {
+      available: false,
+      maxAllocationBytes: 0,
+      testedAllocations: [],
+      failedAllocationBytes: null,
+      adapterLimits: limits,
+      deviceCreated: false,
+      probe: 'unavailable',
+    };
+
+    if (!adapter || typeof adapter.requestDevice !== 'function') return unavailable;
+
+    let device: any;
+    try {
+      device = await adapter.requestDevice();
+    } catch {
+      return unavailable;
+    }
+    if (!device || typeof device.createBuffer !== 'function') {
+      return { ...unavailable, deviceCreated: Boolean(device) };
+    }
+
+    const MIB = 1024 * 1024;
+    const start = Math.max(1, options.startBytes ?? 64 * MIB);
+    const configuredMax = Math.max(start, options.maxBytes ?? 4 * 1024 * MIB);
+    const max = limits.maxBufferSize > 0
+      ? Math.min(configuredMax, limits.maxBufferSize)
+      : configuredMax;
+    const step = Math.max(1, options.stepBytes ?? 32 * MIB);
+    const tested: number[] = [];
+    const STORAGE = 0x80;
+    const COPY_DST = 0x8;
+
+    const attempt = (size: number): boolean => {
+      tested.push(size);
+      try {
+        const buffer = device.createBuffer({ size, usage: STORAGE | COPY_DST });
+        buffer?.destroy?.();
+        return true;
+      } catch {
+        return false;
+      }
+    };
+
+    let low = 0;
+    let high = 0;
+    let size = Math.min(start, max);
+    while (size >= 1) {
+      if (attempt(size)) {
+        low = size;
+        break;
+      }
+      high = size;
+      if (size === 1) break;
+      size = Math.max(1, Math.floor(size / 2));
+    }
+
+    if (low === 0) {
+      device.destroy?.();
+      return {
+        ...unavailable,
+        testedAllocations: tested,
+        failedAllocationBytes: high || null,
+        deviceCreated: true,
+        probe: 'limits',
+      };
+    }
+
+    let next = Math.min(max, low * 2);
+    while (next > low && tested.length < 16) {
+      if (attempt(next)) {
+        low = next;
+        next = Math.min(max, low * 2);
+        if (low === max) break;
+      } else {
+        high = next;
+        break;
+      }
+    }
+    if (high === 0) high = Math.min(max, low * 2);
+
+    while (high - low > step && tested.length < 24) {
+      const mid = Math.floor((low + high) / 2);
+      if (mid <= low || mid >= high) break;
+      if (attempt(mid)) low = mid;
+      else high = mid;
+    }
+
+    device.destroy?.();
+    return {
+      available: true,
+      maxAllocationBytes: low,
+      testedAllocations: tested,
+      failedAllocationBytes: high > low ? high : null,
+      adapterLimits: limits,
+      deviceCreated: true,
+      probe: 'empirical',
+    };
+  }
+
+  /**
+   * Resolve `promise`, or resolve to `null` if it has not settled within `ms`.
+   * Used to stop a hung GPU driver from blocking model loading forever.
+   */
+  private static withTimeout<T>(promise: Promise<T>, ms: number): Promise<T | null> {
+    let timer: any;
+    return Promise.race([
+      promise.then((v) => { clearTimeout(timer); return v; }),
+      new Promise<null>((resolve) => {
+        timer = setTimeout(() => {
+          console.warn(`EdgeInfer: navigator.gpu.requestAdapter() did not settle within ${ms}ms; falling back to WASM.`);
+          resolve(null);
+        }, ms);
+        if (typeof timer?.unref === 'function') timer.unref();
+      }),
+    ]);
+  }
+
+  /**
+   * Estimate if a model of given size (bytes) can fit in available VRAM.
+   */
+  static async canFitModel(modelSizeBytes: number): Promise<boolean> {
+    const caps = await this.detectCapabilities();
+    if (!caps.hasWebGPU) return true; // WASM uses system RAM, not VRAM
+    // Model typically needs 1.5x its size in VRAM (weights + activations)
+    return caps.estimatedVRAM >= modelSizeBytes * 1.5;
+  }
+
+  /**
+   * Reset cached capabilities (useful for testing or re-detection).
+   */
+  static resetCache(): void {
+    this.cachedCapabilities = null;
+  }
+}

package/src/speculative.ts

@@ -0,0 +1,294 @@
+// SPDX-License-Identifier: MIT
+// Copyright (c) 2024-2026 Soumya Debnath
+//
+// Speculative decoding with a small draft model and a larger target model.
+// The target verifies gamma draft tokens in one forward pass. Accepted tokens
+// keep their draft distribution; a rejection resamples from the residual
+// max(0, p_target - p_draft) and truncates both KV-cache states.
+
+export interface ModelState {
+  tokens: number[];
+  /** Opaque KV-cache handle for cache-aware adapters. */
+  cache?: unknown;
+}
+
+export interface DecodeResult {
+  state: ModelState;
+  /**
+   * One next-token distribution for each input position, plus one distribution
+   * after the last input token. Length = tokens.length + 1. Calling decode()
+   * with an empty token array returns the distribution after the current state.
+   */
+  logits: Float32Array[];
+}
+
+export interface LanguageModel {
+  readonly vocabSize: number;
+  prefill(tokens: number[]): Promise<ModelState>;
+  decode(state: ModelState, tokens: number[]): Promise<DecodeResult>;
+  /**
+   * Drop cache entries after `length` total tokens. Cache-aware adapters must
+   * implement this; stateless models can just slice the token list.
+   */
+  truncate(state: ModelState, length: number): ModelState;
+}
+
+export interface SpeculativeOptions {
+  draft: LanguageModel;
+  target: LanguageModel;
+  gamma?: number;
+  rng?: () => number;
+  eosTokenId?: number;
+}
+
+export interface SpeculativeStats {
+  outputTokens: number;
+  draftTokens: number;
+  acceptedTokens: number;
+  bonusTokens: number;
+  rejectedTokens: number;
+  targetCalls: number;
+  draftCalls: number;
+  acceptanceRate: number;
+  targetCallsPerToken: number;
+  elapsedMs: number;
+  tokensPerSecond: number;
+}
+
+export interface SpeculativeGenerationResult {
+  tokens: number[];
+  stats: SpeculativeStats;
+}
+
+export class SpeculativeDecoder {
+  readonly draft: LanguageModel;
+  readonly target: LanguageModel;
+  readonly gamma: number;
+  private readonly rng: () => number;
+  private readonly eosTokenId?: number;
+
+  constructor(options: SpeculativeOptions) {
+    if (options.draft.vocabSize !== options.target.vocabSize) {
+      throw new Error("Draft and target models must share a vocabulary size");
+    }
+    const gamma = options.gamma ?? 5;
+    if (!Number.isInteger(gamma) || gamma < 1) {
+      throw new Error("gamma must be a positive integer");
+    }
+    this.draft = options.draft;
+    this.target = options.target;
+    this.gamma = gamma;
+    this.rng = options.rng ?? Math.random;
+    this.eosTokenId = options.eosTokenId;
+  }
+
+  async generate(
+    prompt: number[],
+    maxNewTokens = 64,
+    { temperature = 1 }: { temperature?: number } = {}
+  ): Promise<SpeculativeGenerationResult> {
+    if (prompt.length === 0) throw new Error("prompt must not be empty");
+    if (maxNewTokens < 1) throw new Error("maxNewTokens must be positive");
+    if (temperature <= 0) throw new Error("temperature must be positive");
+
+    const started = performance.now();
+    let draftState = await this.draft.prefill(prompt);
+    let targetState = await this.target.prefill(prompt);
+
+    const output: number[] = [];
+    let draftTokens = 0;
+    let acceptedTokens = 0;
+    let bonusTokens = 0;
+    let rejectedTokens = 0;
+    let targetCalls = 0;
+    let draftCalls = 0;
+
+    while (output.length < maxNewTokens) {
+      const gamma = Math.min(this.gamma, maxNewTokens - output.length);
+      const draft = await this.draftBatch(draftState, gamma, temperature, draftCalls);
+      draftState = draft.state;
+      draftCalls = draft.calls;
+      draftTokens += draft.tokens.length;
+
+      const prefixLength = targetState.tokens.length;
+      const verification = await this.target.decode(targetState, draft.tokens);
+      targetCalls++;
+      targetState = verification.state;
+
+      let rejected = false;
+      for (let i = 0; i < draft.tokens.length; i++) {
+        const token = draft.tokens[i];
+        const q = draft.probabilities[i][token];
+        const p = probability(verification.logits[i], token, temperature);
+        if (this.rng() < Math.min(1, p / q)) {
+          output.push(token);
+          acceptedTokens++;
+          if (token === this.eosTokenId || output.length >= maxNewTokens) break;
+          continue;
+        }
+
+        const residual = residualDistribution(
+          verification.logits[i],
+          draft.probabilities[i],
+          temperature
+        );
+        const replacement = sample(residual, this.rng());
+        output.push(replacement);
+        rejectedTokens++;
+        rejected = true;
+
+        const keep = prefixLength + i;
+        targetState = this.target.truncate(targetState, keep);
+        draftState = this.draft.truncate(draftState, keep);
+        const targetAdvance = await this.target.decode(targetState, [replacement]);
+        targetCalls++;
+        targetState = targetAdvance.state;
+        const draftAdvance = await this.draft.decode(draftState, [replacement]);
+        draftCalls++;
+        draftState = draftAdvance.state;
+        break;
+      }
+
+      if (rejected) continue;
+      if (output.length >= maxNewTokens || output[output.length - 1] === this.eosTokenId) break;
+
+      const bonusDistribution = softmax(
+        verification.logits[draft.tokens.length],
+        temperature
+      );
+      const bonus = sample(bonusDistribution, this.rng());
+      output.push(bonus);
+      bonusTokens++;
+      const targetAdvance = await this.target.decode(targetState, [bonus]);
+      targetCalls++;
+      targetState = targetAdvance.state;
+      const draftAdvance = await this.draft.decode(draftState, [bonus]);
+      draftCalls++;
+      draftState = draftAdvance.state;
+      if (bonus === this.eosTokenId) break;
+    }
+
+    const elapsedMs = performance.now() - started;
+    const stats: SpeculativeStats = {
+      outputTokens: output.length,
+      draftTokens,
+      acceptedTokens,
+      bonusTokens,
+      rejectedTokens,
+      targetCalls,
+      draftCalls,
+      acceptanceRate: draftTokens === 0 ? 0 : acceptedTokens / draftTokens,
+      targetCallsPerToken: output.length === 0 ? 0 : targetCalls / output.length,
+      elapsedMs,
+      tokensPerSecond: elapsedMs === 0 ? 0 : (output.length / elapsedMs) * 1000,
+    };
+    return { tokens: output, stats };
+  }
+
+  private async draftBatch(
+    state: ModelState,
+    gamma: number,
+    temperature: number,
+    calls: number
+  ): Promise<{ state: ModelState; tokens: number[]; probabilities: Float32Array[]; calls: number }> {
+    let current = state;
+    const tokens: number[] = [];
+    const probabilities: Float32Array[] = [];
+
+    const first = await this.draft.decode(current, []);
+    calls++;
+    let nextDistribution = first.logits[0];
+    for (let i = 0; i < gamma; i++) {
+      const distribution = softmax(nextDistribution, temperature);
+      const token = sample(distribution, this.rng());
+      tokens.push(token);
+      probabilities.push(distribution);
+      const advanced = await this.draft.decode(current, [token]);
+      calls++;
+      current = advanced.state;
+      nextDistribution = advanced.logits[advanced.logits.length - 1];
+    }
+    return { state: current, tokens, probabilities, calls };
+  }
+}
+
+export function softmax(logits: Float32Array, temperature = 1): Float32Array {
+  let max = -Infinity;
+  for (let i = 0; i < logits.length; i++) max = Math.max(max, logits[i]);
+  const output = new Float32Array(logits.length);
+  let sum = 0;
+  for (let i = 0; i < logits.length; i++) {
+    const value = Math.exp((logits[i] - max) / temperature);
+    output[i] = value;
+    sum += value;
+  }
+  if (sum === 0) return output;
+  for (let i = 0; i < output.length; i++) output[i] /= sum;
+  return output;
+}
+
+export function probability(logits: Float32Array, token: number, temperature = 1): number {
+  return softmax(logits, temperature)[token] ?? 0;
+}
+
+export function sample(distribution: Float32Array, random: number): number {
+  let cumulative = 0;
+  for (let i = 0; i < distribution.length; i++) {
+    cumulative += distribution[i];
+    if (random <= cumulative) return i;
+  }
+  return distribution.length - 1;
+}
+
+export function residualDistribution(
+  targetLogits: Float32Array,
+  draftProbabilities: Float32Array,
+  temperature = 1
+): Float32Array {
+  const p = softmax(targetLogits, temperature);
+  const q = draftProbabilities;
+  const residual = new Float32Array(p.length);
+  let sum = 0;
+  for (let i = 0; i < p.length; i++) {
+    const value = Math.max(0, p[i] - q[i]);
+    residual[i] = value;
+    sum += value;
+  }
+  if (sum === 0) return p;
+  for (let i = 0; i < residual.length; i++) residual[i] /= sum;
+  return residual;
+}
+
+export type LogitPolicy = (tokens: number[]) => Float32Array;
+
+/** Deterministic mock used by tests and the throughput benchmark. */
+export class MockLanguageModel implements LanguageModel {
+  constructor(
+    private readonly policy: LogitPolicy,
+    readonly vocabSize: number,
+    private readonly latencyMs = 0
+  ) {}
+
+  async prefill(tokens: number[]): Promise<ModelState> {
+    await delay(this.latencyMs);
+    return { tokens: tokens.slice() };
+  }
+
+  async decode(state: ModelState, tokens: number[]): Promise<DecodeResult> {
+    await delay(this.latencyMs);
+    const logits: Float32Array[] = [];
+    for (let i = 0; i <= tokens.length; i++) {
+      logits.push(this.policy([...state.tokens, ...tokens.slice(0, i)]));
+    }
+    return { state: { tokens: [...state.tokens, ...tokens] }, logits };
+  }
+
+  truncate(state: ModelState, length: number): ModelState {
+    return { ...state, tokens: state.tokens.slice(0, length), cache: undefined };
+  }
+}
+
+function delay(ms: number): Promise<void> {
+  if (ms <= 0) return Promise.resolve();
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}