EAGLE-3 推测解码实战：vLLM 基准测试冠军的技术解析与 7x 加速部署指南

# 1. 创建虚拟环境
conda create -n eagle3 python=3.10
conda activate eagle3

# 2. 安装 vLLM 0.20.2
pip install vllm==0.20.2

# 3. 安装 EAGLE-3 训练工具（可选，用于自定义模型）
pip install eagle3-specdecode

# 4. 验证安装
python -c "import vllm; print(vllm.__version__)"
# 应输出: 0.20.2

"""
EAGLE-3 推测解码完整部署脚本
包含：服务启动、性能监控、自动调优
依赖: pip install vllm==0.20.2 eagle3-specdecode
"""
import time
import json
import asyncio
from dataclasses import dataclass, field
from typing import Optional
from vllm import LLM, SamplingParams

@dataclass
class EAGLE3Config:
    """EAGLE-3 配置"""
    model: str = "meta-llama/Llama-3-70B-Instruct"
    speculator: str = "eagle3-llama3-70b-spec"
    num_speculative_tokens: int = 8
    tensor_parallel_size: int = 4
    max_model_len: int = 8192
    gpu_memory_utilization: float = 0.90
    enable_chunked_prefill: bool = True

@dataclass
class PerfMetrics:
    """性能指标"""
    total_tokens: int = 0
    total_time: float = 0.0
    accepted_tokens: int = 0
    rejected_tokens: int = 0
    requests: int = 0
    latencies: list = field(default_factory=list)

    @property
    def tokens_per_second(self) -> float:
        return self.total_tokens / self.total_time if self.total_time > 0 else 0

    @property
    def acceptance_rate(self) -> float:
        total = self.accepted_tokens + self.rejected_tokens
        return self.accepted_tokens / total if total > 0 else 0

    def report(self) -> str:
        return json.dumps({
            "tokens/s": round(self.tokens_per_second, 1),
            "acceptance_rate": f"{self.acceptance_rate:.1%}",
            "total_tokens": self.total_tokens,
            "avg_latency": f"{sum(self.latencies)/len(self.latencies):.0f}ms" if self.latencies else "N/A",
            "requests": self.requests,
        }, indent=2)

class EAGLE3Server:
    """EAGLE-3 推理服务"""

    def __init__(self, config: EAGLE3Config):
        self.config = config
        self.metrics = PerfMetrics()
        self._init_model()

    def _init_model(self):
        """初始化模型和 speculator"""
        print(f"[EAGLE-3] 加载 target model: {self.config.model}")
        print(f"[EAGLE-3] 加载 speculator: {self.config.speculator}")
        print(f"[EAGLE-3] 推测 token 数: {self.config.num_speculative_tokens}")

        self.llm = LLM(
            model=self.config.model,
            speculative_model=self.config.speculator,
            num_speculative_tokens=self.config.num_speculative_tokens,
            tensor_parallel_size=self.config.tensor_parallel_size,
            max_model_len=self.config.max_model_len,
            gpu_memory_utilization=self.config.gpu_memory_utilization,
            enable_chunked_prefill=self.config.enable_chunked_prefill,
        )
        print("[EAGLE-3] 模型加载完成 ✅")

    def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> dict:
        """生成文本并记录性能指标"""
        params = SamplingParams(
            temperature=temperature,
            max_tokens=max_tokens,
        )

        start = time.perf_counter()
        outputs = self.llm.generate([prompt], params)
        elapsed = time.perf_counter() - start

        output = outputs[0].outputs[0]
        num_tokens = len(output.token_ids)

        # 更新指标
        self.metrics.total_tokens += num_tokens
        self.metrics.total_time += elapsed
        self.metrics.requests += 1
        self.metrics.latencies.append(elapsed * 1000)

        return {
            "text": output.text,
            "tokens": num_tokens,
            "latency_ms": round(elapsed * 1000, 1),
            "tokens_per_second": round(num_tokens / elapsed, 1),
        }

    def auto_tune(self, prompts: list[str], target_acceptance: float = 0.85):
        """自动调优 num_speculative_tokens

        策略：从 5 开始逐步增加，直到接受率低于目标值
        """
        print("[EAGLE-3] 开始自动调优...")

        best_config = 5
        for num_tokens in [5, 6, 8, 10, 12]:
            # 重新配置
            self.llm.llm_engine.model_config.num_speculative_tokens = num_tokens

            # 测试 10 个 prompt
            test_metrics = PerfMetrics()
            for prompt in prompts[:10]:
                start = time.perf_counter()
                self.llm.generate([prompt], SamplingParams(max_tokens=100))
                elapsed = time.perf_counter() - start
                test_metrics.total_tokens += 100
                test_metrics.total_time += elapsed

            tps = test_metrics.tokens_per_second
            print(f"  num_spec_tokens={num_tokens}: {tps:.1f} tokens/s")

            if tps > best_config:
                best_config = num_tokens

        print(f"[EAGLE-3] 最优配置: num_speculative_tokens={best_config}")
        return best_config


