vLLM 高性能推理引擎

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-7b-hf",       # 模型名称或本地路径（必填）
    tokenizer=None,                          # 分词器路径，默认与model相同
    tokenizer_mode="auto",                   # 分词器模式："auto"或"slow"
    tokenizer_revision=None,                 # 分词器的HuggingFace revision
    trust_remote_code=False,                 # 是否信任远程代码（自定义模型需要）
    tensor_parallel_size=1,                  # 张量并行数（GPU数量）
    dtype="auto",                            # 数据类型："auto","float16","bfloat16","float32"
    quantization=None,                       # 量化方式："awq","gptq","squeezellm","fp8"等
    gpu_memory_utilization=0.9,              # GPU显存利用率（0~1），预留部分给其他操作
    swap_space=4,                            # CPU交换空间大小（GB），用于KV cache换出
    enforce_eager=False,                     # 是否强制使用eager模式（禁用CUDA Graph）
    max_seq_len_to_capture=8192,            # CUDA Graph捕获的最大序列长度
    max_model_len=None,                      # 模型最大序列长度，默认从模型配置读取
    speculative_model=None,                  # 推测解码的草稿模型
    num_speculative_tokens=None,             # 推测解码的草稿token数
    enable_prefix_caching=False,             # 是否启用前缀缓存
    enable_lora=False,                       # 是否启用LoRA
    max_lora_rank=0,                         # LoRA最大秩
    max_num_seqs=256,                        # 最大并发序列数
    max_num_batched_tokens=None,             # 每次迭代最大批处理token数
    revision=None,                           # 模型的HuggingFace revision
    code_revision=None,                      # 代码的HuggingFace revision
    seed=0,                                  # 随机种子
)

outputs = llm.generate(
    prompts,                    # 输入提示，可以是字符串列表或TokensPrompt列表
    sampling_params=None,       # SamplingParams对象，控制生成行为
    use_tqdm=True,             # 是否显示进度条
    lora_request=None,         # LoRA请求（需启用enable_lora）
)

for output in outputs:
    prompt = output.prompt           # 原始输入提示
    generated_text = output.outputs[0].text        # 生成的文本
    token_ids = output.outputs[0].token_ids        # 生成的token ID列表
    finish_reason = output.outputs[0].finish_reason  # 结束原因："stop"或"length"
    cumulative_logprob = output.outputs[0].cumulative_logprob  # 累积对数概率

from vllm import LLM, SamplingParams

# 1. 初始化模型
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=2,          # 使用2块GPU进行张量并行
    gpu_memory_utilization=0.9,      # 使用90%的GPU显存
    enable_prefix_caching=True,      # 启用前缀缓存
)

# 2. 设置采样参数
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=512,
)

# 3. 批量推理
prompts = [
    "请解释什么是机器学习？",
    "Python中的装饰器是什么？",
    "量子计算的基本原理是什么？",
]
outputs = llm.generate(prompts, sampling_params)

# 4. 输出结果
for output in outputs:
    print(f"提示: {output.prompt!r}")
    print(f"生成: {output.outputs[0].text!r}")
    print("---")

from vllm import SamplingParams

params = SamplingParams(
    n=1,                          # 每个提示生成的补全数量
    best_of=None,                 # 从n个候选中选择最佳的个数（需n<best_of）
    presence_penalty=0.0,         # 存在惩罚：正值降低重复token的概率
    frequency_penalty=0.0,        # 频率惩罚：基于token出现频率的惩罚
    repetition_penalty=1.0,       # 重复惩罚：>1时惩罚重复token
    temperature=1.0,              # 温度：控制随机性，0为贪心解码
    top_p=1.0,                    # Top-p（核采样）：0.1表示只考虑概率前10%的token
    top_k=-1,                     # Top-k：只从概率最高的k个token中采样，-1表示不限制
    min_p=0.0,                    # Min-p：过滤掉概率低于最高概率token的min_p倍的token
    seed=None,                    # 采样随机种子
    use_beam_search=False,        # 是否使用束搜索
    length_penalty=1.0,           # 长度惩罚：束搜索时对长度的惩罚
    early_terminating=None,       # 早停策略："high"或"low"
    stop=None,                    # 停止词列表，遇到这些字符串时停止生成
    stop_token_ids=None,          # 停止token ID列表
    include_stop_str_in_output=False,  # 是否在输出中包含停止字符串
    ignore_eos=False,             # 是否忽略EOS token
    max_tokens=16,                # 最大生成token数
    min_tokens=0,                 # 最小生成token数
    logprobs=None,                # 返回top-N个token的对数概率
    prompt_logprobs=None,         # 返回提示token的对数概率
    skip_special_tokens=True,     # 是否跳过特殊token
    spaces_between_special_tokens=True,  # 特殊token之间是否加空格
    logits_processors=None,       # 自定义logits处理器列表
    truncate_prompt_tokens=None,  # 截断提示到指定token数
)

