【vLLM 学习】Simple Profiling

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

更多 vLLM 中文文档及教程可访问 →https://hyper.ai/notebooks/37048?utm_source=Distribute&utm_medium=Distribute-vLLM&utm_campaign=Distribute-vLLM-260126

*在线运行 vLLM 入门教程：零基础分步指南

源码 examples/offline_inference/simple_profiling.py

# SPDX-License-Identifier: Apache-2.0

import os
import time

from vllm import LLM, SamplingParams

# 启用 torch 分析器，也可以在命令行设置
os.environ["VLLM_TORCH_PROFILER_DIR"] = "./vllm_profile"

# 样本提示。
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]
# 创建一个采样参数对象。
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)

if __name__ == "__main__":

    # 创建一个 LLM。
    llm = LLM(model="facebook/opt-125m", tensor_parallel_size=1)

    llm.start_profile()

    # 从提示中生成文本。输出是 RequestOutput 的包含提示，生成文本和其他信息的对象列表。
    outputs = llm.generate(prompts, sampling_params)

    llm.stop_profile()

    # 打印输出。
    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

    # 添加一个缓冲区，在后台过程中等待 profiling(如果 MP 为 ON) 完成分析输出。
    time.sleep(10)