# ─── 使用示例 ───

if __name__ == "__main__":
    config = EAGLE3Config(
        model="meta-llama/Llama-3-70B-Instruct",
        speculator="eagle3-llama3-70b-spec",
        num_speculative_tokens=8,
        tensor_parallel_size=4,
    )

    server = EAGLE3Server(config)

    # 测试生成
    result = server.generate("用 Python 实现快速排序算法：")
    print(f"\n生成结果 ({result['tokens']} tokens, {result['tokens_per_second']} tok/s):")
    print(result['text'][:200])

    # 性能报告
    print(f"\n性能报告:")
    print(server.metrics.report())

#!/bin/bash
# EAGLE-3 vs 标准解码 性能对比

echo "=== 测试 1: 标准自回归解码 (baseline) ==="
python -c "
from vllm import LLM, SamplingParams
import time

llm = LLM(model='meta-llama/Llama-3-70B-Instruct', tensor_parallel_size=4)
prompts = ['解释量子计算的基本原理', '用 Python 实现二叉树遍历', '写一首关于春天的诗'] * 10

start = time.perf_counter()
outputs = llm.generate(prompts, SamplingParams(max_tokens=256))
elapsed = time.perf_counter() - start

total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
print(f'标准解码: {total_tokens} tokens in {elapsed:.1f}s = {total_tokens/elapsed:.1f} tok/s')
"

echo ""
echo "=== 测试 2: EAGLE-3 推测解码 ==="
python -c "
from vllm import LLM, SamplingParams
import time

llm = LLM(
    model='meta-llama/Llama-3-70B-Instruct',
    speculative_model='eagle3-llama3-70b-spec',
    num_speculative_tokens=8,
    tensor_parallel_size=4,
)
prompts = ['解释量子计算的基本原理', '用 Python 实现二叉树遍历', '写一首关于春天的诗'] * 10

start = time.perf_counter()
outputs = llm.generate(prompts, SamplingParams(max_tokens=256))
elapsed = time.perf_counter() - start

total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
print(f'EAGLE-3: {total_tokens} tokens in {elapsed:.1f}s = {total_tokens/elapsed:.1f} tok/s')
"

echo ""
echo "=== 预期结果 ==="
echo "标准解码: ~800 tok/s (70B, 4x A100)"
echo "EAGLE-3:  ~5000+ tok/s (6-7x 加速)"

"""
EAGLE-3 领域 Speculator 训练脚本
用于在特定领域数据上训练自定义 draft model
依赖: pip install eagle3-specdecode torch transformers
"""
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
from dataclasses import dataclass
import json
import os

@dataclass
class TrainingConfig:
    """训练配置"""
    target_model: str = "meta-llama/Llama-3-8B-Instruct"
    hidden_layer: int = -4  # 从倒数第 4 层提取 hidden states
    draft_num_layers: int = 2
    draft_hidden_dim: int = 1024  # target hidden_dim / 4
    learning_rate: float = 1e-4
    batch_size: int = 32
    num_epochs: int = 5
    max_seq_len: int = 512
    train_test_split: float = 0.9
    mixed_domain_ratio: float = 0.3  # 30% 通用数据

class DraftModel(nn.Module):
    """EAGLE-3 Draft Model (小型 Transformer)"""

    def __init__(self, target_hidden_dim: int, vocab_size: int,
                 draft_hidden_dim: int = 1024, num_layers: int = 2):
        super().__init__()
        self.projection = nn.Linear(target_hidden_dim, draft_hidden_dim)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=draft_hidden_dim,
                nhead=8,
                dim_feedforward=draft_hidden_dim * 4,
                batch_first=True,
            ),
            num_layers=num_layers,
        )
        self.head = nn.Linear(draft_hidden_dim, vocab_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """
        Args:
            hidden_states: [batch, seq_len, target_hidden_dim]
        Returns:
            logits: [batch, seq_len, vocab_size]
        """
        x = self.projection(hidden_states)
        x = self.transformer(x)
        logits = self.head(x)
        return logits

class HiddenStateDataset(Dataset):
    """Hidden State 数据集"""

    def __init__(self, hidden_states: torch.Tensor, target_tokens: torch.Tensor):
        self.hidden_states = hidden_states
        self.target_tokens = target_tokens

    def __len__(self):
        return len(self.hidden_states)

    def __getitem__(self, idx):
        return self.hidden_states[idx], self.target_tokens[idx]

def extract_hidden_states(
    model: AutoModelForCausalLM,
    tokenizer: AutoTokenizer,
    texts: list[str],
    layer: int = -4,
    max_len: int = 512,
    batch_size: int = 8,
) -> tuple[torch.Tensor, torch.Tensor]:
    """从 target model 提取 hidden states 和 target tokens"""
    all_hidden = []
    all_tokens = []

    model.eval()
    with torch.no_grad():
        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i+batch_size]
            inputs = tokenizer(
                batch_texts,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=max_len,
            ).to(model.device)

            outputs = model(**inputs, output_hidden_states=True)
            # 提取指定层的 hidden states
            hidden = outputs.hidden_states[layer]  # [batch, seq, hidden_dim]
            # target tokens 是右移一位的 input_ids
            target = inputs["input_ids"][:, 1:]
            hidden = hidden[:, :-1, :]  # 对齐长度

            all_hidden.append(hidden.cpu())
            all_tokens.append(target.cpu())

    return torch.cat(all_hidden), torch.cat(all_tokens)

