推测解码#

SGLang 现在提供基于 EAGLE(EAGLE-2/EAGLE-3)的推测解码选项。我们的实现旨在最大化速度和效率,被认为是开源 LLM 引擎中最快的之一。

性能亮点#

请参见以下在 MT bench 上测试的 LLaMA-Instruct 3.1 8B 通过 EAGLE3 解码可实现的吞吐量巨大提升。 有关更多详细信息,请参阅 EAGLE3 论文

方法

吞吐量 (tokens/s)

SGLang(无推测,1x H100)

158.34 tokens/s

SGLang + EAGLE-2(1x H100)

244.10 tokens/s

SGLang + EAGLE-3(1x H100)

373.25 tokens/s

EAGLE 解码#

要启用 EAGLE 推测解码,以下参数是相关的:

  • speculative_draft_model_path:指定草稿模型。此参数是必需的。

  • speculative_num_steps:自回归草稿的深度。增加推测范围但存在拒绝级联的风险。默认值为 5。

  • speculative_eagle_topk:每步的分支因子。提高候选多样性,将导致更高的接受率,但也会导致更高的内存/计算消耗。默认值为 4。

  • speculative_num_draft_tokens:最大并行验证能力。允许更深的树评估,但会导致更高的 GPU 内存使用。默认值为 8。

这些参数对于 EAGLE-2 和 EAGLE-3 是相同的。

您可以使用 bench_speculative.py 找到这些参数的最佳组合。

在下文文档中,我们设置了较小的 --cuda-graph-max-bs 值以实现更快的引擎启动。对于您自己的工作负载,请将上述参数与 --cuda-graph-max-bs--max-running-requests--mem-fraction-static 一起调整以获得最佳性能。

EAGLE-2 解码#

您可以通过设置 --speculative-algorithm EAGLE 并选择合适的模型来启用 EAGLE-2 解码。

[ ]:
from sglang.test.doc_patch import launch_server_cmd
from sglang.utils import wait_for_server, print_highlight, terminate_process

import openai
[ ]:
server_process, port = launch_server_cmd(
    """
python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf  --speculative-algorithm EAGLE \
    --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 3 \
    --speculative-eagle-topk 4 --speculative-num-draft-tokens 16 --cuda-graph-max-bs 8 --log-level warning
"""
)

wait_for_server(f"http://localhost:{port}")
[ ]:
client = openai.Client(base_url=f"http://127.0.0.1:{port}/v1", api_key="None")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-chat-hf",
    messages=[
        {"role": "user", "content": "列出3个国家及其首都。"},
    ],
    temperature=0,
    max_tokens=64,
)

print_highlight(f"响应: {response}")
[ ]:
terminate_process(server_process)

使用 torch.compile 的 EAGLE-2 解码#

您还可以启用 torch.compile 进行进一步优化,并可选设置 --torch-compile-max-bs

[ ]:
server_process, port = launch_server_cmd(
    """
python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf  --speculative-algorithm EAGLE \
    --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 5 \
        --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --mem-fraction 0.6 \
            --enable-torch-compile --torch-compile-max-bs 2 --log-level warning
"""
)

wait_for_server(f"http://localhost:{port}")
[ ]:
client = openai.Client(base_url=f"http://127.0.0.1:{port}/v1", api_key="None")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-chat-hf",
    messages=[
        {"role": "user", "content": "列出3个国家及其首都。"},
    ],
    temperature=0,
    max_tokens=64,
)

print_highlight(f"响应: {response}")
[ ]:
terminate_process(server_process)

基于频率排名推测采样的 EAGLE-2 解码#

通过在草稿模型中使用截断的高频词元词汇表,Eagle 推测解码减少了 lm_head 的计算开销,同时在不降低质量的情况下加速了流水线。有关更多详细信息,请查阅论文

在我们的实现中,设置 --speculative-token-map 以启用此优化。您可以从这个模型中获取 FR-Spec 的高频词元。或者,您可以从此仓库直接下载这些词元。

感谢 Weilin ZhaoZhousx 的贡献。

