adaptive-memory-multi-model-router 1.2.2 → 1.3.1

package/src/utils/speculativeDecoding.ts

@@ -1,344 +0,0 @@
-/**
- * TMLPD Speculative Decoding
- * 
- * Based on Medusa (arXiv:2401.10774) and EAGLE approaches
- * Small draft model proposes tokens, large model verifies in parallel
- * 2-3x faster generation with same quality
- */
-
-export interface SpeculativeConfig {
-  draft_model: string;
-  target_model: string;
-  num_draft_tokens: number;
-  temperature?: number;
-  max_verify_tokens?: number;
-}
-
-export interface SpeculativeResult {
-  accepted: number;
-  rejected: number;
-  draft_tokens: number;
-  speedup: number;
-  final_text: string;
-}
-
-export interface DraftCandidate {
-  token: string;
-  probability: number;
-  position: number;
-}
-
-/**
- * Medusa-style multi-token prediction heads
- * Instead of separate draft model, uses speculative sampling
- */
-export class MedusaPredictor {
-  private num_heads: number;
-  private temperature: number;
-  
-  constructor(options?: {
-    num_heads?: number;
-    temperature?: number;
-  }) {
-    this.num_heads = options?.num_heads || 5;
-    this.temperature = options?.temperature || 0.7;
-  }
-  
-  /**
-   * Generate k draft tokens from one forward pass
-   * In production, this uses actual Medusa prediction heads
-   */
-  async generateDraftTokens(
-    context: string,
-    last_token: string,
-    getLogits: (text: string) => Promise<Record<string, number>>
-  ): Promise<DraftCandidate[]> {
-    // Simulate getting logits for next token predictions
-    // In real Medusa, this comes from extra prediction heads
-    const prompt = context + last_token;
-    const logits = await getLogits(prompt);
-    
-    const candidates: DraftCandidate[] = [];
-    const sorted = Object.entries(logits)
-      .sort((a, b) => b[1] - a[1])
-      .slice(0, this.num_heads);
-    
-    for (let i = 0; i < sorted.length; i++) {
-      const [token, prob] = sorted[i];
-      // Apply temperature
-      const adjusted = Math.pow(prob, 1 / this.temperature);
-      candidates.push({
-        token,
-        probability: adjusted,
-        position: i + 1
-      });
-    }
-    
-    return candidates;
-  }
-  
-  /**
-   * Verify draft tokens against target model
-   * Returns which tokens were accepted
-   */
-  async verifyDraft(
-    context: string,
-    drafts: DraftCandidate[],
-    targetLogits: (text: string) => Promise<Record<string, number>>
-  ): Promise<{ accepted: number[]; rejected: number[] }> {
-    const accepted: number[] = [];
-    const rejected: number[] = [];
-    
-    let current_context = context;
-    
-    for (const draft of drafts) {
-      // Get target model's prediction for this position
-      const target_logits = await targetLogits(current_context);
-      const target_token = Object.entries(target_logits)
-        .sort((a, b) => b[1] - a[1])[0]?.[0];
-      
-      // Accept if matches or probability is high enough
-      if (draft.token === target_token || draft.probability > 0.3) {
-        accepted.push(draft.position);
-        current_context += draft.token;
-      } else {
-        rejected.push(draft.position);
-        break;  // Reject rest of draft
-      }
-    }
-    
-    return { accepted, rejected };
-  }
-}
-
-/**
- * EAGLE-style speculative decoding
- * Uses regression-based draft token prediction
- */
-export class EagleSpeculative {
-  private num_draft_tokens: number;
-  
-  constructor(num_draft_tokens: number = 4) {
-    this.num_draft_tokens = num_draft_tokens;
-  }
-  
-  /**
-   * Generate draft sequence
-   * In production, this uses EAGLE's auto-regressive draft model
-   */
-  async generateDraft(
-    context: string,
-    generateFn: (prompt: string) => Promise<string>
-  ): Promise<string[]> {
-    const drafts: string[] = [];
-    let current = context;
-    
-    for (let i = 0; i < this.num_draft_tokens; i++) {
-      // In EAGLE, draft is generated from a compressed hidden state
-      // Here we simulate with regular generation
-      const next = await generateFn(current);
-      drafts.push(next);
-      current += next;
-      
-      if (next.trim().length === 0) break;
-    }
-    
-    return drafts;
-  }
-  
-  /**
-   * Verify draft with tree-based attention
-   * Multiple drafts are verified simultaneously
-   */
-  async verifyDraftTree(
-    context: string,
-    drafts: string[],
-    targetGenerate: (prompt: string) => Promise<string>
-  ): Promise<{ accepted: number; text: string }> {
-    let current_context = context;
-    let accepted_count = 0;
-    
-    for (const draft of drafts) {
-      // Target model generates one token at this position
-      const target_token = await targetGenerate(current_context);
-      
-      // If draft matches target, accept
-      if (draft.startsWith(target_token) || draft === target_token) {
-        accepted_count++;
-        current_context += target_token;
-      } else {
-        // Rejected - use target token
-        current_context += target_token;
-        if (accepted_count > 0) break;
-      }
-    }
-    
-    return {
-      accepted: accepted_count,
-      text: current_context.slice(context.length)
-    };
-  }
-}
-
-/**
- * Simple speculative decoding wrapper
- * Works with any model pair that supports continued generation
- */
-export class SpeculativeDecoder {
-  private draft_threshold: number;
-  
-  constructor(draft_threshold: number = 0.5) {
-    this.draft_threshold = draft_threshold;
-  }
-  
-  /**
-   * Execute speculative decoding
-   * 
-   * @param prompt - Input prompt
-   * @param draftFn - Function to generate draft completion (fast model)
-   * @param targetFn - Function to generate target completion (slow model)
-   * @param max_draft_tokens - Maximum tokens to draft
-   */
-  async decode(
-    prompt: string,
-    draftFn: (prompt: string, max_tokens: number) => Promise<string>,
-    targetFn: (prompt: string, max_tokens: number) => Promise<string>,
-    max_draft_tokens: number = 5
-  ): Promise<SpeculativeResult> {
-    const start_time = Date.now();
-    
-    // Phase 1: Generate draft with fast model
-    const draft_start = Date.now();
-    const draft_text = await draftFn(prompt, max_draft_tokens * 2);
-    const draft_time = Date.now() - draft_start;
-    
-    // Phase 2: Verify with target model (single pass)
-    // Instead of verifying token-by-token, we use acceptance criteria
-    const target_start = Date.now();
-    const target_text = await targetFn(prompt, max_draft_tokens);
-    const target_time = Date.now() - target_start;
-    
-    // Calculate acceptance rate
-    let accepted = 0;
-    let rejected = 0;
-    
-    const draft_words = draft_text.split(/\s+/);
-    const target_words = target_text.split(/\s+/);
-    
-    for (let i = 0; i < Math.min(draft_words.length, target_words.length); i++) {
-      // Simple word-level acceptance
-      if (draft_words[i].toLowerCase() === target_words[i].toLowerCase()) {
-        accepted++;
-      } else {
-        rejected++;
-        break;  // Stop at first rejection
-      }
-    }
-    
-    // If draft was longer, those are rejected
-    rejected += Math.max(0, draft_words.length - target_words.length);
-    
-    // Speedup: time_target / (time_draft + time_verification)
-    const total_time = draft_time + target_time;
-    const speedup = total_time > 0 ? (target_time / total_time) : 1;
-    
-    // Use target text (higher quality) as final
-    const final_text = target_text;
-    
-    return {
-      accepted,
-      rejected,
-      draft_tokens: draft_words.length,
-      speedup: Math.min(speedup, 3.0),  // Cap at 3x
-      final_text
-    };
-  }
-  
-  /**
-   * Execute with streaming (faster perceived latency)
-   */
-  async decodeStreaming(
-    prompt: string,
-    draftFn: (prompt: string) => Promise<string>,
-    targetFn: (prompt: string) => Promise<string>,
-    onToken: (token: string, is_draft: boolean) => void,
-    max_draft_tokens: number = 5
-  ): Promise<{ accepted: number; final_text: string }> {
-    // Generate drafts first
-    const drafts = await draftFn(prompt + " ");
-    const draft_tokens = drafts.split(/\s+/);
-    
-    let accepted = 0;
-    let final_text = "";
-    
-    // Verify and stream tokens
-    for (const token of draft_tokens) {
-      if (accepted >= max_draft_tokens) break;
-      
-      // Emit draft token immediately (lower quality)
-      onToken(token, true);
-      final_text += token + " ";
-      
-      accepted++;
-    }
-    
-    return { accepted, final_text: final_text.trim() };
-  }
-}
-
-/**
- * Batch speculative decoding
- * Processes multiple prompts with speculative execution
- */
-export async function speculativeBatch(
-  prompts: string[],
-  draftFn: (prompt: string) => Promise<string>,
-  targetFn: (prompt: string) => Promise<string>,
-  options?: {
-    concurrency?: number;
-    max_draft_tokens?: number;
-  }
-): Promise<SpeculativeResult[]> {
-  const concurrency = options?.concurrency || 3;
-  const max_draft_tokens = options?.max_draft_tokens || 5;
-  
-  const decoder = new SpeculativeDecoder();
-  const results: SpeculativeResult[] = [];
-  
-  // Process in batches
-  for (let i = 0; i < prompts.length; i += concurrency) {
-    const batch = prompts.slice(i, i + concurrency);
-    const batch_results = await Promise.all(
-      batch.map(p => decoder.decode(p, draftFn, targetFn, max_draft_tokens))
-    );
-    results.push(...batch_results);
-  }
-  
-  return results;
-}
-
-/**
- * Estimate speedup potential for a given prompt
- */
-export function estimateSpeedupPotential(
-  prompt_length: number,
-  expected_completion_length: number,
-  draft_speed_ms: number,
-  target_speed_ms: number
-): number {
-  // If draft is much faster, potential is higher
-  const draft_vs_target = target_speed_ms / draft_speed_ms;
-  
-  // But speculative decoding has overhead
-  const overhead_factor = 1.2;  // 20% overhead
-  
-  return Math.min(draft_vs_target / overhead_factor, 3.0);
-}
-
-export default {
-  MedusaPredictor,
-  EagleSpeculative,
-  SpeculativeDecoder,
-  speculativeBatch,
-  estimateSpeedupPotential
-};