LLM 推理优化：加速与量化

LLM 推理优化加速与量化1. 技术分析1.1 LLM 推理挑战LLM 推理面临的主要挑战推理挑战 计算量大: O(n²d) 内存占用高: 参数 KV Cache 延迟要求: 实时应用需求1.2 推理优化方法方法原理加速比精度损失量化降低精度2-4x小蒸馏知识迁移1.5-2x小剪枝移除冗余1.5-3x中编译优化算子融合1.2-1.5x无1.3 KV Cache 优化KV Cache 存储中间结果避免重复计算KV Cache 每步推理: 新增1个token的K/V 复用之前所有token的K/V 内存: O(n * d * heads)2. 核心功能实现2.1 量化实现import torch import torch.nn as nn class QuantizedLinear(nn.Module): def __init__(self, weight, biasNone, bits4): super().__init__() self.bits bits self.weight self._quantize(weight, bits) self.bias bias def _quantize(self, weight, bits): scale (weight.max() - weight.min()) / (2**bits - 1) zero_point -weight.min() / scale quantized torch.round(weight / scale zero_point) quantized torch.clamp(quantized, 0, 2**bits - 1) return { values: quantized.to(torch.int8), scale: scale, zero_point: zero_point } def _dequantize(self): return self.weight[values].float() * self.weight[scale] - self.weight[zero_point] def forward(self, x): weight self._dequantize() return F.linear(x, weight, self.bias) class GPTQQuantizer: def __init__(self, model): self.model model def quantize(self, bits4): for name, module in self.model.named_modules(): if isinstance(module, nn.Linear): self._quantize_layer(module, bits) def _quantize_layer(self, layer, bits): H layer.weight.data rows, cols H.shape quantized torch.zeros(rows, cols, dtypetorch.int8) scales torch.zeros(rows) zeros torch.zeros(rows) for i in range(rows): row H[i] max_val row.abs().max() scale max_val / (2**(bits-1) - 1) zeros[i] 0 quantized_row torch.round(row / scale).to(torch.int8) quantized[i] quantized_row scales[i] scale layer.weight nn.Parameter(quantized) layer.register_buffer(scales, scales) layer.register_buffer(zeros, zeros)2.2 推理优化class OptimizedTransformer: def __init__(self, model): self.model model self.kv_cache {} def forward(self, input_ids, past_key_valuesNone): batch_size, seq_len input_ids.shape if past_key_values is None: past_key_values {} outputs [] for layer_idx, layer in enumerate(self.model.layers): key flayer_{layer_idx} if key in past_key_values: past_k, past_v past_key_values[key] else: past_k, past_v None, None output, new_k, new_v layer( input_ids, past_keypast_k, past_valuepast_v ) past_key_values[key] (new_k, new_v) input_ids output return output, past_key_values class FlashAttention: def __init__(self, causalTrue): self.causal causal def forward(self, Q, K, V): batch_size, heads, seq_len, d_k Q.shape Q Q.transpose(1, 2).reshape(batch_size * heads, seq_len, d_k) K K.transpose(1, 2).reshape(batch_size * heads, seq_len, d_k) V V.transpose(1, 2).reshape(batch_size * heads, seq_len, d_k) output self._flash_attention(Q, K, V) output output.reshape(batch_size, heads, seq_len, d_k).transpose(1, 2) return output def _flash_attention(self, Q, K, V): import flash_attn output flash_attn.flash_attn_func( Q, K, V, causalself.causal, return_attn_probsFalse ) return output2.3 编译优化class TorchCompileOptimizer: def __init__(self, model): self.model model def optimize(self, modereduce-overhead): self.model torch.compile(self.model, modemode) def compile_with_inductor(self): self.model torch.compile(self.model, backendinductor) class ONNXExporter: def __init__(self, model): self.model model def export(self, output_path): dummy_input torch.randint(0, 1000, (1, 32)) torch.onnx.export( self.model, dummy_input, output_path, opset_version15, input_names[input_ids], output_names[logits], dynamic_axes{ input_ids: {0: batch_size, 1: seq_len}, logits: {0: batch_size, 1: seq_len} } ) class TensorRTConverter: def __init__(self, model): self.model model def convert(self, output_path, precisionfp16): import tensorrt as trt builder trt.Builder(trt.Logger(trt.Logger.WARNING)) network builder.create_network() parser trt.OnnxParser(network, builder.logger) onnx_path output_path.replace(.engine, .onnx) self._export_onnx(onnx_path) with open(onnx_path, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) if precision fp16: config.set_flag(trt.BuilderFlag.FP16) engine builder.build_engine(network, config) with open(output_path, wb) as f: f.write(engine.serialize()) def _export_onnx(self, path): ONNXExporter(self.model).export(path)3. 性能对比3.1 量化方法对比方法精度加速比内存节省适用场景FP3232位1x0%训练FP1616位2x50%推理BF1616位2x50%训练/推理INT88位4x75%推理INT44位8x87%推理3.2 推理框架对比框架速度兼容性易用性PyTorch1x好高TensorRT2-4x中低ONNX Runtime1.5-2x好中TFLite2-3x中高3.3 KV Cache 影响序列长度KV Cache 内存推理时间(首token)推理时间(后续token)64128MB50ms10ms256512MB150ms10ms10242GB500ms10ms4. 最佳实践4.1 推理优化流程def optimize_for_inference(model, config): optimizer InferenceOptimizer(model) if config.get(quantize, False): optimizer.quantize(config[bits]) if config.get(compile, False): optimizer.compile() if config.get(flash_attention, False): optimizer.enable_flash_attention() return optimizer.model class InferenceOptimizer: def __init__(self, model): self.model model def quantize(self, bits4): quantizer GPTQQuantizer(self.model) quantizer.quantize(bits) def compile(self): compiler TorchCompileOptimizer(self.model) compiler.compile_with_inductor() def enable_flash_attention(self): self._replace_attention() def _replace_attention(self): for name, module in self.model.named_modules(): if attention in name.lower() and isinstance(module, nn.Module): module.forward self._flash_attention_forward4.2 部署建议class LLMDeployer: def __init__(self, model, config): self.model model self.config config def deploy(self): optimized_model optimize_for_inference(self.model, self.config) if self.config[platform] api: self._deploy_as_api(optimized_model) elif self.config[platform] edge: self._deploy_to_edge(optimized_model) def _deploy_as_api(self, model): from fastapi import FastAPI app FastAPI() app.post(/generate) def generate(prompt: str): return {response: model.generate(prompt)} return app def _deploy_to_edge(self, model): converter TensorRTConverter(model) converter.convert(model.engine, precisionfp16)5. 总结LLM 推理优化是部署的关键量化最有效的优化方法KV Cache避免重复计算Flash Attention内存高效的注意力计算编译优化进一步提升性能对比数据如下INT4 量化可实现 8 倍加速Flash Attention 可节省 30-50% 内存TensorRT 可再提升 2-4 倍速度推荐组合使用多种优化方法