1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 加载 target model（大模型）
model_name = "meta-llama/Llama-2-7b-chat-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# 加载 draft model（小模型）
draft_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
draft_model = AutoModelForCausalLM.from_pretrained(
    draft_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

prompt = "Explain the concept of speculative decoding in three sentences."
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# 普通生成（baseline）
output_normal = model.generate(
    **inputs,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
)

# Assisted Generation
output_assisted = model.generate(
    **inputs,
    max_new_tokens=100,
    do_sample=True,
    temperature=0.7,
    assistant_model=draft_model,  # 关键参数
)

print("Normal:", tokenizer.decode(output_normal[0], skip_special_tokens=True))
print("Assisted:", tokenizer.decode(output_assisted[0], skip_special_tokens=True))

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

import torch
import torch.nn.functional as F

def speculative_decoding_step(
    draft_model,
    target_model,
    prefix_ids,
    gamma=5,
    temperature=1.0
):
    """
    一次 speculative decoding 迭代。
    prefix_ids: shape [1, seq_len]
    返回: accepted token IDs list
    """
    device = prefix_ids.device
    
    # Phase 1: Draft — 小模型自回归生成 γ 个候选
    draft_ids = prefix_ids.clone()
    for _ in range(gamma):
        with torch.no_grad():
            logits = draft_model(draft_ids).logits[:, -1, :] / temperature
        probs = F.softmax(logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        draft_ids = torch.cat([draft_ids, next_token], dim=-1)
    
    draft_tokens = draft_ids[:, prefix_ids.shape[1]:]  # [1, gamma]
    
    # Phase 2: Verify — 大模型一次并行前向传播
    with torch.no_grad():
        target_logits = target_model(draft_ids).logits / temperature
    # 获取每个位置的概率分布
    # target_logits shape: [1, prefix_len + gamma, vocab_size]
    target_probs = F.softmax(target_logits[:, prefix_ids.shape[1]-1:, :], dim=-1)
    draft_probs = F.softmax(
        draft_model(draft_ids).logits[:, prefix_ids.shape[1]-1:, :] / temperature,
        dim=-1
    )
    
    # Phase 3: Accept/Reject via Modified Rejection Sampling
    accepted = []
    for i in range(gamma):
        token_id = draft_tokens[0, i].item()
        p_i = target_probs[0, i, token_id].item()
        q_i = draft_probs[0, i, token_id].item()
        
        accept_prob = min(1.0, p_i / (q_i + 1e-10))
        
        if torch.rand(1).item() < accept_prob:
            accepted.append(token_id)
        else:
            # Reject: resample from (p_i - q_i)_+
            adjusted = target_probs[0, i] - draft_probs[0, i]
            adjusted = torch.clamp(adjusted, min=0)
            if adjusted.sum() > 0:
                adjusted /= adjusted.sum()
                resampled = torch.multinomial(adjusted, 1).item()
                accepted.append(resampled)
            break
    
    return accepted