大模型企业私有化推理加速：推测解码（Speculative Decoding）与 Draft Model 协同架构实战

/**
 * 模块：企业私有化大模型推测解码（Speculative Decoding）模拟器
 * 核心功能：模拟草稿模型生成、主模型单步并行核验、拒绝采样（Rejection Sampling）概率对齐与加速比推演
 */

// 概率分布映射字典：Token ID -> 概率值
export type ProbabilityDistribution = Record<number, number>;

// 单步草稿推测结果
export interface DraftCandidate {
  tokenId: number;
  tokenText: string;
  draftProbability: number;
}

// 主模型并行核验的单步结果
export interface TargetStepVerification {
  position: number;
  distribution: ProbabilityDistribution;
  argmaxTokenId: number;
}

// 模拟验证结果报告
export interface SpeculativeDecodingRoundReport {
  roundIndex: number;
  draftTokens: DraftCandidate[];
  acceptedCount: number;
  acceptedTokens: Array<{ tokenId: number; tokenText: string }>;
  rejectionHappened: boolean;
  rejectedIndex: number | null;
  correctedToken: { tokenId: number; tokenText: string } | null;
  finalEmittedTokens: Array<{ tokenId: number; tokenText: string }>;
  effectiveSpeedup: number;
}

/**
 * 推测解码采样器与调度管理类
 */
export class SpeculativeDecodingEngine {
  private vocabMap: Record<number, string>;

  constructor(vocabMap: Record<number, string>) {
    this.vocabMap = vocabMap;
  }

  /**
   * 执行单个推测解码轮次的拒绝采样（Rejection Sampling）判定
   * @param draftTokens 草稿模型生成的 K 个候选 Token
   * @param targetDistributions 主模型单步并行前向给出的 K+1 个位置的概率分布
   * @param temperature 采样温度系数（0 为确定性贪心模式）
   * @param randomSeedFn 可选的随机数发生函数，用于单元测试确定性复现
   */
  public executeSpeculativeRound(
    draftTokens: DraftCandidate[],
    targetDistributions: TargetStepVerification[],
    temperature: number = 0.7,
    randomSeedFn: () => number = Math.random
  ): SpeculativeDecodingRoundReport {
    const k = draftTokens.length;
    if (targetDistributions.length < k + 1) {
      throw new Error(`主模型前向分布长度不足：需要 ${k + 1} 步，实际仅传入 ${targetDistributions.length} 步`);
    }

    const acceptedTokens: Array<{ tokenId: number; tokenText: string }> = [];
    let rejectionHappened = false;
    let rejectedIndex: number | null = null;
    let correctedToken: { tokenId: number; tokenText: string } | null = null;

    // 逐位置进行推测核验
    for (let i = 0; i < k; i++) {
      const draft = draftTokens[i];
      const targetVerify = targetDistributions[i];
      const pDist = targetVerify.distribution;

      const pVal = pDist[draft.tokenId] || 0.0;
      const qVal = draft.draftProbability;

      if (temperature === 0) {
        // 贪心确定性场景判定
        if (draft.tokenId === targetVerify.argmaxTokenId) {
          acceptedTokens.push({ tokenId: draft.tokenId, tokenText: draft.tokenText });
        } else {
          rejectionHappened = true;
          rejectedIndex = i;
          // 修正为大模型贪心最高概率的 Token
          correctedToken = {
            tokenId: targetVerify.argmaxTokenId,
            tokenText: this.vocabMap[targetVerify.argmaxTokenId] || `[Token_${targetVerify.argmaxTokenId}]`
          };
          break;
        }
      } else {
        // 随机采样模式下的拒绝采样（Rejection Sampling）
        const acceptanceRatio = qVal > 0 ? Math.min(1.0, pVal / qVal) : 0.0;
        const diceRoll = randomSeedFn();

        if (diceRoll < acceptanceRatio) {
          // 概率接受
          acceptedTokens.push({ tokenId: draft.tokenId, tokenText: draft.tokenText });
        } else {
          // 触发拒绝，从残差分布 p'(x) = max(0, p(x) - q(x)) 中重新采样
          rejectionHappened = true;
          rejectedIndex = i;
          correctedToken = this.sampleFromResidualDistribution(pDist, draftTokens[i], randomSeedFn);
          break;
        }
      }
    }

    // 构建最终输出序列
    const finalEmittedTokens: Array<{ tokenId: number; tokenText: string }> = [...acceptedTokens];

    if (rejectionHappened && correctedToken) {
      // 若发生了拒绝，追加修正后的 Token，丢弃其后所有候选
      finalEmittedTokens.push(correctedToken);
    } else {
      // 若草稿 K 个全被接受，免费获得大模型第 K+1 步的前向预测 Token（无需额外前向）
      const bonusVerify = targetDistributions[k];
      const bonusTokenId = bonusVerify.argmaxTokenId;
      finalEmittedTokens.push({
        tokenId: bonusTokenId,
        tokenText: this.vocabMap[bonusTokenId] || `[Token_${bonusTokenId}]`
      });
    }

    // 计算本轮等效单次前向输出比（加速比衡量指标）
    const effectiveSpeedup = finalEmittedTokens.length;

    return {
      roundIndex: 1,
      draftTokens,
      acceptedCount: acceptedTokens.length,
      acceptedTokens,
      rejectionHappened,
      rejectedIndex,
      correctedToken,
      finalEmittedTokens,
      effectiveSpeedup
    };
  }

  /**
   * 从残差归一化分布中执行安全采样
   */
  private sampleFromResidualDistribution(
    pDist: ProbabilityDistribution,
    draftCandidate: DraftCandidate,
    randomSeedFn: () => number
  ): { tokenId: number; tokenText: string } {
    const residual: Record<number, number> = {};
    let totalResidual = 0;

    for (const [idStr, p] of Object.entries(pDist)) {
      const id = Number(idStr);
      const q = id === draftCandidate.tokenId ? draftCandidate.draftProbability : 0;
      const diff = Math.max(0, p - q);
      residual[id] = diff;
      totalResidual += diff;
    }

    // 极端边界防空降级
    if (totalResidual <= 0) {
      return {
        tokenId: draftCandidate.tokenId,
        tokenText: draftCandidate.tokenText
      };
    }

    // 累积概率切片采样
    const randThreshold = randomSeedFn() * totalResidual;
    let accumulated = 0;
    for (const [idStr, diff] of Object.entries(residual)) {
      accumulated += diff;
      if (accumulated >= randThreshold) {
        const selectedId = Number(idStr);
        return {
          tokenId: selectedId,
          tokenText: this.vocabMap[selectedId] || `[Token_${selectedId}]`
        };
      }
    }

    // 兜底返回主键
    const fallbackId = Number(Object.keys(pDist)[0]);
    return {
      tokenId: fallbackId,
      tokenText: this.vocabMap[fallbackId] || `[Token_${fallbackId}]`
    };
  }
}