[ ]:
server_process, port = launch_server_cmd(
    """
python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algorithm EAGLE \
    --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \
    --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --speculative-token-map thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt \
    --mem-fraction 0.7 --cuda-graph-max-bs 2 --dtype float16  --log-level warning
"""
)

wait_for_server(f"http://localhost:{port}")
[ ]:
client = openai.Client(base_url=f"http://127.0.0.1:{port}/v1", api_key="None")

response = client.chat.completions.create(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    messages=[
        {"role": "user", "content": "列出3个国家及其首都。"},
    ],
    temperature=0,
    max_tokens=64,
)

print_highlight(f"响应: {response}")
[ ]:
terminate_process(server_process)

EAGLE-3 解码#

您可以通过设置 --speculative-algorithm EAGLE3 并选择合适的模型来启用 EAGLE-3 解码。

[ ]:
server_process, port = launch_server_cmd(
    """
python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct  --speculative-algorithm EAGLE3 \
    --speculative-draft-model-path jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B --speculative-num-steps 5 \
        --speculative-eagle-topk 8 --speculative-num-draft-tokens 32 --mem-fraction 0.6 \
        --cuda-graph-max-bs 2 --dtype float16 --log-level warning
"""
)

wait_for_server(f"http://localhost:{port}")
[ ]:
client = openai.Client(base_url=f"http://127.0.0.1:{port}/v1", api_key="None")

response = client.chat.completions.create(
    model="meta-llama/Meta-Llama-3.1-8B-Instruct",
    messages=[
        {"role": "user", "content": "列出3个国家及其首都。"},
    ],
    temperature=0,
    max_tokens=64,
)

print_highlight(f"响应: {response}")
[ ]:
terminate_process(server_process)

多词元预测#

我们通过使用推测解码在 SGLang 中支持MTP(Multi-Token Prediction)。我们在这里以 Xiaomi/MiMo-7B-RL 模型为例(deepseek mtp 使用参考 deepseek 文档)

[ ]:
server_process, port = launch_server_cmd(
    """
    python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-7B-RL --host 0.0.0.0 --trust-remote-code \
    --speculative-algorithm EAGLE --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 \
    --mem-fraction 0.5 --log-level warning
"""
)

wait_for_server(f"http://localhost:{port}")
[ ]:
import requests

url = f"http://localhost:{port}/v1/chat/completions"

data = {
    "model": "XiaomiMiMo/MiMo-7B-RL",
    "messages": [{"role": "user", "content": "法国的首都是什么?"}],
}

response = requests.post(url, json=data)
print_highlight(response.json())
[ ]:
terminate_process(server_process)

参考#

EAGLE 流程如下:

  • 在 EAGLE 中,草稿模型使用特征序列 \((f_1, ..., f_k)\) 和词元序列 \((t_2, ..., t_{k+1})\) 来预测下一个特征向量,即原始 LLM 的最后一个隐藏状态。

  • 然后从 \(p_{k+2}=\text{LMHead}(f_{k+1})\) 中采样下一个词元。之后,这两个序列以树状方式扩展——分支出多个可能的连续,每步的分支因子由 speculative_eagle_topk 参数控制——以确保上下文更连贯的连接,并再次作为输入。

  • EAGLE-2 额外使用草稿模型评估草稿树中某些分支的可能性,动态停止不太可能分支的扩展。在扩展阶段后,使用重排(reranking)只选择前 speculative_num_draft_tokens 个最终节点作为草稿词元。

  • EAGLE-3 移除了特征预测目标, incorporates 低层和中间层特征,并以在线方式(on-policy)进行训练。

通过在特征而非词元上操作,以及通过传递下一个时间步的词元来最小化采样的随机影响,这提高了草稿的准确性。此外,动态调整草稿树和选择重排的最终节点进一步提高了草稿词元的接受率。有关更多详细信息,请参阅 EAGLE-2EAGLE-3 论文。

有关如何训练您自己的 EAGLE 模型的指导,请参阅 EAGLE 仓库