第13章 Speculative Decoding与推理加速¶

# 直观理解 Memory-bound
# 70B模型 Decode阶段每步：
model_params = 70e9 * 2  # 140GB (FP16)
kv_cache_per_step = 0.5  # ~0.5GB (假设，实际取决于序列长度和hidden_size)
total_read = model_params + kv_cache_per_step  # ~140.5GB

# A100 80GB GPU 带宽 = 2TB/s
bandwidth = 2e12  # bytes/s
min_time_per_token = total_read / bandwidth  # ~70ms
# 注：这是简化假设，实际读取时间受KV Cache命中、模型分片等因素影响
# 即便计算瞬间完成，仅读取权重就需要 ~70ms → ~14 tokens/s

# 但实际计算量极小：
flops_per_token = 2 * 70e9  # ~140 GFLOPs
a100_flops = 312e12  # 312 TFLOPS (FP16)
compute_time = flops_per_token / a100_flops  # ~0.45ms

# 计算时间 0.45ms vs 读取时间 70ms → 计算资源严重浪费！

import torch
import time

def benchmark_matmul():
    """对比单向量 vs 批量矩阵乘的 GPU利用率"""
    device = torch.device('cuda')
    weight = torch.randn(4096, 4096, device=device, dtype=torch.float16)

    # 单 token 解码：矩阵-向量乘 (4096, 4096) x (4096, 1)
    x_single = torch.randn(1, 4096, device=device, dtype=torch.float16)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(1000):
        _ = x_single @ weight.T
    torch.cuda.synchronize()
    t_single = (time.perf_counter() - t0) / 1000

    # 批量验证：矩阵-矩阵乘 (4096, 4096) x (4096, 8)
    x_batch = torch.randn(8, 4096, device=device, dtype=torch.float16)
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(1000):
        _ = x_batch @ weight.T
    torch.cuda.synchronize()
    t_batch = (time.perf_counter() - t0) / 1000

    print(f"单token: {t_single*1e6:.1f} μs (1 token)")
    print(f"批量8:   {t_batch*1e6:.1f} μs (8 tokens)")
    print(f"批量8每token: {t_batch/8*1e6:.1f} μs")
    print(f"吞吐提升: {t_single / (t_batch/8):.1f}x")

# 典型输出（A100）：
# 单token: 42.3 μs (1 token)
# 批量8:   48.7 μs (8 tokens)
# 批量8每token: 6.1 μs
# 吞吐提升: 6.9x

import torch
import torch.nn.functional as F

def speculative_sampling(
    target_logits: torch.Tensor,  # (K+1, vocab_size) - Target模型输出
    draft_probs: torch.Tensor,    # (K, vocab_size)   - Draft模型采样时的概率
    draft_tokens: torch.Tensor,   # (K,)              - Draft模型采样的token
    temperature: float = 1.0
) -> tuple[torch.Tensor, int]:
    """
    执行 Speculative Decoding 的验证步骤。

    Returns:
        accepted_tokens: 最终接受的 token 序列
        n_accepted: 接受的 Draft token 数量
    """
    K = draft_tokens.shape[0]
    vocab_size = target_logits.shape[-1]  # [-1]负索引取最后元素

    # Target 模型的概率分布
    target_probs = F.softmax(target_logits[:K] / temperature, dim=-1)  # (K, V)  # F.xxx PyTorch函数式API

    accepted_tokens = []
    n_accepted = 0

    for i in range(K):
        token = draft_tokens[i]
        p_i = target_probs[i, token].item()   # Target 概率
        q_i = draft_probs[i, token].item()    # Draft 概率

        # 接受概率
        accept_prob = min(1.0, p_i / (q_i + 1e-10))

        r = torch.rand(1).item()
        if r < accept_prob:
            # 接受 Draft 的猜测
            accepted_tokens.append(token.item())
            n_accepted += 1
        else:
            # 拒绝：从修正分布采样
            # p'(x) = norm(max(0, p(x) - q(x)))
            corrected = torch.clamp(target_probs[i] - draft_probs[i], min=0)
            corrected = corrected / (corrected.sum() + 1e-10)
            replacement = torch.multinomial(corrected, 1).item()
            accepted_tokens.append(replacement)
            break

    # 如果全部接受，从 p_{K+1} 采样 bonus token
    if n_accepted == K:
        bonus_probs = F.softmax(target_logits[K] / temperature, dim=-1)
        bonus_token = torch.multinomial(bonus_probs, 1).item()
        accepted_tokens.append(bonus_token)

    return torch.tensor(accepted_tokens), n_accepted