from vllm import SamplingParams

# 贪心解码（确定性输出）
greedy_params = SamplingParams(temperature=0, max_tokens=256)

# 创意写作（高随机性）
creative_params = SamplingParams(
    temperature=1.2,
    top_p=0.95,
    top_k=100,
    max_tokens=1024,
    repetition_penalty=1.1,
)

# 代码生成（低随机性）
code_params = SamplingParams(
    temperature=0.2,
    top_p=0.9,
    max_tokens=512,
    stop=["\n\n", "def ", "class "],  # 遇到这些模式时停止
)

# 多候选生成
multi_params = SamplingParams(
    n=3,                  # 生成3个候选
    temperature=0.8,
    top_p=0.95,
    max_tokens=256,
)

import openai

client = openai.OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="not-needed",  # vLLM默认不需要API key
)

# 文本补全
response = client.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    prompt="请解释什么是深度学习：",
    max_tokens=256,
    temperature=0.7,
    top_p=0.9,
)
print(response.choices[0].text)

# 多轮对话
response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[
        {"role": "system", "content": "你是一个有帮助的AI助手。"},
        {"role": "user", "content": "请解释什么是Transformer架构？"},
    ],
    max_tokens=512,
    temperature=0.7,
)
print(response.choices[0].message.content)

# 流式输出
stream = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[
        {"role": "user", "content": "写一首关于春天的诗"},
    ],
    max_tokens=256,
    stream=True,  # 启用流式输出
)
for chunk in stream:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-hf")
params = SamplingParams(temperature=0, max_tokens=128)

# 读取数据集并批量推理
import json
with open("dataset.jsonl") as f:
    prompts = [json.loads(line)["prompt"] for line in f]

outputs = llm.generate(prompts, params)

# 保存结果
results = []
for output in outputs:
    results.append({
        "prompt": output.prompt,
        "completion": output.outputs[0].text,
    })

with open("results.jsonl", "w") as f:
    for r in results:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")

# 终端1：启动服务器
# vllm serve meta-llama/Llama-2-7b-hf --port 8000

# 终端2：客户端调用
import openai
import time

client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")

# 实时请求
start = time.time()
response = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[{"role": "user", "content": "Hello!"}],
    max_tokens=64,
)
latency = time.time() - start
print(f"响应: {response.choices[0].message.content}")
print(f"延迟: {latency:.2f}s")

from vllm import LLM, SamplingParams

llm = LLM(model="meta-llama/Llama-2-7b-chat-hf", enable_prefix_caching=True)
params = SamplingParams(temperature=0.3, max_tokens=512)

# 共享的系统提示会通过前缀缓存复用KV cache
system_prompt = "你是一个专业的问答助手。根据以下上下文回答问题，如果上下文中没有答案，请说明。"

def rag_query(question: str, context: str) -> str:
    prompt = f"[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n上下文：{context}\n\n问题：{question} [/INST]"
    outputs = llm.generate([prompt], params)
    return outputs[0].outputs[0].text

# 多个查询共享system_prompt的KV cache
answer1 = rag_query("什么是BERT？", "BERT是Google提出的预训练语言模型...")
answer2 = rag_query("什么是GPT？", "GPT是OpenAI提出的生成式预训练模型...")

from vllm import LLM, SamplingParams

# 使用张量并行加载大模型
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,          # 4卡并行
    gpu_memory_utilization=0.92,
    quantization="awq",              # 使用AWQ量化降低显存
)

# 不同任务的采样参数
factual_params = SamplingParams(temperature=0, max_tokens=256)
creative_params = SamplingParams(temperature=1.0, top_p=0.95, max_tokens=1024)

# 事实性问答
factual_prompts = ["法国的首都是什么？", "水的化学式是什么？"]
factual_outputs = llm.generate(factual_prompts, factual_params)

# 创意写作
creative_prompts = ["写一篇关于人工智能未来的科幻短文"]
creative_outputs = llm.generate(creative_prompts, creative_params)

# 服务端启动：
# vllm serve meta-llama/Llama-2-7b-chat-hf --enable-prefix-caching

# 客户端流式调用
import openai

client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")

