你应该懂的AI大模型十三 之 推理框架在前面的系列文章中我们已经讨论了模型架构、训练技巧、微调方法等。但真正让大模型从“实验室玩具”变成“生产工具”的关键一环是推理框架。想象一下你训练了一个千亿参数的模型但用户点击“发送”后却要等30秒才能看到第一个字——这体验是灾难性的。推理框架就是专门解决这类问题的“性能引擎”。### 什么是推理框架简单来说推理框架是部署大模型并执行生成请求的软件系统。它负责加载模型权重、管理GPU显存、优化计算图、调度请求最终以最低延迟和最高吞吐量输出结果。它不负责训练只负责“跑起来”。为什么不能直接用PyTorch呢因为PyTorch的推理模式存在几个痛点-显存浪费每个请求都要复制一份模型权重。-计算低效没有对KV Cache、Attention进行专门优化。-无批处理不同请求无法共享计算。推理框架的核心目标就是解决这三个问题。我们按从易到难的顺序逐个击破。—## 基础从零实现一个简单的批处理推理首先我们理解一下最朴素的方式——单请求推理。此时模型一次只处理一个用户输入。代码如下pythonimport torchfrom transformers import AutoModelForCausalLM, AutoTokenizer# 加载一个小模型如GPT-2model_name gpt2tokenizer AutoTokenizer.from_pretrained(model_name)model AutoModelForCausalLM.from_pretrained(model_name).to(cuda)# 单请求生成def generate_single(prompt: str, max_new_tokens: int 50): inputs tokenizer(prompt, return_tensorspt).to(cuda) outputs model.generate(**inputs, max_new_tokensmax_new_tokens) return tokenizer.decode(outputs[0], skip_special_tokensTrue)print(generate_single(The capital of France is))问题如果同时有10个用户请求这个代码只能串行处理GPU利用率极低。我们改进一下引入动态批处理Dynamic Batching——把多个请求拼成一个批次共享一次前向计算。pythondef generate_batch(prompts: list[str], max_new_tokens: int 50): # 将多个prompt编码成同一batch注意padding inputs tokenizer(prompts, return_tensorspt, paddingTrue, truncationTrue).to(cuda) outputs model.generate(**inputs, max_new_tokensmax_new_tokens) return [tokenizer.decode(out, skip_special_tokensTrue) for out in outputs]# 模拟4个并发请求prompts [ What is the weather in Beijing?, Explain quantum computing briefly., Write a haiku about autumn., List three benefits of exercise.]results generate_batch(prompts)for r in results: print(r[:50] ...)这一步已经比单请求快了不少因为GPU并行处理多个序列。但这只是“表面功夫”——真正的推理框架会做更多细致优化。—## 进阶KV Cache 与显存优化大模型生成时每生成一个新token都需要重新计算前面所有token的Key和Value向量用于Attention。如果不缓存复杂度是O(n²)。KV Cache就是将这些中间结果缓存下来把复杂度降为O(n)。下面我们手写一个简化版的KV Cache实现pythondef inference_with_kv_cache(model, tokenizer, prompt, max_new_tokens30): inputs tokenizer(prompt, return_tensorspt).to(cuda) input_ids inputs[input_ids] past_key_values None # 初始无缓存 generated input_ids.tolist()[0] for _ in range(max_new_tokens): # 只输入最后一个token配合past_key_values out model(input_idsinput_ids, past_key_valuespast_key_values, use_cacheTrue) logits out.logits[:, -1, :] # 取最后一个位置的logits next_token torch.argmax(logits, dim-1).unsqueeze(0) generated.append(next_token.item()) # 更新past_key_values为模型返回的缓存 past_key_values out.past_key_values input_ids next_token # 下一轮只输入新token return tokenizer.decode(generated, skip_special_tokensTrue)显存优化方面推理框架常用的技术包括-模型量化将FP16降到INT8/INT4把显存占用缩至1/4甚至1/8。-连续批处理不再等待整个batch完成而是动态插入新请求移除已完成的请求。-PagedAttention如vLLM像操作系统分页一样管理KV Cache减少碎片浪费。—## 高级主流推理框架实战对比现在我们介绍两个最主流的开源推理框架vLLM和TensorRT-LLM。### vLLM易用性之王vLLM基于PagedAttention在吞吐量上比HuggingFace Transformers快数倍且接口极简。python# 安装pip install vllmfrom vllm import LLM, SamplingParams# 加载模型llm LLM(modelmeta-llama/Llama-2-7b-chat-hf, dtypefloat16)# 设置采样参数sampling_params SamplingParams(temperature0.8, max_tokens100)# 批量推理prompts [ What is the meaning of life?, Give me a recipe for chocolate cake.,]outputs llm.generate(prompts, sampling_params)for output in outputs: print(output.outputs[0].text)vLLM的LLM类自动处理了批处理、缓存、量化等细节。你只需关注业务逻辑。它还支持OpenAI兼容的API服务可以直接替代/v1/completions接口。### TensorRT-LLM极致性能NVIDIA的TensorRT-LLM则更底层它把模型编译成TensorRT引擎在GPU上运行速度极快但配置复杂。它常用于需要超低延迟的生产环境如金融交易、实时对话。python# 安装pip install tensorrt_llm# 通常需要先导出ONNX模型再用trtllm-build命令编译# 伪代码示例import tensorrt_llm as tllmfrom tensorrt_llm.runtime import ModelRunner# 假设已经编译好了engine文件runner ModelRunner.from_dir( engine_dir/path/to/trt_engine, lora_dirNone, is_enc_decFalse, tensor_parallel_size1, use_gpt_attention_pluginTrue,)# 推理output_ids runner.generate([Tell me a joke.], max_new_tokens50)print(output_ids)对比总结| 框架 | 易用性 | 性能 | 适用场景 ||------|--------|------|----------|| vLLM | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐ | 快速部署、高吞吐API || TensorRT-LLM | ⭐⭐ | ⭐⭐⭐⭐⭐ | 极致延迟、定制生产 |—## 实战部署一个高并发聊天服务结合vLLM我们可以快速搭建一个能承受高并发的推理服务。下面是一个微型示例python# 使用FastAPI vLLMfrom fastapi import FastAPI, Requestfrom vllm import LLM, SamplingParamsimport uvicornapp FastAPI()llm LLM(modelmistralai/Mistral-7B-Instruct-v0.2, dtypefloat16)app.post(/chat)async def chat(request: Request): data await request.json() prompt data[prompt] params SamplingParams(max_tokens200, temperature0.7) result llm.generate([prompt], params) return {reply: result[0].outputs[0].text}if __name__ __main__: uvicorn.run(app, host0.0.0.0, port8000)运行这个服务后你可以用curl测试bashcurl -X POST http://localhost:8000/chat -H Content-Type: application/json -d {prompt:Explain AI in one sentence.}—## 总结推理框架是大模型落地的“最后一公里”。我们从最基础的批处理讲起理解了KV Cache和显存优化的必要性再对比了vLLM和TensorRT-LLM两大主流方案最后实现了一个可用的推理服务。核心要点1.动态批处理是提升吞吐量的第一级台阶。2.KV Cache和量化是降低延迟和显存的关键。3. 选择框架时vLLM适合快速开发TensorRT-LLM适合极致性能。4. 生产环境还需考虑并发调度、容错、模型热加载等。推理框架的技术仍在快速演进例如投机采样、并行解码等新方法不断涌现。掌握这些原理你就能根据业务需求灵活调优让大模型真正跑得又快又稳。希望这篇文章能帮你建立对推理框架的系统认知。下一期我们将深入探讨模型量化技术敬请期待