import numpy as np

def compute_speedup(alpha: float, K: int, c: float) -> float:
    """计算 Speculative Decoding 的理论加速比"""
    expected_tokens = (1 - alpha**(K+1)) / (1 - alpha)
    wall_time_ratio = K * c + 1
    return expected_tokens / wall_time_ratio

# 不同参数下的加速比
print("=" * 60)
print(f"{'接受率α':<10} {'K':<5} {'速度比c':<10} {'加速比':<10}")
print("=" * 60)
for alpha in [0.6, 0.7, 0.8, 0.9]:
    for K in [3, 5, 8]:
        for c in [0.05, 0.1]:
            s = compute_speedup(alpha, K, c)
            print(f"{alpha:<10} {K:<5} {c:<10} {s:<10.2f}")
    print("-" * 60)

# 典型输出：
# α=0.8, K=5, c=0.05 → 加速比 ≈ 2.95x
# α=0.9, K=5, c=0.05 → 加速比 ≈ 3.69x
# α=0.7, K=5, c=0.10 → 加速比 ≈ 1.94x

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
from typing import Optional

class SpeculativeDecoder:
    """
    完整的 Speculative Decoding 实现。
    使用小模型 (draft) 猜测 + 大模型 (target) 验证。
    """
    def __init__(
        self,
        target_model_name: str,
        draft_model_name: str,
        device: str = "cuda",
        dtype: torch.dtype = torch.float16,
    ):
        print(f"Loading target model: {target_model_name}")
        self.target_model = AutoModelForCausalLM.from_pretrained(
            target_model_name, torch_dtype=dtype, device_map=device
        )
        self.target_model.eval()  # eval()评估模式

        print(f"Loading draft model: {draft_model_name}")
        self.draft_model = AutoModelForCausalLM.from_pretrained(
            draft_model_name, torch_dtype=dtype, device_map=device
        )
        self.draft_model.eval()

        self.tokenizer = AutoTokenizer.from_pretrained(target_model_name)
        self.device = device

    @torch.no_grad()  # 禁用梯度计算，节省内存
    def draft_generate(
        self,
        input_ids: torch.Tensor,
        K: int,
        temperature: float = 1.0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Draft 模型串行生成 K 个候选 token。
        Returns:
            draft_tokens: (K,) 候选 token ids
            draft_probs:  (K, vocab_size) 候选 token 的概率分布
        """
        draft_tokens = []
        draft_probs = []
        current_ids = input_ids.clone()

        for _ in range(K):
            outputs = self.draft_model(current_ids)
            logits = outputs.logits[:, -1, :]  # (1, vocab_size)
            probs = F.softmax(logits / temperature, dim=-1)
            token = torch.multinomial(probs, 1)  # (1, 1)

            draft_tokens.append(token.squeeze())  # squeeze压缩维度
            draft_probs.append(probs.squeeze())
            current_ids = torch.cat([current_ids, token], dim=-1)  # torch.cat沿已有维度拼接张量

        return torch.stack(draft_tokens), torch.stack(draft_probs)  # torch.stack沿新维度拼接张量

    @torch.no_grad()
    def target_verify(
        self,
        input_ids: torch.Tensor,
        draft_tokens: torch.Tensor,
    ) -> torch.Tensor:
        """
        Target 模型一次前向传播验证 K 个候选 token。
        Returns:
            logits: (K+1, vocab_size) — K个验证位置 + 1个bonus位置
        """
        # 将 draft tokens 拼接到输入
        candidate_ids = torch.cat([
            input_ids,
            draft_tokens.unsqueeze(0)  # unsqueeze增加一个维度
        ], dim=-1)  # (1, seq_len + K)

        outputs = self.target_model(candidate_ids)
        # 取最后 K+1 个位置的 logits
        # 位置 -(K+1): 验证第1个draft token
        # 位置 -1: bonus token 位置
        logits = outputs.logits[0, -(len(draft_tokens)+1):, :]
        return logits

    @torch.no_grad()
    def speculative_sample(
        self,
        target_logits: torch.Tensor,
        draft_probs: torch.Tensor,
        draft_tokens: torch.Tensor,
        temperature: float = 1.0,
    ) -> tuple[torch.Tensor, int]:
        """接受-拒绝采样"""
        K = draft_tokens.shape[0]
        target_probs = F.softmax(target_logits[:K] / temperature, dim=-1)

        accepted = []
        n_accepted = 0

        for i in range(K):
            tok = draft_tokens[i]
            p = target_probs[i, tok].item()
            q = draft_probs[i, tok].item()

            if torch.rand(1).item() < min(1.0, p / (q + 1e-10)):
                accepted.append(tok.item())
                n_accepted += 1
            else:
                # 修正采样
                corrected = torch.clamp(target_probs[i] - draft_probs[i], min=0)
                corrected /= corrected.sum() + 1e-10
                replacement = torch.multinomial(corrected, 1).item()
                accepted.append(replacement)
                break

        if n_accepted == K:
            bonus_probs = F.softmax(target_logits[K] / temperature, dim=-1)
            bonus = torch.multinomial(bonus_probs, 1).item()
            accepted.append(bonus)

        return torch.tensor(accepted, device=self.device), n_accepted

    @torch.no_grad()
    def generate(
        self,
        prompt: str,
        max_new_tokens: int = 128,
        K: int = 5,
        temperature: float = 1.0,
    ) -> str:
        """
        使用 Speculative Decoding 生成文本。

        Args:
            prompt: 输入提示
            max_new_tokens: 最大生成 token 数
            K: 投机长度（每轮猜测的 token 数）
            temperature: 采样温度
        """
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
        generated_tokens = 0
        total_draft = 0
        total_accepted = 0

        while generated_tokens < max_new_tokens:
            # Step 1: Draft 模型猜测 K 个 token
            draft_tokens, draft_probs = self.draft_generate(
                input_ids, K, temperature
            )

            # Step 2: Target 模型一次性验证
            target_logits = self.target_verify(input_ids, draft_tokens)

            # Step 3: 接受-拒绝采样
            accepted_tokens, n_accepted = self.speculative_sample(
                target_logits, draft_probs, draft_tokens, temperature
            )

            # Step 4: 更新输入序列
            input_ids = torch.cat([
                input_ids,
                accepted_tokens.unsqueeze(0)
            ], dim=-1)

            generated_tokens += len(accepted_tokens)
            total_draft += K
            total_accepted += n_accepted

            # 检查是否生成了 EOS
            if self.tokenizer.eos_token_id in accepted_tokens.tolist():
                break

        # 统计信息
        accept_rate = total_accepted / total_draft if total_draft > 0 else 0
        print(f"\n[Stats] 生成 {generated_tokens} tokens | "
              f"接受率 {accept_rate:.1%} | "
              f"Draft {total_draft} tokens, 接受 {total_accepted}")

        return self.tokenizer.decode(input_ids[0], skip_special_tokens=True)

# 使用示例
if __name__ == "__main__":
    # 注意：meta-llama 模型需要在 HuggingFace 上申请访问权限
    # 替代方案：使用公开模型如 "openlm-research/open_llama_7b"
    decoder = SpeculativeDecoder(
        target_model_name="meta-llama/Llama-3.1-8B-Instruct",  # 需要HuggingFace访问权限
        draft_model_name="meta-llama/Llama-3.2-1B-Instruct",   # 需要HuggingFace访问权限
    )

    result = decoder.generate(
        prompt="Explain quantum computing in simple terms:",
        max_new_tokens=200,
        K=5,
        temperature=0.7,
    )
    print(result)

import torch
import torch.nn as nn
import torch.nn.functional as F

class DraftDistillationTrainer:
    """
    从 Target 模型知识蒸馏训练 Draft 模型。
    使用 KL 散度让 Draft 尽可能模仿 Target 的输出分布。
    """
    def __init__(self, target_model, draft_model, temperature=2.0):
        self.target = target_model
        self.draft = draft_model
        self.temperature = temperature

    def distillation_loss(self, input_ids, attention_mask=None):
        """KL(target || draft) — 让 Draft 模拟 Target 的分布"""
        with torch.no_grad():
            target_logits = self.target(
                input_ids, attention_mask=attention_mask
            ).logits
            target_probs = F.softmax(
                target_logits / self.temperature, dim=-1
            )

        draft_logits = self.draft(
            input_ids, attention_mask=attention_mask
        ).logits
        draft_log_probs = F.log_softmax(
            draft_logits / self.temperature, dim=-1
        )

        # KL 散度
        kl_loss = F.kl_div(
            draft_log_probs, target_probs, reduction='batchmean'
        ) * (self.temperature ** 2)

        return kl_loss

from collections import defaultdict

class NGramDraft:
    """
    基于 N-gram 统计的超轻量 Draft。
    几乎零延迟，但接受率较低。
    适合重复性高的场景（代码生成、模板文本）。
    """
    def __init__(self, n: int = 3):
        self.n = n
        # 两层嵌套默认字典：外层key不存在时自动创建新的defaultdict(int)，内层key不存在时自动初始化为0，用于N-gram→next_token→计数
        self.counts = defaultdict(lambda: defaultdict(int))  # defaultdict访问不存在的键时返回默认值

    def fit(self, token_sequences: list[list[int]]):
        """从训练语料构建 N-gram 统计"""
        for seq in token_sequences:
            for i in range(len(seq) - self.n):
                context = tuple(seq[i:i+self.n])
                next_token = seq[i+self.n]
                self.counts[context][next_token] += 1

    def draft(self, context_tokens: list[int], K: int) -> list[int]:
        """基于 N-gram 猜测接下来 K 个 token"""
        guesses = []
        ctx = list(context_tokens)
        for _ in range(K):
            key = tuple(ctx[-self.n:])
            candidates = self.counts.get(key, {})
            if candidates:
                # 选最高频的 token
                best = max(candidates, key=candidates.get)
                guesses.append(best)
                ctx.append(best)
            else:
                break
        return guesses

def prompt_lookup_draft(
    input_ids: list[int],
    prompt_length: int,
    K: int = 5,
    n_match: int = 3,
) -> list[int]:
    """
    Prompt Lookup Decoding: 在 prompt 中寻找匹配当前后缀的片段，
    将其后续 token 作为 draft 候选。

    特别适合：摘要、翻译、重写等与输入高度相关的任务。
    """
    if len(input_ids) < n_match:
        return []

    # 当前生成序列的最后 n_match 个 token
    suffix = input_ids[-n_match:]

    # 在 prompt 中搜索相同的 n-gram
    candidates = []
    for i in range(prompt_length - n_match):
        if input_ids[i:i+n_match] == suffix:
            # 找到匹配！取后续 K 个 token 作为猜测
            end = min(i + n_match + K, prompt_length)
            draft = input_ids[i+n_match:end]
            if draft:
                candidates.append(draft)

    # 返回最长的匹配候选
    if candidates:
        return max(candidates, key=len)[:K]
    return []

import torch
import torch.nn as nn

class MedusaHead(nn.Module):  # 继承nn.Module定义网络层
    """
    单个 Medusa 预测头。
    结构：1层残差MLP + 共享的LM Head权重。
    """
    def __init__(self, hidden_size: int, vocab_size: int):
        super().__init__()  # super()调用父类方法
        self.fc = nn.Linear(hidden_size, hidden_size, bias=False)
        self.act = nn.SiLU()
        self.ln = nn.LayerNorm(hidden_size)
        # 通常复用原始模型的 lm_head 权重
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: (batch, seq_len, hidden_size)
        Returns:
            logits: (batch, seq_len, vocab_size)
        """
        x = self.act(self.fc(hidden_states))
        x = self.ln(x + hidden_states)  # 残差连接
        return self.lm_head(x)

class MedusaModel(nn.Module):
    """
    在基础模型上添加 K 个 Medusa 预测头。
    """
    def __init__(self, base_model, num_heads: int = 4, hidden_size: int = 4096, vocab_size: int = 32000):
        super().__init__()
        self.base_model = base_model
        self.medusa_heads = nn.ModuleList([
            MedusaHead(hidden_size, vocab_size)
            for _ in range(num_heads)
        ])
        # 共享原始模型的 lm_head 权重
        for head in self.medusa_heads:
            head.lm_head.weight = base_model.lm_head.weight

    @torch.no_grad()
    def forward(self, input_ids, attention_mask=None):
        outputs = self.base_model.model(
            input_ids, attention_mask=attention_mask
        )
        hidden = outputs.last_hidden_state  # (B, L, D)

        # 原始 LM Head 预测 t+1
        base_logits = self.base_model.lm_head(hidden)

        # Medusa Heads 预测 t+2, t+3, ...
        medusa_logits = [head(hidden) for head in self.medusa_heads]

        return base_logits, medusa_logits

def build_medusa_tree(
    medusa_logits: list[torch.Tensor],
    top_k: int = 3,
) -> list[list[int]]:
    """
    构建 Medusa 候选树。

    Args:
        medusa_logits: K 个 Medusa Head 的输出 logits
        top_k: 每个 Head 取 top-k 候选

    Returns:
        candidates: 所有候选路径 [[t1, t2, t3], [t1, t2', t3'], ...]
    """
    K = len(medusa_logits)
    # 每个 Head 取 top-k
    topk_per_head = []
    for logits in medusa_logits:
        probs = torch.softmax(logits[:, -1, :], dim=-1)
        topk = torch.topk(probs, top_k, dim=-1)
        topk_per_head.append(topk.indices[0].tolist())  # (top_k,)

    # 构建所有路径组合（可用剪枝策略减少路径数）
    # 简化示例：仅取联合概率最高的若干路径
    candidates = []
    for t1 in topk_per_head[0]:
        for t2 in topk_per_head[1]:
            for t3 in topk_per_head[2] if K > 2 else [None]:
                path = [t1, t2]
                if t3 is not None:
                    path.append(t3)
                candidates.append(path)

    return candidates

def tree_attention_mask(num_candidates: int, seq_len: int) -> torch.Tensor:
    """
    构建 Tree Attention 的因果掩码。
    每条候选路径只能看到自己的祖先节点和共享前缀。
    """
    # 简化示例：具体实现需根据树的拓扑结构
    mask = torch.zeros(num_candidates, seq_len + num_candidates)
    # 所有候选都能看到完整 prefix
    mask[:, :seq_len] = 1
    # 每个候选只能看到自己路径上的祖先
    for i in range(num_candidates):
        mask[i, seq_len + i] = 1
    return mask

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader

def train_medusa_heads(
    model: MedusaModel,
    dataloader: DataLoader,
    num_epochs: int = 3,
    lr: float = 1e-4,
):
    """
    Medusa-1 训练：冻结基础模型，仅训练 Medusa Heads。
    训练目标：让 Head_i 预测位置 t+i+2 的 token。
    """
    # 冻结基础模型
    for param in model.base_model.parameters():
        param.requires_grad = False

    # 只优化 Medusa Heads（不含共享的 lm_head.weight）
    medusa_params = []
    for head in model.medusa_heads:
        medusa_params.extend([head.fc.weight, head.ln.weight, head.ln.bias])

    optimizer = torch.optim.AdamW(medusa_params, lr=lr)

    for epoch in range(num_epochs):
        total_loss = 0
        for batch in dataloader:
            input_ids = batch["input_ids"].cuda()
            labels = batch["labels"].cuda()  # 与 input_ids 对齐的标签

            base_logits, medusa_logits = model(input_ids)

            loss = 0
            for i, m_logits in enumerate(medusa_logits):  # enumerate同时获取索引和元素
                # Head_i 预测 t+i+2 的 token
                shift = i + 2  # head_0 预测 t+2, head_1 预测 t+3, ...
                if shift < labels.shape[1]:
                    pred = m_logits[:, :-shift, :]  # 预测部分
                    target = labels[:, shift:]       # 目标部分
                    loss += F.cross_entropy(  # F.cross_entropy PyTorch函数式交叉熵损失
                        pred.reshape(-1, pred.shape[-1]),  # 重塑张量形状
                        target.reshape(-1),
                        ignore_index=-100,
                    )

            optimizer.zero_grad()  # 清零梯度
            loss.backward()  # 反向传播计算梯度
            optimizer.step()  # 更新参数
            total_loss += loss.item()

        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch+1}/{num_epochs} | Loss: {avg_loss:.4f}")

