adaptive-memory-multi-model-router 2.2.5 → 2.2.6

package/src/inference/speculativeDecoding.ts

@@ -0,0 +1,373 @@
+/**
+ * A3M Router - Speculative Decoding Interface
+ *
+ * Interface for integrating speculative decoding.
+ * Currently a stub/interface for future Medusa/Lookahead integration.
+ *
+ * Speculative decoding uses a smaller "draft" model to predict
+ * multiple tokens ahead, which are then verified in parallel by
+ * the main model. This can provide 2-3x speedup in generation.
+ *
+ * Usage:
+ *   const specDec = new SpeculativeDecoding();
+ *   if (specDec.shouldUse(true)) {
+ *     const draftModel = specDec.getDraftModel('medusa');
+ *     // ... use draft model to generate and verify
+ *   }
+ */
+
+import { ProviderTier } from '../providers/providerConfig';
+
+// ============================================================
+// Types
+// ============================================================
+
+export type DraftModelType = 'medusa' | 'lookahead' | 'eagle' | 'spark';
+
+export interface SpeculativeConfig {
+  /** Enable speculative decoding */
+  enabled: boolean;
+  /** Draft model type */
+  draftModelType?: DraftModelType;
+  /** Number of speculative tokens to generate */
+  speculationWindow?: number;
+  /** Temperature for draft model */
+  temperature?: number;
+  /** Provider to use for draft model */
+  provider?: string;
+  /** Model to use for draft model */
+  model?: string;
+}
+
+export interface DraftModel {
+  /** Model identifier */
+  id: string;
+  /** Model type */
+  type: DraftModelType;
+  /** Provider for this model */
+  provider: string;
+  /** Model size description */
+  size: string;
+  /** Supported speculation windows */
+  supportedWindows: number[];
+}
+
+export interface SpeculativeResult {
+  /** Whether speculative decoding was used */
+  used: boolean;
+  /** Number of tokens in draft */
+  draftTokens: number;
+  /** Number of tokens accepted */
+  acceptedTokens: number;
+  /** Acceptance rate */
+  acceptanceRate: number;
+  /** Time saved (ms, estimated) */
+  estimatedTimeSaved: number;
+}
+
+export interface VerificationResult {
+  /** All tokens verified successfully */
+  allAccepted: boolean;
+  /** Indices of accepted tokens */
+  acceptedIndices: number[];
+  /** Indices of rejected tokens */
+  rejectedIndices: number[];
+  /** Actual tokens to use (may differ from draft) */
+  actualTokens: string[];
+  /** Number of tokens to rewind */
+  rewindCount: number;
+}
+
+// ============================================================
+// Speculative Decoding Interface
+// ============================================================
+
+export class SpeculativeDecoding {
+  private config: SpeculativeConfig;
+  private draftModels: Map<string, DraftModel> = new Map();
+
+  constructor(config?: Partial<SpeculativeConfig>) {
+    this.config = {
+      enabled: config?.enabled ?? false,
+      draftModelType: config?.draftModelType ?? 'medusa',
+      speculationWindow: config?.speculationWindow ?? 4,
+      temperature: config?.temperature ?? 0.0,
+      provider: config?.provider,
+      model: config?.model,
+    };
+
+    // Register available draft models
+    this.registerDraftModels();
+  }
+
+  /**
+   * Determine if speculative decoding should be used.
+   * Checks config, provider support, and model availability.
+   */
+  shouldUse(forceEnable?: boolean): boolean {
+    if (forceEnable) return true;
+    if (!this.config.enabled) return false;
+
+    // Speculative decoding beneficial for longer outputs
+    // Check if provider supports it
+    return this.config.enabled;
+  }
+
+  /**
+   * Get the draft model configuration.
+   */
+  getDraftModel(type?: DraftModelType): DraftModel | null {
+    const modelType = type || this.config.draftModelType || 'medusa';
+
+    const model = this.draftModels.get(modelType);
+    if (!model) {
+      // Return a generic configuration
+      return {
+        id: `draft-${modelType}`,
+        type: modelType,
+        provider: this.config.provider || 'auto',
+        size: 'small',
+        supportedWindows: [2, 4, 6, 8],
+      };
+    }
+
+    return model;
+  }
+
+  /**
+   * Generate draft tokens using the draft model.
+   * This is an async method to support API-based draft models.
+   */
+  async generateDraft(
+    prompt: string,
+    maxTokens: number
+  ): Promise<{ tokens: string[]; scores: number[] }> {
+    const draftModel = this.getDraftModel();
+    if (!draftModel) {
+      throw new Error('No draft model available');
+    }
+
+    // Stub: In production, this would call the draft model API
+    // For now, return empty draft
+    return {
+      tokens: [],
+      scores: [],
+    };
+  }
+
+  /**
+   * Verify draft tokens against the main model.
+   * Returns which tokens were accepted and corrections if needed.
+   */
+  async verifyDraft(
+    prompt: string,
+    draftTokens: string[]
+  ): Promise<VerificationResult> {
+    if (draftTokens.length === 0) {
+      return {
+        allAccepted: true,
+        acceptedIndices: [],
+        rejectedIndices: [],
+        actualTokens: [],
+        rewindCount: 0,
+      };
+    }
+
+    // Stub: In production, this would verify each draft token
+    // For now, accept all tokens (optimistic)
+    return {
+      allAccepted: true,
+      acceptedIndices: draftTokens.map((_, i) => i),
+      rejectedIndices: [],
+      actualTokens: draftTokens,
+      rewindCount: 0,
+    };
+  }
+
+  /**
+   * Calculate speedup from speculative decoding results.
+   */
+  calculateSpeedup(result: SpeculativeResult): number {
+    if (!result.used || result.draftTokens === 0) return 1.0;
+
+    // Approximate speedup based on acceptance rate
+    // Higher acceptance = more speedup
+    const acceptanceRate = result.acceptanceRate;
+
+    // Theoretical 2-3x speedup at 90%+ acceptance
+    if (acceptanceRate >= 0.9) return 2.5;
+    if (acceptanceRate >= 0.8) return 2.0;
+    if (acceptanceRate >= 0.7) return 1.7;
+    if (acceptanceRate >= 0.5) return 1.4;
+    if (acceptanceRate >= 0.3) return 1.2;
+    return 1.1;
+  }
+
+  /**
+   * Update configuration.
+   */
+  updateConfig(updates: Partial<SpeculativeConfig>): void {
+    this.config = { ...this.config, ...updates };
+  }
+
+  /**
+   * Get current configuration.
+   */
+  getConfig(): SpeculativeConfig {
+    return { ...this.config };
+  }
+
+  /**
+   * Check if a provider supports speculative decoding.
+   */
+  supportsSpeculativeDecoding(provider: string): boolean {
+    // Local providers (Ollama, vLLM) typically support it
+    const localProviders = ['ollama', 'lmstudio', 'vllm'];
+    if (localProviders.includes(provider)) return true;
+
+    // Check if provider has known speculative support
+    const supportedProviders = ['together', 'fireworks', 'anyscale'];
+    return supportedProviders.includes(provider);
+  }
+
+  /**
+   * Get recommended speculation window based on model size.
+   */
+  getRecommendedWindow(modelSize: 'small' | 'medium' | 'large'): number {
+    switch (modelSize) {
+      case 'small': return 6;
+      case 'medium': return 4;
+      case 'large': return 2;
+      default: return 4;
+    }
+  }
+
+  // ---- Private helpers ----
+
+  private registerDraftModels(): void {
+    // Medusa-style models (multiple draft heads)
+    this.draftModels.set('medusa', {
+      id: 'medusa-7b',
+      type: 'medusa',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8, 10],
+    });
+
+    // Lookahead decoding
+    this.draftModels.set('lookahead', {
+      id: 'lookahead-7b',
+      type: 'lookahead',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6],
+    });
+
+    // EAGLE decoding
+    this.draftModels.set('eagle', {
+      id: 'eagle-7b',
+      type: 'eagle',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8],
+    });
+
+    // Spark (speculative decoding for transformers)
+    this.draftModels.set('spark', {
+      id: 'spark-7b',
+      type: 'spark',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8, 10, 12],
+    });
+  }
+}
+
+// ============================================================
+// Speculative Decoding Wrapper for LLM calls
+// ============================================================
+
+export class SpeculativeDecodingWrapper {
+  private specDec: SpeculativeDecoding;
+  private mainModelCall: (prompt: string, options?: any) => Promise<string>;
+
+  constructor(
+    mainModelCall: (prompt: string, options?: any) => Promise<string>,
+    config?: Partial<SpeculativeConfig>
+  ) {
+    this.mainModelCall = mainModelCall;
+    this.specDec = new SpeculativeDecoding(config);
+  }
+
+  /**
+   * Generate with speculative decoding.
+   */
+  async generate(
+    prompt: string,
+    maxTokens: number,
+    options?: { temperature?: number; useSpecDec?: boolean }
+  ): Promise<{ text: string; result: SpeculativeResult }> {
+    if (!this.specDec.shouldUse(options?.useSpecDec)) {
+      // Fallback to normal generation
+      const text = await this.mainModelCall(prompt, options);
+      return {
+        text,
+        result: {
+          used: false,
+          draftTokens: 0,
+          acceptedTokens: 0,
+          acceptanceRate: 0,
+          estimatedTimeSaved: 0,
+        },
+      };
+    }
+
+    // Speculative generation
+    const startTime = Date.now();
+    const draftTokens = await this.specDec.generateDraft(prompt, maxTokens);
+    const verification = await this.specDec.verifyDraft(prompt, draftTokens.tokens);
+
+    let text: string;
+
+    if (verification.allAccepted) {
+      // All draft tokens accepted, continue with main model
+      text = verification.actualTokens.join('') +
+        await this.mainModelCall(prompt, options);
+    } else {
+      // Some tokens rejected, need to rewind
+      text = verification.actualTokens.join('');
+      // Re-prompt with corrected context
+      text += await this.mainModelCall(prompt + text, options);
+    }
+
+    const endTime = Date.now();
+    const speedup = this.specDec.calculateSpeedup({
+      used: true,
+      draftTokens: draftTokens.tokens.length,
+      acceptedTokens: verification.acceptedIndices.length,
+      acceptanceRate: verification.acceptedIndices.length / draftTokens.tokens.length,
+      estimatedTimeSaved: endTime - startTime,
+    });
+
+    return {
+      text,
+      result: {
+        used: true,
+        draftTokens: draftTokens.tokens.length,
+        acceptedTokens: verification.acceptedIndices.length,
+        acceptanceRate: verification.acceptedIndices.length / draftTokens.tokens.length,
+        estimatedTimeSaved: Math.round((1 - 1 / speedup) * (endTime - startTime)),
+      },
+    };
+  }
+}
+
+// ============================================================
+// Factory
+// ============================================================
+
+export function createSpeculativeDecoding(
+  config?: Partial<SpeculativeConfig>
+): SpeculativeDecoding {
+  return new SpeculativeDecoding(config);
+}