03 - 推理服务部署¶

# 基础推理示例
from vllm import LLM, SamplingParams

# 加载模型
llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    tensor_parallel_size=1,  # 单卡
    gpu_memory_utilization=0.9,  # GPU内存使用率
)

# 设置采样参数
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256,
)

# 生成
prompts = [
    "The future of AI is",
    "In the beginning,",
    "Once upon a time,",
]

outputs = llm.generate(prompts, sampling_params)

# 打印结果
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {generated_text!r}")
    print("-" * 50)

from vllm import LLM, SamplingParams

class VLLMDeployment:
    """
    vLLM高级部署配置
    """

    def __init__(self, model_path, config=None):
        self.config = config or {}

        # 初始化LLM
        self.llm = LLM(
            model=model_path,

            # 并行配置
            tensor_parallel_size=self.config.get('tensor_parallel_size', 1),
            pipeline_parallel_size=self.config.get('pipeline_parallel_size', 1),

            # 内存配置
            gpu_memory_utilization=self.config.get('gpu_memory_utilization', 0.9),
            max_num_seqs=self.config.get('max_num_seqs', 256),
            max_model_len=self.config.get('max_model_len', 4096),

            # 量化配置
            quantization=self.config.get('quantization', None),
            # 可选: 'awq', 'gptq', 'squeezellm'

            # 其他配置
            dtype=self.config.get('dtype', 'auto'),
            # 可选: 'float16', 'bfloat16', 'float32'

            trust_remote_code=self.config.get('trust_remote_code', True),
        )

    def generate(self, prompts, **sampling_kwargs):
        """
        批量生成
        """
        sampling_params = SamplingParams(
            temperature=sampling_kwargs.get('temperature', 0.7),
            top_p=sampling_kwargs.get('top_p', 0.9),
            top_k=sampling_kwargs.get('top_k', -1),
            max_tokens=sampling_kwargs.get('max_tokens', 256),
            presence_penalty=sampling_kwargs.get('presence_penalty', 0.0),
            frequency_penalty=sampling_kwargs.get('frequency_penalty', 0.0),
            stop=sampling_kwargs.get('stop', None),
        )

        outputs = self.llm.generate(prompts, sampling_params)

        # 格式化输出
        results = []
        for output in outputs:
            results.append({
                'prompt': output.prompt,
                'text': output.outputs[0].text,
                'tokens': len(output.outputs[0].token_ids),
            })

        return results

    def chat(self, messages, **kwargs):  # **kwargs收集关键字参数
        """
        对话模式
        """
        # 构建prompt
        prompt = self._build_chat_prompt(messages)

        # 生成
        result = self.generate([prompt], **kwargs)

        return result[0]['text']

    def _build_chat_prompt(self, messages):
        """
        构建对话prompt
        """
        # 使用模型特定的chat template
        if hasattr(self.llm.get_tokenizer(), 'apply_chat_template'):  # hasattr检查对象是否有某属性
            prompt = self.llm.get_tokenizer().apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
        else:
            # 简单的对话格式
            prompt = ""
            for msg in messages:
                role = msg['role']
                content = msg['content']
                prompt += f"{role}: {content}\n"
            prompt += "assistant:"

        return prompt

# 配置示例
VLLM_CONFIG = {
    'tensor_parallel_size': 2,  # 2卡并行
    'gpu_memory_utilization': 0.85,
    'max_num_seqs': 128,
    'max_model_len': 8192,
    'quantization': None,
    'dtype': 'bfloat16',
}

# 使用示例
deployment = VLLMDeployment("meta-llama/Llama-2-7b-hf", VLLM_CONFIG)

# 批量生成
results = deployment.generate(
    ["Hello, how are you?", "What is machine learning?"],
    temperature=0.7,
    max_tokens=100
)

# 对话模式
response = deployment.chat([
    {'role': 'user', 'content': 'Explain quantum computing'}
])