def train_domain_speculator(config: TrainingConfig, domain_texts: list[str]):
    """训练领域 speculator"""
    print(f"[训练] 加载 target model: {config.target_model}")
    tokenizer = AutoTokenizer.from_pretrained(config.target_model)
    model = AutoModelForCausalLM.from_pretrained(
        config.target_model,
        torch_dtype=torch.float16,
        device_map="auto",
    )

    # Step 1: 提取 hidden states
    print(f"[训练] 提取 hidden states from {len(domain_texts)} texts...")
    hidden_states, target_tokens = extract_hidden_states(
        model, tokenizer, domain_texts,
        layer=config.hidden_layer,
        max_len=config.max_seq_len,
    )
    print(f"[训练] 提取完成: {hidden_states.shape}")

    # Step 2: 创建数据集
    dataset = HiddenStateDataset(hidden_states, target_tokens)
    train_size = int(len(dataset) * config.train_test_split)
    train_set, val_set = torch.utils.data.random_split(dataset, [train_size, len(dataset) - train_size])

    train_loader = DataLoader(train_set, batch_size=config.batch_size, shuffle=True)
    val_loader = DataLoader(val_set, batch_size=config.batch_size)

    # Step 3: 初始化 draft model
    target_hidden_dim = hidden_states.shape[-1]
    vocab_size = model.config.vocab_size
    draft_model = DraftModel(
        target_hidden_dim=target_hidden_dim,
        vocab_size=vocab_size,
        draft_hidden_dim=config.draft_hidden_dim,
        num_layers=config.draft_num_layers,
    ).cuda()

    optimizer = torch.optim.AdamW(draft_model.parameters(), lr=config.learning_rate)
    criterion = nn.CrossEntropyLoss()

    # Step 4: 训练循环
    print(f"[训练] 开始训练 ({config.num_epochs} epochs)...")
    best_val_loss = float('inf')

    for epoch in range(config.num_epochs):
        draft_model.train()
        total_loss = 0
        correct = 0
        total = 0

        for batch_hidden, batch_target in train_loader:
            batch_hidden = batch_hidden.cuda()
            batch_target = batch_target.cuda()

            logits = draft_model(batch_hidden)
            loss = criterion(logits.reshape(-1, vocab_size), batch_target.reshape(-1))

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            total_loss += loss.item()
            correct += (logits.argmax(-1) == batch_target).sum().item()
            total += batch_target.numel()

        train_acc = correct / total
        avg_loss = total_loss / len(train_loader)

        # 验证
        draft_model.eval()
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for batch_hidden, batch_target in val_loader:
                logits = draft_model(batch_hidden.cuda())
                val_correct += (logits.argmax(-1) == batch_target.cuda()).sum().item()
                val_total += batch_target.numel()

        val_acc = val_correct / val_total
        print(f"  Epoch {epoch+1}: loss={avg_loss:.4f} train_acc={train_acc:.1%} val_acc={val_acc:.1%}")

        # 早停
        if avg_loss < best_val_loss:
            best_val_loss = avg_loss
            torch.save(draft_model.state_dict(), "eagle3_domain_speculator.pt")
            print(f"  ✅ 保存最优模型 (val_acc={val_acc:.1%})")

    print("[训练] 完成！模型已保存到 eagle3_domain_speculator.pt")

# ─── 使用示例 ───

if __name__ == "__main__":
    # 示例：法律领域 speculator 训练
    domain_texts = [
        "根据《民法典》第一百四十三条，具备下列条件的民事法律行为有效...",
        "本合同自双方签字盖章之日起生效，有效期为三年...",
        # ... 更多法律文本
    ]

    config = TrainingConfig(
        target_model="meta-llama/Llama-3-8B-Instruct",
        num_epochs=5,
        batch_size=16,
    )
    train_domain_speculator(config, domain_texts)