LLM批处理推理实战：异步队列+动态批处理优化吞吐量

import json
import time
import redis
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

class BatchInferenceWorker:
    def __init__(self, model_path, redis_url="redis://localhost:6379",
                 max_batch_size=8, max_wait_ms=100):
        self.model = AutoModelForCausalLM.from_pretrained(
            model_path, torch_dtype=torch.float16, device_map="auto"
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.redis = redis.from_url(redis_url)
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms

    def collect_batch(self):
        """从队列收集一个批次的请求"""
        batch = []
        start = time.time()

        while len(batch) < self.max_batch_size:
            # 用BLPOP阻塞等待，避免忙轮询
            elapsed_ms = (time.time() - start) * 1000
            remaining = self.max_wait_ms - elapsed_ms
            if remaining <= 0:
                break

            result = self.redis.blpop("inference:queue", timeout=remaining / 1000)
            if result:
                _, data = result
                batch.append(json.loads(data))
            else:
                break  # 超时，用当前批次推理

        return batch

    def run_batch(self, batch):
        """批量推理"""
        prompts = [item["prompt"] for item in batch]
        inputs = self.tokenizer(prompts, return_tensors="pt", padding=True, truncation=True).to(self.model.device)

        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=512)

        results = []
        for i, output in enumerate(outputs):
            # 只取新生成的token
            new_tokens = output[inputs["input_ids"].shape[1]:]
            text = self.tokenizer.decode(new_tokens, skip_special_tokens=True)
            results.append({
                "task_id": batch[i]["task_id"],
                "result": text
            })

        return results

    def process_loop(self):
        """主循环"""
        while True:
            batch = self.collect_batch()
            if not batch:
                continue

            results = self.run_batch(batch)

            # 把结果写回Redis
            for r in results:
                self.redis.hset("inference:results", r["task_id"], json.dumps(r, ensure_ascii=False))
                self.redis.expire(f"inference:results:{r['task_id']}", 3600)

class DynamicBatcher:
    def __init__(self, min_batch=1, max_batch=16, max_wait_ms=200):
        self.min_batch = min_batch
        self.max_batch = max_batch
        self.max_wait_ms = max_wait_ms
        self.recent_latencies = []

    def get_batch_size(self, queue_depth):
        """根据队列深度和历史延迟动态调整批次大小"""
        if queue_depth > 100:
            return self.max_batch
        elif queue_depth > 20:
            return min(self.max_batch, queue_depth // 3)
        else:
            return self.min_batch

    def get_wait_time(self, queue_depth):
        """队列深时少等，队列浅时多等"""
        if queue_depth > 50:
            return 50  # 50ms
        elif queue_depth > 10:
            return 100
        else:
            return self.max_wait_ms

vllm serve Qwen/Qwen2.5-7B-Instruct \
    --max-num-seqs 16 \
    --max-num-batched-tokens 8192 \
    --gpu-memory-utilization 0.9

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1")

# 多个请求并发提交，vLLM自动做批处理
import asyncio
from concurrent.futures import ThreadPoolExecutor

def analyze_one(text):
    response = client.chat.completions.create(
        model="Qwen/Qwen2.5-7B-Instruct",
        messages=[
            {"role": "system", "content": "分析用户反馈的情感，返回JSON格式。"},
            {"role": "user", "content": text}
        ],
        temperature=0
    )
    return response.choices[0].message.content

# 并发提交50万条
with ThreadPoolExecutor(max_workers=32) as pool:
    results = list(pool.map(analyze_one, feedback_texts))