投机解码深度解析：让 LLM 推理快 3 倍

1. Draft Model 快速生成 k 个候选 token: [t1, t2, ..., tk]
2. Target Model 一次前向传播计算这 k 个位置的 logits
3. 对每个候选 token，按拒绝采样决定接受/拒绝:
   - 若 p_target(ti) / p_draft(ti) >= 1: 接受
   - 否则以概率 p_target(ti)/p_draft(ti) 接受，否则拒绝并从调整分布重采样
4. 接受的 token 加入序列，从拒绝点重新开始

class MedusaHead(nn.Module):
    def __init__(self, d_model, vocab_size):
        super().__init__()
        self.linear = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.SiLU(),
            nn.Linear(d_model, vocab_size)
        )

    def forward(self, hidden_states):
        return self.linear(hidden_states)

class MedusaModel(nn.Module):
    def __init__(self, base_model, num_heads=4):
        super().__init__()
        self.base_model = base_model
        self.medusa_heads = nn.ModuleList([
            MedusaHead(base_model.config.d_model, 
                       base_model.config.vocab_size)
            for _ in range(num_heads)
        ])
    
    def forward(self, input_ids):
        # 基础模型前向传播
        hidden = self.base_model(input_ids, output_hidden_states=True)
        h = hidden.hidden_states[-1]
        
        # 主头预测 next token (heads[0] = 原始 LM head)
        main_logits = self.base_model.lm_head(h)
        
        # Medusa 头预测 future tokens
        future_logits = [head(h) for head in self.medusa_heads]
        # future_logits[0] = t+2 的预测
        # future_logits[1] = t+3 的预测
        # ...
        return main_logits, future_logits

          [the]
         /  |  \
     [cat] [dog] [bird]
      / \    |
   [sat] [ran] [barked]

def tree_attention_verify(candidates_tree, target_model):
    # 将树展平为序列
    flat_tokens = tree.flatten()  # [the, cat, dog, bird, sat, ran, barked]
    # 构造因果掩码：只允许父→子方向
    attention_mask = build_tree_mask(tree)  # (n, n) 三角矩阵
    
    logits = target_model(flat_tokens, attention_mask=attention_mask)
    
    # 从根到叶逐层验证，找到最长接受路径
    best_path = find_longest_accepted_path(tree, logits)
    return best_path

class EAGLEDraftNet(nn.Module):
    def __init__(self, d_model, vocab_size):
        super().__init__()
        # 融合隐状态 + embedding
        self.fuse = nn.Linear(d_model + d_model, d_model)
        # 轻量自回归层
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(d_model, nhead=8, 
                                        batch_first=True)
            for _ in range(2)
        ])
        self.lm_head = nn.Linear(d_model, vocab_size)

    def forward(self, hidden_states, token_embeddings, num_steps=5):
        x = self.fuse(torch.cat([hidden_states, token_embeddings], dim=-1))
        drafts = []
        h = x
        for _ in range(num_steps):
            for layer in self.layers:
                h = layer(h.unsqueeze(0)).squeeze(0)
            logits = self.lm_head(h)
            next_token = logits.argmax(dim=-1)
            drafts.append(next_token)
            h = self.fuse(torch.cat([h, token_embeddings[next_token]], dim=-1))
        return drafts

def lookahead_decoding(model, prompt, n_window=5, max_steps=100):
    seq = prompt
    for step in range(max_steps):
        # 并行预测 n_window 个 token
        candidate = seq + [0] * n_window  # padding
        logits = model(candidate)
        
        # Jacobi 更新
        changed = False
        for i in range(n_window):
            new_token = logits[len(seq) + i].argmax()
            if candidate[len(seq) + i] != new_token:
                candidate[len(seq) + i] = new_token
                changed = True
        
        if not changed:
            break  # 收敛
        seq = candidate[:len(seq) + n_window]
    
    return seq