import torch
import torch.nn as nn

class EAGLEDraftHead(nn.Module):
    """
    EAGLE 的特征级别 Draft 网络。
    输入: Target 模型的隐藏特征 + token embedding
    输出: 预测下一步的隐藏特征
    """
    def __init__(self, hidden_size: int = 4096):
        super().__init__()
        # 融合 token embedding 和 hidden state
        self.fc_in = nn.Linear(hidden_size * 2, hidden_size, bias=False)
        # 轻量级 Transformer 层（通常1层）
        self.transformer_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=32,
            dim_feedforward=hidden_size * 4,
            batch_first=True,
            norm_first=True,
        )
        self.ln = nn.LayerNorm(hidden_size)

    def forward(
        self,
        hidden_states: torch.Tensor,
        token_embeddings: torch.Tensor,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states: Target 模型的隐藏特征 (B, L, D)
            token_embeddings: 对应 token 的 embedding (B, L, D)
        Returns:
            预测的下一步隐藏特征 (B, L, D)
        """
        # 拼接特征和 embedding
        fused = torch.cat([hidden_states, token_embeddings], dim=-1)
        x = self.fc_in(fused)
        x = self.transformer_layer(x)
        return self.ln(x)

class EAGLEDecoder:
    """EAGLE Speculative Decoding 推理流程"""

    def __init__(self, target_model, eagle_head, lm_head, embed_tokens):
        self.target = target_model
        self.eagle_head = eagle_head
        self.lm_head = lm_head        # 复用 Target 的 LM Head
        self.embed = embed_tokens      # 复用 Target 的 Embedding

    @torch.no_grad()
    def draft_with_features(
        self,
        last_hidden: torch.Tensor,  # Target 模型最后一层输出
        last_token_id: int,
        K: int = 5,
    ) -> tuple[list[int], list[torch.Tensor]]:
        """
        使用 EAGLE Head 在特征空间做 K 步外推。
        """
        draft_tokens = []
        draft_features = []

        current_hidden = last_hidden[:, -1:, :]  # (1, 1, D)
        current_embed = self.embed(
            torch.tensor([[last_token_id]], device=last_hidden.device)
        )  # (1, 1, D)

        for _ in range(K):
            # 预测下一步特征
            predicted_hidden = self.eagle_head(current_hidden, current_embed)
            draft_features.append(predicted_hidden)

            # 用 Target 的 LM Head 将特征映射到 token
            logits = self.lm_head(predicted_hidden)  # (1, 1, V)
            token = logits.argmax(dim=-1).item()
            draft_tokens.append(token)

            # 更新 Draft 状态
            current_hidden = predicted_hidden
            current_embed = self.embed(
                torch.tensor([[token]], device=last_hidden.device)
            )

        return draft_tokens, draft_features

class DynamicDraftTree:
    """EAGLE-2 的动态 Draft 树策略"""

    def __init__(
        self,
        max_depth: int = 8,
        max_width: int = 5,
        confidence_threshold: float = 0.4,
        max_candidates: int = 60,
    ):
        self.max_depth = max_depth
        self.max_width = max_width
        self.conf_threshold = confidence_threshold
        self.max_candidates = max_candidates

    def build_tree(self, eagle_head, hidden, embed, lm_head):
        """
        动态构建 Draft 树。
        高置信度分支: 继续深入
        低置信度分支: 剪枝停止
        """
        tree_nodes = []  # [(token, parent_idx, depth, confidence)]
        active_leaves = [(hidden[:, -1:, :], embed[:, -1:, :], -1, 0)]

        while active_leaves and len(tree_nodes) < self.max_candidates:
            new_leaves = []
            for h, e, parent_idx, depth in active_leaves:
                if depth >= self.max_depth:
                    continue

                # EAGLE Head 预测
                pred_hidden = eagle_head(h, e)
                logits = lm_head(pred_hidden).squeeze()
                probs = torch.softmax(logits, dim=-1)

                # 取 top-k
                topk_probs, topk_ids = probs.topk(self.max_width)

                for prob, tid in zip(topk_probs, topk_ids):  # zip按位置配对
                    conf = prob.item()
                    if conf < self.conf_threshold:
                        continue  # 置信度低，剪枝

                    node_idx = len(tree_nodes)
                    tree_nodes.append((tid.item(), parent_idx, depth + 1, conf))

                    # 准备该节点的特征用于下一层
                    new_embed = embed.new_zeros(1, 1, embed.shape[-1])
                    new_leaves.append((pred_hidden, new_embed, node_idx, depth + 1))

            active_leaves = new_leaves

        return tree_nodes

import torch

def lookahead_decoding_step(model, input_ids, n_lookahead=5, n_iterations=3):
    """
    Lookahead Decoding 的核心步骤。
    维护一个 N-token 的猜测窗口，通过 Jacobi 迭代并行求解。

    优点: 不需要 Draft 模型
    缺点: 需要多次迭代，加速比通常低于 Speculative Decoding

    Args:
        model: 语言模型
        input_ids: 当前序列 (1, seq_len)
        n_lookahead: 前瞻窗口大小
        n_iterations: Jacobi 迭代次数
    """
    device = input_ids.device
    seq_len = input_ids.shape[1]

    # 初始化猜测窗口（随机或启发式）
    guess = torch.randint(0, model.config.vocab_size, (1, n_lookahead), device=device)

    for iteration in range(n_iterations):
        # 将猜测拼接到序列
        full_ids = torch.cat([input_ids, guess], dim=-1)

        # 构造因果掩码（标准自回归）
        outputs = model(full_ids)
        logits = outputs.logits  # (1, seq_len + n_lookahead, vocab_size)

        # 并行更新所有猜测位置
        new_guess = logits[0, seq_len-1:seq_len+n_lookahead-1, :].argmax(dim=-1)
        new_guess = new_guess.unsqueeze(0)

        # 检查是否收敛
        if torch.equal(new_guess, guess):
            break
        guess = new_guess

    # 验证：找到第一个不一致的位置
    final_ids = torch.cat([input_ids, guess], dim=-1)
    verify_outputs = model(final_ids)
    verify_logits = verify_outputs.logits

    n_accepted = 0
    for i in range(n_lookahead):
        predicted = verify_logits[0, seq_len + i - 1, :].argmax()
        if predicted == guess[0, i]:
            n_accepted += 1
        else:
            break

    return guess[0, :n_accepted+1], n_accepted

class ContinuousBatchingScheduler:
    """
    连续批处理调度器（简化示意）。
    每次 Decode 迭代后检查是否有请求完成或新请求到达。
    """
    def __init__(self, model, max_batch_size=64):
        self.model = model
        self.max_batch = max_batch_size
        self.active_requests = []     # 正在生成的请求
        self.waiting_queue = []       # 等待队列

    def iteration_step(self):
        """执行一次 Decode 迭代"""
        # 1. 移除已完成的请求
        finished = [r for r in self.active_requests if r.is_done()]
        for r in finished:
            self.active_requests.remove(r)
            r.callback(r.generated_text)

        # 2. 从等待队列填充空位
        while (len(self.active_requests) < self.max_batch
               and self.waiting_queue):
            new_req = self.waiting_queue.pop(0)
            self.active_requests.append(new_req)

        if not self.active_requests:
            return

        # 3. 组装 batch，执行一次前向传播
        batch_ids = self._pad_batch()
        with torch.no_grad():
            outputs = self.model(batch_ids)

        # 4. 每个请求采样下一个 token
        for i, req in enumerate(self.active_requests):
            next_token = outputs.logits[i, -1, :].argmax().item()
            req.append_token(next_token)

    def _pad_batch(self):
        """将不同长度的请求 padding 到相同长度"""
        max_len = max(len(r.tokens) for r in self.active_requests)
        batch = torch.zeros(len(self.active_requests), max_len, dtype=torch.long)
        for i, r in enumerate(self.active_requests):
            batch[i, :len(r.tokens)] = torch.tensor(r.tokens)
        return batch.cuda()

class PrefixCacheManager:
    """
    前缀 KV-Cache 管理器。
    适用场景：所有请求共享相同 system prompt。
    """
    def __init__(self, model, max_cache_entries=100):
        self.model = model
        self.cache = {}  # hash(prefix_tokens) → KV-Cache
        self.max_entries = max_cache_entries

    def get_or_compute_prefix_cache(self, prefix_tokens: tuple):
        """查找或计算前缀的 KV-Cache"""
        cache_key = hash(prefix_tokens)

        if cache_key in self.cache:
            return self.cache[cache_key]  # 缓存命中！

        # 缓存未命中，计算 KV-Cache
        input_ids = torch.tensor([list(prefix_tokens)]).cuda()
        with torch.no_grad():
            outputs = self.model(input_ids, use_cache=True)

        kv_cache = outputs.past_key_values

        # LRU 淘汰
        if len(self.cache) >= self.max_entries:
            oldest = next(iter(self.cache))
            del self.cache[oldest]

        self.cache[cache_key] = kv_cache
        return kv_cache

    def generate_with_cached_prefix(self, prefix_tokens, continuation_tokens):
        """使用缓存的前缀 KV-Cache 直接从续写部分开始计算"""
        kv_cache = self.get_or_compute_prefix_cache(tuple(prefix_tokens))

        # 仅需处理续写部分的 token
        cont_ids = torch.tensor([continuation_tokens]).cuda()
        with torch.no_grad():
            outputs = self.model(
                cont_ids,
                past_key_values=kv_cache,
                use_cache=True,
            )
        return outputs

class ChunkedPrefillScheduler:
    """分块预填充调度器"""

    def __init__(self, model, chunk_size=1024):
        self.model = model
        self.chunk_size = chunk_size  # 每块处理的 token 数

    def chunked_prefill(self, prompt_tokens, decode_requests=None):
        """
        将长 prompt 的 Prefill 拆分为多个 chunk 交替执行。
        减少 Prefill 对 Decode 延迟的影响 (TTFT ↓, ITL 稳定)。
        """
        total_len = len(prompt_tokens)
        num_chunks = (total_len + self.chunk_size - 1) // self.chunk_size
        kv_cache = None

        for i in range(num_chunks):
            start = i * self.chunk_size
            end = min(start + self.chunk_size, total_len)
            chunk = prompt_tokens[start:end]

            # 处理一块 Prefill
            chunk_ids = torch.tensor([chunk]).cuda()
            with torch.no_grad():
                outputs = self.model(
                    chunk_ids,
                    past_key_values=kv_cache,
                    use_cache=True,
                )
            kv_cache = outputs.past_key_values

            # 在 chunk 之间插入 Decode 步骤
            if decode_requests:
                self._serve_decode_batch(decode_requests)

        return kv_cache

    def _serve_decode_batch(self, requests):
        """服务当前等待中的 Decode 请求"""
        # 对活跃请求执行一步 Decode
        pass

from vllm import LLM, SamplingParams

# ===== 方法1：使用独立 Draft 模型 =====
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="meta-llama/Llama-3.2-1B-Instruct",
    num_speculative_tokens=5,      # 每轮猜测 5 个 token
    max_model_len=4096,
    tensor_parallel_size=4,         # Target 模型 4 卡并行
    # Draft 模型自动放在单卡上
)

# ===== 方法2：Prompt Lookup Decoding（不需要 Draft 模型） =====
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="[ngram]",    # 使用 N-gram 作为 Draft
    num_speculative_tokens=5,
    ngram_prompt_lookup_max=4,      # N-gram 匹配的最大长度
    ngram_prompt_lookup_min=2,
)

# ===== 方法3：EAGLE 投机解码 =====
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-70B",
    num_speculative_tokens=5,
    max_model_len=4096,
)

# 使用方式与普通 vLLM 完全一致
sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
outputs = llm.generate(["Explain quantum computing:"], sampling_params)
print(outputs[0].outputs[0].text)

# TensorRT-LLM Speculative Decoding 配置示例
import tensorrt_llm
from tensorrt_llm import BuildConfig
from tensorrt_llm.executor import ExecutorConfig

# 构建 Target 引擎
target_build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_seq_len=4096,
    speculative_decoding_mode="DRAFT_TOKENS_EXTERNAL",
)

# 构建 Draft 引擎
draft_build_config = BuildConfig(
    max_batch_size=64,
    max_input_len=2048,
    max_seq_len=4096,
)

# 运行时配置
executor_config = ExecutorConfig(
    speculative_config={
        "num_draft_tokens": 5,
        "acceptance_threshold": 0.0,  # 使用标准 Speculative Sampling
    }
)

# 在 trtllm-build 命令行中:
# trtllm-build --checkpoint_dir ./target_ckpt \
#              --output_dir ./target_engine \
#              --speculative_decoding_mode draft_tokens_external \
#              --max_draft_len 5

def should_use_speculative_decoding(
    batch_size: int,
    task_type: str,
    temperature: float,
    gpu_memory_free_gb: float,
    draft_model_size_gb: float,
) -> tuple[bool, str]:
    """
    决策函数：是否应该使用 Speculative Decoding
    """
    reasons = []

    if batch_size > 16:
        reasons.append(f"batch_size={batch_size} 过大，GPU已充分利用")
        return False, "; ".join(reasons)

    if temperature > 1.5:
        reasons.append(f"temperature={temperature} 过高，接受率将很低")
        return False, "; ".join(reasons)

    if gpu_memory_free_gb < draft_model_size_gb * 1.2:
        reasons.append("显存不足以加载 Draft 模型")
        return False, "; ".join(reasons)

    # 按任务类型预估加速比
    expected_speedup = {
        "code": 2.5, "translation": 2.2, "summary": 2.0,
        "chat": 1.8, "creative": 1.2, "math": 1.5,
    }.get(task_type, 1.5)

    if expected_speedup < 1.3:
        reasons.append(f"预估加速比 {expected_speedup}x 不值得额外复杂度")
        return False, "; ".join(reasons)

    return True, f"建议使用，预估加速 {expected_speedup}x"

class AdaptiveSpeculativeLength:
    """运行时自适应调整投机长度"""
    def __init__(self, K_min=2, K_max=10, target_accept_rate=0.75):
        self.K = 5  # 初始值
        self.K_min = K_min
        self.K_max = K_max
        self.target_rate = target_accept_rate
        self.accept_history = []

    def update(self, n_accepted, K_used):
        rate = n_accepted / K_used
        self.accept_history.append(rate)

        # 用最近 10 轮的平均接受率调整
        if len(self.accept_history) >= 10:
            avg_rate = sum(self.accept_history[-10:]) / 10
            if avg_rate > self.target_rate + 0.1:
                self.K = min(self.K + 1, self.K_max)
            elif avg_rate < self.target_rate - 0.1:
                self.K = max(self.K - 1, self.K_min)

        return self.K