# 启动vLLM OpenAI兼容API服务
# 命令行方式（vLLM 0.6+推荐使用 vllm serve）
"""
vllm serve meta-llama/Llama-2-7b-hf \
    --tensor-parallel-size 2 \
    --gpu-memory-utilization 0.85 \
    --max-num-seqs 256 \
    --port 8000

# 旧版本也可使用：
# python -m vllm.entrypoints.openai.api_server \
#     --model meta-llama/Llama-2-7b-hf ...
"""

# 使用Python自定义服务
import uvicorn
from fastapi import FastAPI
from vllm import LLM, SamplingParams

app = FastAPI()

# 全局模型实例
llm_engine = None

@app.on_event("startup")
async def startup_event():  # async def定义协程函数
    global llm_engine
    llm_engine = LLM(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=2,
    )

@app.post("/v1/completions")
async def create_completion(request: CompletionRequest):
    """
    OpenAI兼容的completions API
    """
    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        top_p=request.top_p,
    )

    outputs = llm_engine.generate(request.prompt, sampling_params)

    return {
        "id": "cmpl-" + str(uuid.uuid4()),
        "object": "text_completion",
        "created": int(time.time()),
        "model": request.model,
        "choices": [
            {
                "text": output.outputs[0].text,
                "index": i,
                "logprobs": None,
                "finish_reason": "stop"
            }
            for i, output in enumerate(outputs)  # enumerate同时获取索引和元素
        ]
    }

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    """
    OpenAI兼容的chat completions API
    """
    # 构建prompt
    prompt = build_chat_prompt(request.messages)

    sampling_params = SamplingParams(
        temperature=request.temperature,
        max_tokens=request.max_tokens,
    )

    outputs = llm_engine.generate(prompt, sampling_params)

    return {
        "id": "chatcmpl-" + str(uuid.uuid4()),
        "object": "chat.completion",
        "created": int(time.time()),
        "model": request.model,
        "choices": [
            {
                "index": i,
                "message": {
                    "role": "assistant",
                    "content": output.outputs[0].text
                },
                "finish_reason": "stop"
            }
            for i, output in enumerate(outputs)
        ]
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

# Python客户端
from text_generation import Client

client = Client("http://localhost:8080")

# 文本生成
text = client.generate(
    "The future of AI is",
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.9,
).generated_text

print(text)

# 流式生成
for response in client.generate_stream(
    "Explain machine learning:",
    max_new_tokens=100
):
    print(response.token.text, end="", flush=True)

# 基础推理示例
import sglang as sgl

# 方式1：离线批量推理
llm = sgl.Engine(model_path="Qwen/Qwen2.5-7B-Instruct")

prompts = [
    "解释什么是大模型推理优化",
    "介绍 KV Cache 的作用"
]
outputs = llm.generate(
    prompts,
    sampling_params={"temperature": 0.7, "max_new_tokens": 256}
)
for out in outputs:
    print(out["text"])

llm.shutdown()

# 方式2：启动 OpenAI 兼容服务器
# 终端执行：
# python -m sglang.launch_server --model Qwen/Qwen2.5-7B-Instruct --port 8000

# 客户端调用
import openai
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="none")

response = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[{"role": "user", "content": "用JSON格式输出三个中国城市的信息"}],
    response_format={"type": "json_object"},
)
print(response.choices[0].message.content)

from sglang import Engine

class SGLangDeployment:
    """SGLang 高级部署配置"""

    def __init__(self, model_path, config=None):
        self.config = config or {}

        # 初始化 Engine
        self.llm = Engine(
            model_path=model_path,

            # 并行配置
            tensor_parallel_size=self.config.get('tensor_parallel_size', 1),
            pipeline_parallel_size=self.config.get('pipeline_parallel_size', 1),

            # 内存配置
            mem_fraction_static=self.config.get('mem_fraction_static', 0.9),
            max_running_requests=self.config.get('max_running_requests', 256),
            max_total_num_tokens=self.config.get('max_total_num_tokens', 4096),

            # 优化配置
            enable_torch_compile=self.config.get('enable_torch_compile', True),
            enable_mla=self.config.get('enable_mla', False),  # DeepSeek MLA 优化

            # 量化配置
            quantization=self.config.get('quantization', None),
            # 可选: 'fp8', 'fp8_kv_cache', 'awq', 'gptq'

            # 其他配置
            dtype=self.config.get('dtype', 'auto'),
            trust_remote_code=True,
        )

    def generate(self, prompts, **sampling_kwargs):
        """批量生成"""
        outputs = self.llm.generate(prompts, sampling_params=sampling_kwargs)
        return [out["text"] for out in outputs]

    def shutdown(self):
        """关闭引擎"""
        self.llm.shutdown()


