1. 大模型推理加速的整体设计思路
1.1 为什么推理速度是大模型落地的第一道门槛
做过大模型应用的人都有一个共同体会：训练一个模型可能要几天甚至几周，但推理是每天、每时、每刻都在发生的事情。用户问一句话，模型要在一两秒内给出回答，这背后涉及的是一整套从显存管理、算子优化到批处理调度的系统工程。推理速度直接决定了产品的可用性——响应超过三秒，用户就会觉得“卡”；超过五秒，大概率直接关掉页面。
我最初接触大模型推理优化的时候，犯过一个很典型的错误：以为把模型权重加载进去、调一下model.generate()就完事了。实测下来，一个70亿参数的模型在单张消费级显卡上，生成一段200字的回答要将近8秒。这个速度放在Demo里勉强能看，放到真实产品里完全不可接受。后来一步步拆解，才发现推理速度的瓶颈分布在好几个层面，每一个层面都有对应的优化手段。
这篇文章要聊的，就是怎么把大模型的推理速度压到极致。我会从整体架构设计讲到具体的代码实现，包括显存优化、注意力机制加速、量化压缩、批处理策略等核心环节。适合有一定Python和深度学习基础、正在做或准备做大模型应用落地的开发者参考。即使你之前没接触过推理优化，跟着思路走也能理解每一步在做什么、为什么要这么做。
1.2 推理加速的核心矛盾：显存、算力与延迟的三角博弈
大模型推理优化本质上是在三个约束之间找平衡：显存容量、计算吞吐和响应延迟。这三个指标互相牵制，你优化了其中一个，往往会给另外两个带来压力。
显存容量决定了你能加载多大的模型、能支持多大的批处理。一个FP16精度的70亿参数模型，光权重就要占大约14GB显存，加上KV Cache和中间激活值，20GB起步。如果显存不够，要么换更小的模型，要么做量化压缩，要么把部分计算卸载到内存——但卸载会带来延迟飙升。
计算吞吐关注的是单位时间内能处理多少token。这跟GPU的算力利用率有关，如果算子实现得不好，GPU大量时间在等数据搬运，算力就浪费了。批处理是提升吞吐的经典手段，但批处理会增大显存占用，同时可能拉高单个请求的延迟。
响应延迟是用户直接感知的指标，通常用首token延迟和每token生成时间来衡量。首token延迟取决于预填充阶段的效率，每token生成时间取决于解码阶段的优化程度。这两个阶段的瓶颈完全不同，需要分开处理。
我个人的经验是：先明确你的场景最看重哪个指标。如果是对话类应用，首token延迟最关键；如果是批量内容生成，吞吐量优先；如果是端侧部署，显存容量是第一约束。目标不同，优化策略的优先级完全不一样。
1.3 从模型加载到token输出的完整推理链路
要优化推理速度，首先得搞清楚一次完整的推理请求到底经历了什么。以自回归生成模型为例，整个链路大致分为以下几个阶段：
模型加载阶段：把权重从磁盘读入显存，初始化计算图。这个阶段通常只发生一次，但如果使用动态量化或按需加载策略，可能会在推理过程中反复触发。
预填充阶段：把用户输入的prompt一次性喂给模型，计算所有输入token的Key和Value，存入KV Cache。这个阶段是计算密集型的，矩阵乘法的规模跟输入长度成正比。
解码阶段：逐个生成新token，每生成一个token都要跟之前所有token的KV做注意力计算。这个阶段是内存带宽密集型的，因为每步的计算量不大，但需要频繁读取KV Cache。
后处理阶段：把生成的token解码成文本，可能还包括采样、过滤、格式化等操作。这个阶段通常不是瓶颈，但如果采样策略复杂（比如beam search），也会消耗不少时间。
理解了这条链路，就知道优化该从哪里下手了。预填充阶段重点优化矩阵计算效率，解码阶段重点优化KV Cache的访问模式，模型加载阶段重点优化显存占用和加载速度。
1.4 不同规模模型的优化策略选择
不是所有模型都适合同一套优化方案。70亿参数的模型和700亿参数的模型，瓶颈完全不一样。
对于中小规模模型（1B-13B），显存压力相对可控，优化重点应该放在算子效率和批处理调度上。这类模型通常可以完整加载到单张显卡，甚至可以用消费级显卡跑起来。量化到INT8或INT4后，显存占用可以压到很低，适合端侧或边缘部署。
对于大规模模型（30B-70B），单卡显存往往不够，需要做模型并行或张量并行。这时候通信开销成为新的瓶颈，优化重点转向减少卡间通信、重叠计算与通信。量化在这里也很关键，但精度损失需要更仔细地评估。
对于超大规模模型（100B以上），通常需要多机多卡部署，推理框架的选择比单点优化更重要。这类场景下，批处理策略、请求调度、负载均衡对整体性能的影响远大于单个算子的优化。
我下面要讲的代码实现，主要以中小规模模型在单卡上的推理优化为例，因为这是大多数开发者能直接上手复现的场景。大规模部署的思路会顺带提到，但不会展开太多细节。
2. 核心加速技术的原理与代码实现
2.1 KV Cache：自回归生成的第一加速利器
KV Cache是大模型推理中最基础也是最重要的优化手段。没有它，每生成一个token都要重新计算所有历史token的Key和Value，计算量会随着生成长度呈平方级增长。
原理其实很直观：在自回归生成中，每个新token的注意力计算需要用到之前所有token的Key和Value。如果不缓存，每步都要重新算一遍，做了大量重复计算。KV Cache的做法是：预填充阶段把所有输入token的K、V算好存起来，解码阶段每生成一个新token，只算这个新token的K、V，然后跟缓存里的拼接。
代码实现上，核心是维护两个张量：
class KVCache: def __init__(self, max_batch_size, max_seq_len, num_heads, head_dim, dtype, device): self.k_cache = torch.zeros( max_batch_size, num_heads, max_seq_len, head_dim, dtype=dtype, device=device ) self.v_cache = torch.zeros( max_batch_size, num_heads, max_seq_len, head_dim, dtype=dtype, device=device ) self.seq_len = 0 def update(self, new_k, new_v): batch_size, num_heads, new_len, head_dim = new_k.shape self.k_cache[:, :, self.seq_len:self.seq_len + new_len, :] = new_k self.v_cache[:, :, self.seq_len:self.seq_len + new_len, :] = new_v self.seq_len += new_len return self.k_cache[:, :, :self.seq_len, :], self.v_cache[:, :, :self.seq_len, :]这段代码看起来简单，但有几个关键细节容易踩坑。第一，预分配固定大小的缓存比动态扩容快得多，因为避免了反复申请显存和拷贝数据。第二，缓存的形状布局会影响后续注意力计算的效率，通常把head_dim放在最后一维对GPU更友好。第三，max_seq_len要设得合理，太大浪费显存，太小会截断生成长度。
实测数据：在一个70亿参数模型上，开启KV Cache后，生成512个token的时间从约12秒降到了约3.5秒，加速比超过3倍。生成长度越长，加速效果越明显。
2.2 PagedAttention：解决显存碎片化的关键方案
KV Cache虽然好用，但有个致命问题：显存碎片化。每个请求的KV Cache长度不一样，如果按最大长度预分配，短请求会浪费大量显存；如果动态分配，显存碎片又会降低利用率。
PagedAttention的思路借鉴了操作系统的虚拟内存分页机制。把KV Cache切成固定大小的块（比如每块存16个token的KV），不同请求的块可以离散存储在显存中，通过一个块表来索引。这样显存利用率可以接近100%，而且支持不同请求之间共享块（比如相同的system prompt）。
实现上，核心是块管理器和注意力计算时的块索引：
class BlockManager: def __init__(self, num_blocks, block_size, num_heads, head_dim, dtype, device): self.block_size = block_size self.free_blocks = list(range(num_blocks)) self.k_blocks = torch.zeros( num_blocks, num_heads, block_size, head_dim, dtype=dtype, device=device ) self.v_blocks = torch.zeros( num_blocks, num_heads, block_size, head_dim, dtype=dtype, device=device ) self.block_tables = {} def allocate(self, request_id, num_tokens): num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size if len(self.free_blocks) < num_blocks_needed: raise RuntimeError("Out of KV cache blocks") blocks = [self.free_blocks.pop() for _ in range(num_blocks_needed)] self.block_tables[request_id] = blocks return blocks def free(self, request_id): if request_id in self.block_tables: self.free_blocks.extend(self.block_tables.pop(request_id))PagedAttention的注意力计算需要根据块表来 gather 对应的KV块，比连续内存的注意力计算多了一步索引操作，但换来的是显存利用率的大幅提升。在vLLM的论文中，PagedAttention把吞吐量提升了2到4倍，主要就是靠显存利用率的改善。
2.3 量化压缩：用精度换速度的工程取舍
量化是把模型权重和激活值从高精度（FP16/FP32）转成低精度（INT8/INT4）的过程。好处很直接：显存占用成倍降低，内存带宽压力减小，计算速度提升。代价是精度损失，需要评估对下游任务的影响。
目前主流的量化方案分两类：训练后量化（PTQ）和量化感知训练（QAT）。PTQ不需要重新训练，直接对训练好的模型做量化，实现简单但精度损失可能较大。QAT在训练过程中模拟量化误差，精度保持更好但需要训练资源。
对于推理加速场景，PTQ更常用。以INT8量化为例，核心是把浮点权重映射到8位整数：
def quantize_weight(weight, num_bits=8): qmin = -(2 ** (num_bits - 1)) qmax = 2 ** (num_bits - 1) - 1 scale = (weight.max() - weight.min()) / (qmax - qmin) zero_point = qmin - round(weight.min() / scale) zero_point = torch.clamp(zero_point, qmin, qmax) q_weight = torch.clamp( torch.round(weight / scale) + zero_point, qmin, qmax ).to(torch.int8) return q_weight, scale, zero_point def dequantize_weight(q_weight, scale, zero_point): return (q_weight.to(torch.float32) - zero_point) * scale实际工程中，量化不是简单地对整个模型一刀切。通常的做法是：对权重做INT8量化，对激活值保持FP16，注意力计算中的softmax保持FP16，LayerNorm等对精度敏感的算子保持FP16。这种混合精度策略能在精度损失可控的前提下获得大部分加速收益。
踩过的坑：一开始我把所有层都量化到INT8，结果模型输出质量明显下降，重复生成、逻辑混乱的问题很突出。后来改成只量化线性层的权重，激活值和其他算子保持FP16，精度就恢复到了可接受范围。量化粒度和量化对象的选择，比量化位数本身更重要。
2.4 FlashAttention：注意力计算的IO优化
标准注意力计算的问题在于，它需要把Q、K、V矩阵完整地读写一遍，中间还要存储注意力分数矩阵。对于长序列，这个中间矩阵非常大，显存读写成为瓶颈。
FlashAttention的核心思想是分块计算+在线softmax。把Q、K、V切成小块，每次只加载一小块到SRAM（GPU的高速缓存）中计算，避免了大矩阵的反复读写。在线softmax的技巧使得不需要预先知道所有分数就能正确计算归一化。
def flash_attention_forward(q, k, v, block_size=128): batch, heads, seq_len, dim = q.shape scale = 1.0 / math.sqrt(dim) output = torch.zeros_like(q) for i in range(0, seq_len, block_size): q_block = q[:, :, i:i+block_size, :] acc = torch.zeros_like(q_block) m = torch.full( (batch, heads, q_block.shape[2], 1), float('-inf'), device=q.device ) l = torch.zeros_like(m) for j in range(0, seq_len, block_size): k_block = k[:, :, j:j+block_size, :] v_block = v[:, :, j:j+block_size, :] s = torch.matmul(q_block, k_block.transpose(-2, -1)) * scale m_new = torch.maximum(m, s.max(dim=-1, keepdim=True).values) p = torch.exp(s - m_new) l = l * torch.exp(m - m_new) + p.sum(dim=-1, keepdim=True) acc = acc * torch.exp(m - m_new) + torch.matmul(p, v_block) m = m_new output[:, :, i:i+block_size, :] = acc / l return output这段代码是FlashAttention的简化版，实际生产中使用的是经过高度优化的CUDA kernel。但理解这个逻辑很重要：它解释了为什么FlashAttention能大幅减少显存访问——因为它在SRAM中完成了大部分计算，只把最终结果写回显存。
实测中，FlashAttention在序列长度超过512时开始显现优势，序列越长优势越明显。在2048长度的序列上，相比标准注意力可以提速2到3倍，显存占用降低5到10倍。
2.5 连续批处理：提升GPU利用率的调度策略
静态批处理的问题是：一个批次里所有请求必须等最长的那个生成完才能释放。如果批次里有一个请求要生成500个token，其他请求只要生成50个，那GPU就有大量时间在等。
连续批处理（Continuous Batching）的思路是：不等整个批次完成，而是每生成一步就检查有没有请求完成，完成了就移出，同时把等待队列里的新请求加进来。这样GPU始终在处理有效计算，利用率大幅提升。
class ContinuousBatchScheduler: def __init__(self, max_batch_size, max_seq_len): self.max_batch_size = max_batch_size self.max_seq_len = max_seq_len self.running = [] self.waiting = [] def step(self): finished = [r for r in self.running if r.is_finished()] for r in finished: self.running.remove(r) while len(self.running) < self.max_batch_size and self.waiting: new_req = self.waiting.pop(0) self.running.append(new_req) if not self.running: return None batch_input = self.prepare_batch(self.running) outputs = self.model_forward(batch_input) for req, out in zip(self.running, outputs): req.append_token(out) return outputs连续批处理的关键在于批次的动态重组。每步推理的批次组成可能都不一样，这要求KV Cache的管理必须支持动态增删。PagedAttention天然适合这种场景，因为块可以独立分配和释放。
在实际部署中，连续批处理通常能把GPU利用率从30%提升到70%以上，吞吐量提升2到3倍。对于请求长度差异大的场景，效果尤其明显。
3. 完整推理引擎的搭建与实操
3.1 环境准备与依赖选型
搭建推理引擎之前，先把环境理清楚。我推荐的配置是：
- Python 3.10+：太老的版本对新的CUDA库支持不好
- PyTorch 2.1+：自带编译优化和更好的算子支持
- CUDA 12.1+：匹配显卡驱动，支持最新的FlashAttention
- Triton：用于编写自定义GPU kernel，FlashAttention依赖它
安装命令大致如下：
pip install torch==2.1.0 --index-url https://download.pytorch.org/whl/cu121 pip install triton==2.1.0 pip install transformers==4.36.0 pip install accelerate==0.25.0如果你的显卡支持FP8计算（比如较新的数据中心卡），还可以装transformer-engine来启用FP8推理。消费级显卡一般用FP16或INT8就够了。
注意：PyTorch版本和CUDA版本必须匹配，否则会出现各种奇怪的运行时错误。我遇到过因为CUDA版本不匹配导致FlashAttention kernel编译失败的情况，排查了大半天才发现是版本问题。建议用
torch.cuda.is_available()和torch.version.cuda先确认环境正常。
3.2 模型加载与显存预分配策略
模型加载阶段有几个优化点容易被忽略。第一，用torch_dtype=torch.float16加载，直接以半精度读入，避免先加载FP32再转换的额外显存开销。第二，用device_map控制权重加载到哪张卡，多卡场景下合理分配。第三，预分配KV Cache的显存池，避免推理过程中反复申请释放。
import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "your-model-path" tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype=torch.float16, device_map="auto", trust_remote_code=True, use_cache=True, ) model.eval() # 预分配KV Cache显存池 max_batch_size = 8 max_seq_len = 2048 num_layers = model.config.num_hidden_layers num_heads = model.config.num_attention_heads head_dim = model.config.hidden_size // num_heads kv_cache_pool = [] for _ in range(num_layers): k_cache = torch.zeros( max_batch_size, num_heads, max_seq_len, head_dim, dtype=torch.float16, device="cuda" ) v_cache = torch.zeros_like(k_cache) kv_cache_pool.append((k_cache, v_cache))预分配的好处是避免了推理时的动态显存分配开销。PyTorch的缓存分配器虽然已经很快了，但在高频推理场景下，预分配仍然能带来10%到20%的延迟降低。
3.3 推理主循环的代码实现
把前面讲的各个模块串起来，推理主循环大致长这样：
@torch.inference_mode() def generate(model, tokenizer, prompt, max_new_tokens=256, temperature=0.7): input_ids = tokenizer.encode(prompt, return_tensors="pt").cuda() batch_size, prompt_len = input_ids.shape # 预填充阶段 with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model(input_ids, use_cache=True) past_key_values = outputs.past_key_values next_token_logits = outputs.logits[:, -1, :] generated_ids = [input_ids] for _ in range(max_new_tokens): # 采样 probs = torch.softmax(next_token_logits / temperature, dim=-1) next_token = torch.multinomial(probs, num_samples=1) generated_ids.append(next_token) # 检查是否遇到结束符 if next_token.item() == tokenizer.eos_token_id: break # 解码阶段 with torch.autocast(device_type="cuda", dtype=torch.float16): outputs = model( next_token, past_key_values=past_key_values, use_cache=True ) past_key_values = outputs.past_key_values next_token_logits = outputs.logits[:, -1, :] return tokenizer.decode( torch.cat(generated_ids, dim=-1)[0], skip_special_tokens=True )这段代码是单请求版本，逻辑清晰但性能一般。生产环境需要改成批处理版本，并且把KV Cache的管理从HuggingFace的默认实现换成自定义的PagedAttention实现。
3.4 性能测试与瓶颈定位方法
优化不能靠猜，得有数据支撑。我常用的性能测试方法是：
首token延迟测试：记录从输入prompt到生成第一个token的时间。这个指标反映预填充阶段的效率。
每token生成时间测试：记录解码阶段每个token的平均生成时间。这个指标反映解码阶段的效率。
吞吐量测试：在固定时间内，统计能完成多少个请求。这个指标反映整体系统的处理能力。
import time def benchmark(model, tokenizer, prompts, max_new_tokens=128): # 预热 for _ in range(3): generate(model, tokenizer, prompts[0], max_new_tokens=10) # 首token延迟 torch.cuda.synchronize() start = time.perf_counter() input_ids = tokenizer.encode(prompts[0], return_tensors="pt").cuda() outputs = model(input_ids, use_cache=True) torch.cuda.synchronize() ttft = time.perf_counter() - start # 每token时间 past = outputs.past_key_values next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) times = [] for _ in range(max_new_tokens): torch.cuda.synchronize() t0 = time.perf_counter() outputs = model(next_token, past_key_values=past, use_cache=True) torch.cuda.synchronize() times.append(time.perf_counter() - t0) past = outputs.past_key_values next_token = outputs.logits[:, -1, :].argmax(dim=-1, keepdim=True) return { "ttft_ms": ttft * 1000, "tpot_ms": sum(times) / len(times) * 1000, "throughput_tokens_per_sec": len(times) / sum(times) }用torch.cuda.synchronize()确保GPU操作完成后再计时，否则测到的是CPU提交任务的时间，不是实际计算时间。这个坑我踩过，一开始测出来的数据好得离谱，后来发现是没同步。
3.5 多卡部署与张量并行的基本配置
单卡显存不够的时候，就需要多卡部署。最简单的方式是张量并行：把每一层的权重切分到多张卡上，每张卡算一部分，然后通过通信汇总结果。
# 使用accelerate进行简单的模型并行 from accelerate import dispatch_model, infer_auto_device_map device_map = infer_auto_device_map( model, max_memory={0: "20GiB", 1: "20GiB"}, no_split_module_classes=["LlamaDecoderLayer"] ) model = dispatch_model(model, device_map=device_map)张量并行的通信开销主要在注意力计算和全连接层。每层都要做all-reduce，卡越多通信占比越高。实测下来，2卡张量并行的加速比大约1.7倍，4卡大约2.8倍，8卡大约4倍——不是线性增长，因为通信开销在增加。
如果模型更大，单机多卡也放不下，就需要流水线并行或者专家并行。这些方案更复杂，通常直接用现成的推理框架（如vLLM、TensorRT-LLM）比自己实现更划算。
4. 常见问题与排查技巧实录
4.1 显存溢出（OOM）的排查与解决
OOM是推理中最常见的问题。排查思路是：先确认是哪个阶段OOM，再定位具体原因。
| 阶段 | 常见原因 | 解决方案 |
|---|---|---|
| 模型加载 | 权重精度太高 | 用FP16或INT8加载 |
| 预填充 | 输入序列太长 | 截断输入或分块处理 |
| 解码 | KV Cache太大 | 限制生成长度或用量化KV |
| 批处理 | 批次太大 | 减小batch size或动态调整 |
一个容易被忽略的点是：PyTorch的缓存分配器会保留已释放的显存，导致nvidia-smi显示的占用比实际使用高。可以用torch.cuda.memory_allocated()看实际使用量，用torch.cuda.empty_cache()释放缓存。
实操技巧：在推理循环中定期调用
torch.cuda.empty_cache()可以缓解碎片化问题，但不要每步都调，因为empty_cache本身有开销。我一般每处理100个请求调一次。
4.2 生成质量下降的归因分析
量化或优化之后，如果发现生成质量下降，需要系统性地排查。常见原因和排查方法：
量化精度损失：对比量化前后的输出，看是普遍性下降还是特定类型输入下降。如果是普遍下降，降低量化位数或改用混合精度。
KV Cache精度问题：KV Cache如果也用INT8存储，长序列生成时误差会累积。可以只对Key做量化，Value保持FP16。
注意力实现差异：FlashAttention和标准注意力的数值结果有细微差异，通常不影响质量，但如果模型对数值敏感，可能需要调整。
采样参数不匹配：优化过程中如果改了temperature、top_p等参数，输出风格会变。确保对比实验时采样参数一致。
4.3 批处理场景下的请求调度问题
批处理虽然提升吞吐，但会引入调度问题。典型场景：一个长请求和多个短请求混在一起，长请求会拖慢短请求的响应。
解决方案是优先级调度+分桶。把请求按预估生成长度分到不同桶里，同桶内的请求一起批处理。预估长度可以用prompt长度加上一个经验值。
def bucket_requests(requests, num_buckets=4): buckets = [[] for _ in range(num_buckets)] for req in requests: estimated_len = len(req.prompt) + req.max_new_tokens bucket_idx = min(estimated_len // 512, num_buckets - 1) buckets[bucket_idx].append(req) return buckets另一个技巧是给每个请求设置超时时间，超过时间还没生成完就强制返回当前结果。这样避免个别请求无限占用资源。
4.4 推理速度不达预期的检查清单
优化做完之后速度还是不理想，按这个清单逐项检查：
- 确认GPU利用率：用
nvidia-smi -l 1观察推理时的GPU利用率。如果低于50%，说明有CPU瓶颈或调度问题。 - 确认是否用了KV Cache：检查
use_cache=True是否生效，past_key_values是否在传递。 - 确认注意力实现：是否用了FlashAttention或等效的优化实现。
- 确认数据类型：权重和激活值是否都是FP16，有没有意外的FP32转换。
- 确认批处理策略：是否用了连续批处理，批次大小是否合理。
- 确认没有同步操作：推理循环中是否有不必要的
.item()或.cpu()调用，这些会强制同步，打断GPU流水线。 - 确认CUDA Graph：对于固定形状的推理，用CUDA Graph可以消除kernel启动开销，通常能再提升10%到20%。
4.5 不同硬件平台的适配经验
不同显卡的优化重点不一样。消费级显卡（如RTX系列）显存带宽相对有限，量化带来的收益更明显。数据中心卡（如A100/H100）算力强、带宽高，重点应该放在批处理和并行策略上。
Apple Silicon芯片（M系列）用统一内存架构，显存和内存共享，适合跑中小模型。用MPS后端可以调用GPU加速，但算子支持不如CUDA完善，部分优化技术用不了。
我的建议是：先在你实际部署的硬件上做基准测试，再决定优化方向。不要盲目照搬别人在别的硬件上的优化方案，瓶颈可能完全不一样。
5. 进阶优化方向与个人实践体会
5.1 投机解码：用小模型加速大模型
投机解码（Speculative Decoding）的思路很巧妙：用一个小模型快速生成多个候选token，然后用大模型一次性验证这些token是否正确。如果小模型猜得准，大模型一次就能确认多个token，相当于把多次解码合并成一次。
def speculative_decode(draft_model, target_model, input_ids, max_new_tokens=128, gamma=5): generated = [] for _ in range(max_new_tokens // gamma): # 小模型生成gamma个候选 draft_tokens = [] draft_probs = [] draft_input = input_ids for _ in range(gamma): with torch.no_grad(): out = draft_model(draft_input, use_cache=True) probs = torch.softmax(out.logits[:, -1, :], dim=-1) token = torch.multinomial(probs, 1) draft_tokens.append(token) draft_probs.append(probs) draft_input = token # 大模型验证 candidate = torch.cat(draft_tokens, dim=-1) with torch.no_grad(): target_out = target_model( torch.cat([input_ids, candidate], dim=-1), use_cache=True ) # 逐个验证并接受 accepted = 0 for i, token in enumerate(draft_tokens): target_prob = torch.softmax( target_out.logits[:, input_ids.shape[1] + i - 1, :], dim=-1 ) accept_prob = min(1.0, (target_prob[0, token] / draft_probs[i][0, token]).item()) if torch.rand(1).item() < accept_prob: accepted += 1 else: break generated.extend(draft_tokens[:accepted]) input_ids = torch.cat([input_ids] + draft_tokens[:accepted], dim=-1) return generated投机解码的加速比取决于小模型和大模型输出分布的重合度。如果小模型跟大模型风格接近，接受率高，加速比可以到2到3倍。如果接受率低，反而可能变慢，因为验证的开销也不小。
5.2 动态批处理与请求优先级的平衡
生产环境中，请求的优先级往往不一样。系统prompt的预填充、用户的实时对话、后台的批量生成，对延迟的敏感度完全不同。动态批处理需要把这些因素考虑进去。
我的做法是给每个请求打两个标签：延迟敏感度和预估计算量。调度器根据这两个标签决定请求进入哪个批次、什么时候执行。延迟敏感度高的请求优先处理，计算量大的请求单独成批避免拖慢其他请求。
这套逻辑实现起来不复杂，但需要跟业务方对齐优先级规则。技术上的优化空间是一回事，业务上的取舍是另一回事。
5.3 我在实际项目中的优化顺序建议
回顾我做过的几个推理优化项目，总结出一个比较稳妥的优化顺序：
第一步，先把KV Cache用上。这是最基础的优化，投入小收益大，通常能带来2到3倍的加速。
第二步，换用FlashAttention。如果序列长度超过512，这个优化的收益很明显，而且实现成本不高，很多框架已经内置了。
第三步，做量化。INT8量化通常能再提升30%到50%的速度，同时显存占用减半。注意做好精度评估。
第四步，上连续批处理。如果是服务端部署，这个优化对吞吐量的提升最大，但实现复杂度也最高。
第五步，考虑投机解码或CUDA Graph。这些是锦上添花的优化，在前几步都做完之后再考虑。
不要一上来就追求最复杂的方案。我见过有人直接上投机解码，结果因为小模型没选好，整体速度反而比基线还慢。先把基础优化做扎实，再逐步叠加高级技术。
5.4 推理框架选型：自研还是用现成的
最后聊一下框架选型。如果只是做实验或小规模部署，HuggingFace Transformers加上FlashAttention和量化就够了。如果要上生产环境，vLLM或TensorRT-LLM是更成熟的选择。
vLLM的优势是PagedAttention和连续批处理都是内置的，开箱即用，社区活跃，支持模型多。TensorRT-LLM的优势是NVIDIA官方优化，在N卡上性能极致，但编译流程复杂，模型支持相对少。
自研推理引擎适合有特殊需求的场景，比如要支持自定义算子、要跟现有系统深度集成、或者要对调度策略做精细控制。但自研的成本很高，没有足够的工程资源不建议走这条路。
我个人在实际操作中的体会是：优化的收益遵循边际递减规律。从零到KV Cache，可能提升3倍；从KV Cache到FlashAttention，再提升1.5倍；从FlashAttention到量化，再提升1.3倍。越往后，投入产出比越低。所以要先做收益最大的优化，把基础打牢，再根据实际瓶颈决定要不要继续深入。
