Skip to main content

投机解码

Speculative Engine 模块架构分析​

1. 整体架构​

┌─────────────────────────────────────────────────────────────────┐
│                     SpeculativeEngine                           │
│  (继承自 EngineBase，是推测解码引擎的主入口)                      │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│   ┌──────────────────┐  ┌──────────────────┐  ┌────────────────┐│
│   │  ProposeExecutor │  │  ScoreExecutor   │  │SpeculativeSampler│
│   │  (草稿模型执行器)  │  │  (验证模型执行器) │  │  (采样器)       ││
│   └────────┬─────────┘  └────────┬─────────┘  └────────────────┘│
│            │                     │                              │
│   ┌────────┴─────────────────────┴───────────┐                  │
│   │              具体实现类                    │                  │
│   │  ┌─────────────┐  ┌─────────────┐        │                  │
│   │  │VanillaExec  │  │ MTPExecutor │        │                  │
│   │  │(标准草稿模型)│  │(多token预测) │        │                  │
│   │  └─────────────┘  └──────┬──────┘        │                  │
│   │                          │               │                  │
│   │                   ┌──────┴──────┐        │                  │
│   │                   │EagleExecutor│        │                  │
│   │                   │(EAGLE算法)  │        │                  │
│   │                   └─────────────┘        │                  │
│   └──────────────────────────────────────────┘                  │
│                                                                 │
│   ┌──────────────────────────────────────────┐                  │
│   │         SpeculativeScheduler             │                  │
│   │         (推测解码调度器)                   │                  │
│   └──────────────────────────────────────────┘                  │
└─────────────────────────────────────────────────────────────────┘

2. 核心组件详解​

2.1 SpeculativeEngine (主引擎)​

职责: 协调整个推测解码流程

关键方法:

  • step() - 主循环的一次迭代，根据不同模式调用不同的处理流程
  • spStep() - 标准推测解码步骤 (Vanilla 模式)
  • mtpStep() - MTP/EAGLE 模式的推测解码步骤
  • normStep() - 普通解码步骤 (不使用推测)

支持的推测解码类型 (sp_type_):

类型说明
vanilla传统小模型草稿
mtpMulti-Token Prediction
eagleEAGLE 算法
eagle3EAGLE3 算法

2.2 ProposeExecutor (提议执行器)​

抽象接口:

class ProposeExecutor {
    virtual absl::Status propose(streams, skip_check) = 0;  // 生成草稿 tokens
    virtual size_t reserveStep() const = 0;                  // 预留步数
    virtual absl::Status normalProcess(streams) = 0;         // 普通处理
};

三种实现:

实现类原理特点
VanillaExecutor使用独立的小型草稿模型传统方法，需要额外模型
MTPExecutor多 token 预测头，共享主干每个 MTP 模块预测下一个位置的 token
EagleExecutor继承 MTPExecutor，使用 EAGLE 树结构更高效的投机采样

2.3 ScoreExecutor (评分执行器)​

职责: 使用目标模型验证草稿 tokens

class ScoreExecutor {
    absl::Status score(streams, skip_check);      // 验证提议的 tokens
    absl::Status normalProcess(streams);          // 普通前向传播
};
  • 内部包含两个 NormalExecutor：一个用于 score，一个用于普通处理
  • 使用 ScoreBatchStreamProcessor 进行批处理

2.4 SpeculativeSampler (推测采样器)​

职责: 实现接受/拒绝采样逻辑

两种采样模式:

模式触发条件特点
streamSampletop-k=1 或强制流式采样逐流处理，简单高效
batchSample包含 top-p 采样批量处理，支持随机采样

采样算法:

  • top1Sample: 贪婪采样，直接比较 token ID
  • stochasticSample: 随机采样，基于概率拒绝
// 拒绝采样核心逻辑
while (accept_len < propose_step) {
    if (random > score_prob / propose_prob) {
        // 拒绝，使用修正后的概率重新采样
        new_p = max(score_probs - propose_probs, 0)
        resample from normalized new_p
        break;
    }
    accept_token();
}

2.5 SpeculativeScheduler (调度器)​

职责: 管理请求调度，区分推测和普通解码

class SpeculativeScheduler : public FIFOScheduler {
    std::list<GenerateStreamPtr> pending_sp_run_streams_;  // 等待推测解码的流
    
    schedule(reserve_step) {
        // 1. 优先返回待处理的推测流
        // 2. 分离禁用推测的流和启用推测的流
    }
};

3. 推测解码流程​

3.1 标准推测解码流程 (spStep)​

┌─────────────┐     ┌─────────────┐     ┌─────────────┐
│   Propose   │────▶│    Score    │────▶│   Sample    │
│ (草稿模型)   │     │  (验证模型)  │     │  (接受/拒绝) │
└─────────────┘     └─────────────┘     └─────────────┘
      │                    │                   │
      │ 生成 k 个          │ 并行验证          │ 接受 n 个
      │ 候选 tokens        │ k+1 个位置        │ tokens (n≤k+1)
      ▼                    ▼                   ▼

3.2 MTP/EAGLE 流程 (mtpStep)​

分类 streams
      │
      ├─── propose_streams (有 hidden states) ──▶ propose()
      │
      ├─── prefill_streams (无 hidden states) ──▶ 创建空 propose stream
      │
      └─── pre_propose_streams (已有 propose token)
                    │
                    ▼
            score() 验证
                    │
                    ▼
            sample() 采样

4. 性能指标​

模块收集以下指标用于监控:

struct SpeculativeEngineStepMetrics {
    int64_t propose_time_us;    // 提议耗时
    int64_t score_time_us;      // 验证耗时
    int64_t sampler_time_us;    // 采样耗时
    int64_t propose_token_num;  // 提议 token 数
    int64_t accept_token_num;   // 接受 token 数
    int64_t stream_num;         // 处理的流数
};

接受率 = accept_token_num / propose_token_num

5. 设计特点​

特点说明
模块化Propose/Score/Sample 解耦，易于扩展新算法
多策略支持同时支持 Vanilla、MTP、EAGLE 等多种推测策略
批处理优化支持批量采样加速 GPU 利用
TP 同步支持张量并行场景下的同步 (tpSyncDisableSPRun)
Prefill 优化MTP 模式下支持 prefill 和 decode 分离
回退机制可禁用推测解码，退化为普通解码

6. 文件结构总结​

speculative_engine/
├── SpeculativeEngine.h/cc      # 主引擎
├── SpeculativeScheduler.h/cc   # 调度器
├── SpeculativeGatherBatchScheduler.h # 批量调度器
├── propose_executor/
│   ├── ProposeExecutor.h/cc    # 提议执行器基类
│   ├── VanillaExecutor.h/cc    # 标准草稿模型
│   ├── MTPExecutor.h/cc        # MTP 执行器
│   ├── EagleExecutor.h/cc      # EAGLE 执行器
│   ├── MTPStream.h             # MTP 流
│   ├── EagleStream.h           # EAGLE 流
│   └── MTPBatchStreamProcessor.h/cc  # MTP 批处理器
├── score_executor/
│   ├── ScoreExecutor.h/cc      # 验证执行器
│   ├── ScoreStream.h           # 验证流
│   └── ScoreBatchStreamProcessor.h/cc # 批处理器
├── speculative_sampler/
│   └── SpeculativeSampler.h/cc # 推测采样器
└── test/                       # 测试文件