# 使用示例
deployment = SGLangDeployment(
    model_path="meta-llama/Llama-3-8B-Instruct",
    config={
        "tensor_parallel_size": 2,
        "enable_torch_compile": True,
        "quantization": "fp8",
    }
)

results = deployment.generate(
    prompts=["Hello, how are you?"],
    temperature=0.7,
    max_new_tokens=256
)
print(results[0])

# TensorRT-LLM Python 高层 API 示例
from tensorrt_llm import LLM, BatchTokenizer
from tensorrt_llm.runtime import ModelConfig, SamplingConfig

# 方式1：使用高层 API（推荐）
llm = LLM(
    model_dir="/models/llama-7b",
    model_type="llama",
    tensor_parallel_size=2,  # 多卡并行
    dtype="float16",
)

# 批量推理
tokenizer = BatchTokenizer(model_dir="/models/llama-7b")

prompts = [
    "TensorRT-LLM is",
    "The future of AI is",
    "Large language models",
]

# 编码输入
input_ids = tokenizer.encode_batch(prompts)

# 执行推理
outputs = llm.generate(
    input_ids,
    max_new_tokens=100,
    temperature=0.7,
    top_p=0.9,
)

# 解码输出
for output in outputs:
    print(tokenizer.decode(output))

# 方式2：使用底层 API 进行自定义优化
from tensorrt_llm.builder import Builder
from tensorrt_llm.network import net
from tensorrt_llm.layers import AttentionLayer, MlpLayer

# 定义模型配置
model_config = ModelConfig(
    vocab_size=32000,
    hidden_size=4096,
    num_layers=32,
    num_heads=32,
    max_seq_len=4096,
    tensor_parallel_size=2,
)

# 构建 TensorRT 网络
with net(network_name="llama-7b"):
    # 嵌入层
    embedding = EmbeddingLayer(vocab_size=32000, hidden_size=4096)

    # Transformer 层
    for layer_idx in range(32):
        # 自注意力层
        attention = AttentionLayer(
            hidden_size=4096,
            num_heads=32,
            q_scaling=1.0,
        )

        # FFN 层（融合 MLP）
        ffn = MlpLayer(
            hidden_size=4096,
            ffn_hidden_size=11008,
            dtype="float16",
        )

    # 输出层
    lm_head = LinearLayer(vocab_size=32000, dtype="float16")

# 编译 TensorRT Engine
builder = Builder()
engine = builder.build_engine(net, model_config)
engine.save("llama-7b.engine")

# TensorRT-LLM 模型构建完整流程
from tensorrt_llm.builder import Builder
from tensorrt_llm.models import LLaMAModel
from tensorrt_llm.network import net

# 1. 下载并转换模型
# 使用 HuggingFace 模型
huggingface_model = "meta-llama/Llama-2-7b-hf"

# 转换模型格式
# python tensorrt_llm/models/convert.py \
#     --model_dir ${hf_model_dir} \
#     --output_dir ./llama-7b-trt \
#     --dtype float16 \
#     --tp_size 2

# 2. 构建 TensorRT Engine
from tensorrt_llm.builder import Builder

builder = Builder()

# 创建网络配置
network_config = {
    'model_name': 'llama',
    'vocab_size': 32000,
    'hidden_size': 4096,
    'num_layers': 32,
    'num_heads': 32,
    'num_kv_heads': 32,  # GQA 支持
    'hidden_act': 'silu',
    'max_seq_len': 4096,
    'dtype': 'float16',
}

