【vLLM 学习】Mlpspeculator

神经星星2025-08-27 10:56

vLLM 是一款专为大语言模型推理加速而设计的框架，实现了 KV 缓存内存几乎零浪费，解决了内存管理瓶颈问题。

更多 vLLM 中文文档及教程可访问 →vllm.hyper.ai/

源码 examples/offline_inference/mlpspeculator.py

ini 复制代码

# SPDX-License-Identifier: Apache-2.0

import gc
import time

from vllm import LLM, SamplingParams


def time_generation(llm: LLM, prompts: list[str],
                    sampling_params: SamplingParams):
    # 从提示中生成文本。输出是 RequestOutput 的包含提示，生成文本和其他信息的对象列表。
    # 首先预热
    llm.generate(prompts, sampling_params)
    llm.generate(prompts, sampling_params)
    start = time.time()
    outputs = llm.generate(prompts, sampling_params)
    end = time.time()
    print((end - start) / sum([len(o.outputs[0].token_ids) for o in outputs]))
    # 打印输出。
    for output in outputs:
        generated_text = output.outputs[0].text
        print(f"text: {generated_text!r}")


if __name__ == "__main__":

    template = (
        "Below is an instruction that describes a task. Write a response "
        "that appropriately completes the request.\n\n### Instruction:\n{}"
        "\n\n### Response:\n")

    # 样本提示。
    prompts = [
        "Write about the president of the United States.",
    ]
    prompts = [template.format(prompt) for prompt in prompts]
    # 创建一个采样参数对象。
    sampling_params = SamplingParams(temperature=0.0, max_tokens=200)

    # 创建一个不使用规格解码的 LLM
    llm = LLM(model="meta-llama/Llama-2-13b-chat-hf")

    print("Without speculation")
    time_generation(llm, prompts, sampling_params)

    del llm
    gc.collect()

    # 与规格解码创建一个 LLM
    llm = LLM(
        model="meta-llama/Llama-2-13b-chat-hf",
        speculative_config={
            "model": "ibm-ai-platform/llama-13b-accelerator",
        },
    )

    print("With speculation")
    time_generation(llm, prompts, sampling_params)