def chat_stream(message: str):
    """流式聊天函数，实时输出模型响应"""
    stream = client.chat.completions.create(
        model="meta-llama/Llama-2-7b-chat-hf",
        messages=[
            {"role": "system", "content": "你是一个有帮助的AI助手。"},
            {"role": "user", "content": message},
        ],
        max_tokens=512,
        temperature=0.7,
        stream=True,
    )
    full_response = ""
    for chunk in stream:
        if chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            full_response += content
            print(content, end="", flush=True)
    print()
    return full_response

chat_stream("请详细解释Transformer中的自注意力机制")

# 启动支持LoRA的服务器
# vllm serve meta-llama/Llama-2-7b-hf \
#     --enable-lora \
#     --max-lora-rank 16

# 客户端使用LoRA
import openai
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="not-needed")

# 使用基础模型
response_base = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[{"role": "user", "content": "写一首诗"}],
    max_tokens=128,
)

# 使用LoRA适配器（假设已加载）
response_lora = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[{"role": "user", "content": "写一首诗"}],
    max_tokens=128,
    extra_body={"lora_name": "poetry-lora"},
)

# 伪代码：每次迭代的调度逻辑
def schedule():
    # 1. 优先处理running队列（继续解码）
    for seq in running:
        if gpu_has_room(seq):
            allocate_new_block(seq)
        else:
            swap_to_cpu(seq)  # KV cache换出到CPU

    # 2. 尝试从swapped队列换回
    for seq in swapped:
        if gpu_has_room(seq):
            swap_from_cpu(seq)

    # 3. 从waiting队列调度新请求
    for seq in waiting:
        if gpu_has_room(seq):
            preempt_if_needed()  # 抢占低优先级序列
            add_to_running(seq)

# BlockManager 核心逻辑（简化版）
class BlockTableManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size     # 每个block的token数（默认16）
        self.free_blocks = list(range(num_blocks))  # 空闲block列表
        self.block_tables = {}           # seq_id -> BlockTable

    def allocate(self, seq_id, num_tokens):
        """为序列分配block"""
        num_blocks_needed = ceil(num_tokens / self.block_size)
        if len(self.free_blocks) < num_blocks_needed:
            return False  # 显存不足
        blocks = [self.free_blocks.pop() for _ in range(num_blocks_needed)]
        self.block_tables[seq_id].extend(blocks)
        return True

    def free(self, seq_id):
        """释放序列的所有block"""
        for block in self.block_tables[seq_id]:
            self.free_blocks.append(block)
        del self.block_tables[seq_id]

# 推荐设置
llm = LLM(
    model="model-name",
    gpu_memory_utilization=0.9,    # 留10%给CUDA和其他操作
    swap_space=4,                   # 设置适当的CPU交换空间
)

# 避免设置过高
# gpu_memory_utilization=0.99  # 危险！可能导致OOM

# 根据模型大小选择GPU数量
# 7B模型: 1张GPU (80GB) 或 2张GPU (24GB each)
# 13B模型: 2张GPU
# 70B模型: 4-8张GPU

llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,
    quantization="awq",  # 量化可减少GPU数量需求
)

# 适用场景：多个请求共享相同系统提示
llm = LLM(
    model="model-name",
    enable_prefix_caching=True,  # RAG、Agent等场景强烈推荐
)

# 不适用场景：每个请求的提示完全不同时，前缀缓存无效果
# 此时启用反而会带来额外的哈希计算开销

# 1. 需要确定性输出时
params = SamplingParams(temperature=0)  # 贪心解码，结果可复现

# 2. 需要多样性时，避免temperature和top_p同时极端
params = SamplingParams(temperature=0.7, top_p=0.9)  # 推荐组合

# 3. 避免重复输出
params = SamplingParams(
    temperature=0.7,
    frequency_penalty=0.3,    # 轻度频率惩罚
    repetition_penalty=1.1,   # 轻度重复惩罚
)

# 4. 控制输出长度
params = SamplingParams(
    max_tokens=256,
    min_tokens=50,            # 避免过短的无意义输出
    stop=["\n\n"],            # 合理设置停止条件
)

# AWQ量化：4-bit量化，精度损失小
llm = LLM(model="TheBloke/Llama-2-7B-AWQ", quantization="awq")

# GPTQ量化：另一种4-bit量化方式
llm = LLM(model="TheBloke/Llama-2-7B-GPTQ", quantization="gptq")

# FP8量化：需要H100/AMD MI300等支持FP8的GPU
llm = LLM(model="model-name", quantization="fp8")