# 构建引擎
engine = builder.build(
    network=net,
    config=network_config,
    optimization_level=3,  # 最高优化级别
)

# 保存引擎
engine.save('llama-7b-fp16.engine')

# 3. 量化引擎构建（INT8/FP8）
from tensorrt_llm.quantization import QuantMode

quant_config = {
    'quant_mode': QuantMode.INT8_WEIGHT_ONLY,  # INT8 权重量化
    'kv_cache_quant_mode': QuantMode.INT8,
}

quantized_engine = builder.build(
    network=net,
    config=quant_config,
    optimization_level=3,
)

quantized_engine.save('llama-7b-int8.engine')

# 使用AutoGPTQ进行量化
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

# 量化配置
quantize_config = BaseQuantizeConfig(
    bits=4,  # 4-bit量化
    group_size=128,  # 分组大小
    desc_act=False,  # 是否降序激活
)

# 加载并量化模型
model = AutoGPTQForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf",
    quantize_config,
)

# 准备校准数据
calib_data = [
    "auto-gptq is an easy-to-use model quantization library",
    "with user-friendly apis",
    # ... 更多校准数据
]

# 执行量化
model.quantize(calib_data)

# 保存量化模型
model.save_quantized("Llama-2-7b-4bit-gptq")

# 加载量化模型进行推理
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

model = AutoGPTQForCausalLM.from_quantized(
    "Llama-2-7b-4bit-gptq",
    device="cuda:0",
)

# 生成
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to("cuda")
outputs = model.generate(**inputs, max_new_tokens=50)
print(tokenizer.decode(outputs[0]))

# AWQ量化（通常比GPTQ质量更好）
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

# 加载模型
model_path = "meta-llama/Llama-2-7b-hf"
quant_path = "Llama-2-7b-awq"

# 量化配置
quant_config = {
    "zero_point": True,
    "q_group_size": 128,
    "w_bit": 4,
    "version": "GEMM"
}

# 加载模型
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 准备校准数据
examples = [
    tokenizer("auto-gptq is an easy-to-use model quantization library")
    for _ in range(8)
]

# 量化
model.quantize(
    tokenizer,
    quant_config=quant_config,
    calib_data=examples,
)

# 保存
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# 加载量化模型
model = AutoAWQForCausalLM.from_quantized(
    quant_path,
    fuse_layers=True,  # 融合层以加速
)

# vLLM支持多种量化方式

# 1. AWQ量化
llm = LLM(
    model="TheBloke/Llama-2-7B-AWQ",
    quantization="awq",
    tensor_parallel_size=1,
)

# 2. GPTQ量化
llm = LLM(
    model="TheBloke/Llama-2-7B-GPTQ",
    quantization="gptq",
    tensor_parallel_size=1,
)

# 3. SqueezeLLM量化
llm = LLM(
    model="path/to/squeezellm/model",
    quantization="squeezellm",
)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional
import uvicorn
from vllm import LLM, SamplingParams

app = FastAPI(title="LLM Inference API")

# 全局模型
llm = None

class CompletionRequest(BaseModel):  # Pydantic BaseModel：自动数据验证和序列化
    model: str
    prompt: str
    max_tokens: Optional[int] = 256  # Optional表示值可以为None
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    top_k: Optional[int] = -1
    stop: Optional[List[str]] = None
    stream: Optional[bool] = False

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[dict]
    max_tokens: Optional[int] = 256
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    stream: Optional[bool] = False

@app.on_event("startup")
async def load_model():
    global llm
    llm = LLM(
        model="meta-llama/Llama-2-7b-hf",
        tensor_parallel_size=1,
        gpu_memory_utilization=0.9,
    )
    print("Model loaded successfully")

