adaptive-memory-multi-model-router 2.2.9 → 2.4.0

package/src/inference/speculativeDecoding.ts

@@ -1,373 +0,0 @@
-/**
- * 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);
-}