@app.post("/v1/completions")
async def create_completion(request: CompletionRequest):
    try:  # try/except捕获异常，防止程序崩溃
        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            top_k=request.top_k,
            max_tokens=request.max_tokens,
            stop=request.stop,
        )

        outputs = llm.generate(request.prompt, sampling_params)

        return {
            "id": f"cmpl-{uuid.uuid4()}",
            "object": "text_completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "text": output.outputs[0].text,
                    "index": i,
                    "logprobs": None,
                    "finish_reason": "stop"
                }
                for i, output in enumerate(outputs)
            ],
            "usage": {
                "prompt_tokens": len(outputs[0].prompt_token_ids),
                "completion_tokens": len(outputs[0].outputs[0].token_ids),
                "total_tokens": len(outputs[0].prompt_token_ids) + len(outputs[0].outputs[0].token_ids)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def create_chat_completion(request: ChatCompletionRequest):
    try:
        # 构建chat prompt
        prompt = build_chat_prompt(request.messages)

        sampling_params = SamplingParams(
            temperature=request.temperature,
            top_p=request.top_p,
            max_tokens=request.max_tokens,
        )

        outputs = llm.generate(prompt, sampling_params)

        return {
            "id": f"chatcmpl-{uuid.uuid4()}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "index": i,
                    "message": {
                        "role": "assistant",
                        "content": output.outputs[0].text
                    },
                    "finish_reason": "stop"
                }
                for i, output in enumerate(outputs)
            ],
            "usage": {
                "prompt_tokens": len(outputs[0].prompt_token_ids),
                "completion_tokens": len(outputs[0].outputs[0].token_ids),
                "total_tokens": len(outputs[0].prompt_token_ids) + len(outputs[0].outputs[0].token_ids)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

def build_chat_prompt(messages):
    """构建chat prompt"""
    prompt = ""
    for msg in messages:
        role = msg.get('role', '')
        content = msg.get('content', '')
        if role == 'system':
            prompt += f"[INST] <<SYS>>\n{content}\n<</SYS>>\n\n"
        elif role == 'user':
            prompt += f"{content} [/INST]"
        elif role == 'assistant':
            prompt += f" {content} </s><s>[INST] "
    return prompt

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model_loaded": llm is not None}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)

import requests
import json

class LLMClient:
    """
    LLM API客户端
    """
    def __init__(self, base_url="http://localhost:8000"):
        self.base_url = base_url

    def complete(self, prompt, **kwargs):
        """
        文本补全
        """
        response = requests.post(
            f"{self.base_url}/v1/completions",
            json={
                "model": "llama-2-7b",
                "prompt": prompt,
                **kwargs
            }
        )
        return response.json()

    def chat(self, messages, **kwargs):
        """
        对话
        """
        response = requests.post(
            f"{self.base_url}/v1/chat/completions",
            json={
                "model": "llama-2-7b",
                "messages": messages,
                **kwargs
            }
        )
        return response.json()

# 使用示例
client = LLMClient()

# 文本补全
result = client.complete(
    "The future of AI is",
    max_tokens=100,
    temperature=0.7
)
print(result['choices'][0]['text'])

# 对话
result = client.chat([
    {"role": "user", "content": "Explain quantum computing"}
])
print(result['choices'][0]['message']['content'])

class PerformanceMonitor:
    """
    性能监控器
    """

    def __init__(self):
        self.metrics = {
            'latency': [],  # 延迟（秒）
            'throughput': [],  # 吞吐量（tokens/秒）
            'queue_size': [],  # 队列大小
        }

    def measure_latency(self, func, *args, **kwargs):  # func为传入的函数，*args/**kwargs透传任意参数，实现通用性能测量
        """
        测量延迟
        """
        import time

        start = time.time()
        result = func(*args, **kwargs)  # 将收集的参数解包后原样传给被测量的函数
        end = time.time()

        latency = end - start
        self.metrics['latency'].append(latency)

        return result, latency

    def calculate_throughput(self, num_tokens, latency):
        """
        计算吞吐量
        """
        throughput = num_tokens / latency
        self.metrics['throughput'].append(throughput)
        return throughput

    def get_statistics(self):
        """
        获取统计信息
        """
        import numpy as np

        stats = {}
        for key, values in self.metrics.items():
            if values:
                stats[key] = {
                    'mean': np.mean(values),
                    'median': np.median(values),
                    'p95': np.percentile(values, 95),
                    'p99': np.percentile(values, 99),
                    'min': np.min(values),
                    'max': np.max(values),
                }

        return stats

# 基准测试
def benchmark_inference(model, prompts, max_tokens=100):
    """
    推理基准测试
    """
    monitor = PerformanceMonitor()

    for prompt in prompts:
        # 测量延迟
        outputs, latency = monitor.measure_latency(
            model.generate,
            prompt,
            max_tokens=max_tokens
        )

        # 计算吞吐量
        num_tokens = len(outputs[0].outputs[0].token_ids)
        throughput = monitor.calculate_throughput(num_tokens, latency)

        print(f"Prompt: {prompt[:50]}...")
        print(f"Latency: {latency:.3f}s")
        print(f"Throughput: {throughput:.1f} tokens/s")
        print("-" * 50)

    # 打印统计
    stats = monitor.get_statistics()
    print("\n=== Benchmark Results ===")
    for metric, values in stats.items():
        print(f"\n{metric.upper()}:")
        for stat, value in values.items():
            print(f"  {stat}: {value:.3f}")

from vllm import LLM, SamplingParams

# 启用 Automatic Prefix Caching（APC）
llm = LLM(
    model="meta-llama/Llama-3-70B-Instruct",
    enable_prefix_caching=True,          # 关键参数
    gpu_memory_utilization=0.90,
    tensor_parallel_size=4,
    max_num_seqs=256,
    # 缓存预热：启动时预计算固定 System Prompt
    # preemption_mode="recompute",       # 显存不足时重算（默认）
    # preemption_mode="swap",            # 显存不足时换出到 CPU
)

# ---- 场景 1：多用户共享 System Prompt ----
SYSTEM_PROMPT = """
You are an expert AI assistant with deep knowledge in...
[此处填充 2000 tokens 的详细系统提示]
"""

def build_messages(user_question: str):
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user",   "content": user_question},
    ]

# 第一批请求：System Prompt 被计算并缓存
prompts_batch1 = [
    llm.get_tokenizer().apply_chat_template(
        build_messages(q), tokenize=False, add_generation_prompt=True
    )
    for q in ["What is RAG?", "Explain Attention", "How does RLHF work?"]
]

sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
# 首次生成：System Prompt 进行 Prefill 并缓存
outputs1 = llm.generate(prompts_batch1, sampling_params)

# 第二批请求：System Prompt KV Cache 命中，TTFT 大幅下降
prompts_batch2 = [
    llm.get_tokenizer().apply_chat_template(
        build_messages(q), tokenize=False, add_generation_prompt=True
    )
    for q in ["What is MoE?", "Explain vLLM", "How does LoRA work?"]
]
# 再次生成：System Prompt 命中缓存，只 Prefill 用户输入部分
outputs2 = llm.generate(prompts_batch2, sampling_params)


# ---- 场景 2：多轮对话的前缀复用 ----
class CachedMultiTurnChat:
    """多轮对话：历史轮次命中 Prefix Cache"""

    def __init__(self, llm_engine, system_prompt: str):
        self.llm = llm_engine
        self.system_prompt = system_prompt
        self.history = []

    def chat(self, user_input: str) -> str:
        self.history.append({"role": "user", "content": user_input})

        messages = [{"role": "system", "content": self.system_prompt}] + self.history
        prompt = self.llm.get_tokenizer().apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        # 每次请求，历史对话前缀命中缓存，只需计算新增 token
        output = self.llm.generate(
            [prompt],
            SamplingParams(temperature=0.7, max_tokens=256)
        )
        response = output[0].outputs[0].text
        self.history.append({"role": "assistant", "content": response})
        return response