当我们讨论大模型推理加速时绕不开的一个思路就是“用一个更小更快的模型先写草稿再用大模型验证”也就是推测解码Speculative Decoding。这套方案在自回归模型上已经非常成熟但当草稿模型变成扩散模型Diffusion Model时事情就开始变得复杂起来。最近在阅读相关论文时看到了 xPress 这个思路它专门用来解决扩散模型作为草稿器时的验证效率问题。本文结合自己的理解把 xPress 的动机、核心思想、与现有方案的差异以及如何在 PyTorch 环境下做一个最小的原型验证完整拆解一遍。1. 背景与核心概念1.1 为什么需要推测解码大语言模型LLM的生成过程本质上是逐 token 自回归。每一步生成一个 token都需要把当前序列重新喂给模型做一次完整的前向计算。这个过程有两个明显的瓶颈一是显存带宽受限二是每一步的并行度极低。假设我们用的是 7B 参数的模型在 A100 上单卡推理每次前向计算虽然只需要几毫秒但生成 2000 个 token 就需要 2000 次串行的前向计算。总延迟就是 2000 乘以单步延迟。这个延迟对很多实时交互场景来说是不可接受的。于是研究者提出了一个很直观的想法能不能让一个更快的小模型先“猜”几个 token然后让大模型一次性验证这些 token如果猜对了就一次性接受多个 token这样大模型的前向计算次数就大幅减少。这就是推测解码的核心思想。1.2 草稿-验证框架推测解码的完整框架包含两部分草稿模型Draft Model是一个小模型速度很快用来生成候选 token 序列。目标模型Target Model是真正负责“质量”的大模型用来验证草稿模型生成的 token 是否可接受。整个过程可以简单描述为草稿模型快速生成 K 个候选 token。目标模型对这 K 个 token 做一次前向计算得到这 K 个位置的真实概率分布。按照拒绝采样规则逐个接受草稿 token直到遇到第一个被拒绝的 token。用目标模型在该位置重新采样一个 token然后继续下一轮。这种方式的核心收益在于目标模型一次前向计算可以验证多个 token而不是每生成一个 token 就前向一次。如果草稿模型足够聪明接受率足够高那么整体加速比就会非常可观。1.3 扩散模型草稿器的特殊性传统的草稿模型通常也是自回归模型比如一个小规模的 Transformer。但自回归草稿模型存在一个问题它在“猜”的时候也是逐 token 生成的只不过模型小、速度快但本质上仍是串行过程。扩散模型则完全不同。扩散模型在生成时是对整个候选序列同时去噪天然具备并行生成多个 token 的能力。这就让它成为草稿模型的理想候选者——因为草稿阶段往往是并行生成 K 个 token而不是串行生成 K 次。然而问题也随之而来。自回归草稿模型生成的 token 之间是有依赖关系的而扩散草稿模型一次生成的 K 个 token 之间往往缺乏足够的自回归依赖。这就导致目标模型在验证时后面位置的 token 接受率可能非常低。如果接受率低那么草稿-验证流程的收益就会被大幅削减。1.4 xPress 要解决什么问题xPress 的全称是 Parallel Refinement for Diffusion Drafters in Speculative Decoding从名字可以看出核心是“并行细化”。它解决的核心问题是当扩散模型作为草稿器时如何通过并行细化的方式提高草稿 token 的接受率从而提升推测解码的整体加速比。这里需要区分两个概念草稿生成Draft Generation扩散模型一次性生成 K 个候选 token。草稿细化Draft Refinement在验证之前对草稿 token 进行进一步修正让它们更接近目标模型的分布。xPress 的重点在第二个环节。它试图在验证阶段之前增加一个并行的细化阶段让草稿 token 在被目标模型验证之前就已经拥有更高的质量和更好的自回归一致性。2. 推测解码的数学基础与验证逻辑2.1 拒绝采样机制要理解 xPress必须先理解推测解码中的验证逻辑。假设目标模型记为 M草稿模型记为 D。当前已有序列为 s。草稿模型生成 K 个 token记为 x_1, x_2, ..., x_K。目标模型一次前向计算后得到每个位置的条件概率分布P_target(x_i | s, x_1, ..., x_{i-1})同时草稿模型也给出了每个 token 的生成概率P_draft(x_i | s, x_1, ..., x_{i-1})验证第 i 个 token 时计算接受概率accept_prob min(1, P_target / P_draft)然后以 accept_prob 的概率接受该 token否则拒绝。如果拒绝就在该位置用目标模型的分布重新采样一个 token本轮验证结束。这个机制的数学保证是最终生成序列的分布恰好等于目标模型的真实分布。也就是说推测解码不改变输出分布只改变计算方式。2.2 接受率与加速比加速比近似公式为speedup (K 1) / (1 K * (1 - acceptance_rate))当接受率为 1 时加速比约为 K1也就是草稿长度越长越好。但当接受率接近 0 时加速比会跌破 1也就是比直接自回归还慢。因此推测解码的加速效果完全取决于草稿模型的接受率。这也是 xPress 选择“并行细化”的直接原因——与其试图让扩散模型一次生成完美的 K 个 token不如在验证前先对草稿做一次修正。2.3 扩散草稿器的天然劣势扩散模型在生成时通常会对整个序列施加一个“全局规划”式的去噪过程。对于图像生成来说这是优势因为图像的空间结构是高度全局化的。但对文本生成来说token 之间的语义连贯性主要依赖自回归依赖。扩散草稿器一次性生成 K 个 token 时后面 token 的生成条件中并没有包括前面 token 的“真实值”而是使用了一些粗略的引导。这就会导致一个典型现象第一个 token 的接受率很高。第二个到第 K 个 token 的接受率逐步下降。到序列中后段时接受率可能低于 0.1。这种现象让扩散草稿器的实际收益大打折扣。3. xPress 核心思想并行细化3.1 细化阶段的设计动机xPress 的核心设计思路是在草稿生成之后、目标模型验证之前插入一个“并行细化”阶段。想象这样一个场景扩散模型已经生成了一整段 K 个候选 token。此时这套序列整体看起来可能“差不多”但有些 token 不一定是最优的。传统做法是直接交给目标模型验证接受率可能不高。xPress 的思路是先用某种方式对这批 token 做一次并行的修正。修正的方向是让每个 token 的分布更接近目标模型的分布。由于所有 token 的修正是并行进行的所以引入的开销很小。为什么不能直接让目标模型做这个修正因为如果让目标模型做了一次完整的前向计算那就失去了“节省一次前向”的意义。xPress 的修正过程应该使用更轻量的手段或者复用同一个修正模型进行多轮迭代。3.2 并行细化的具体过程把 xPress 的细化过程拆开看可以分成以下几个步骤扩散草稿器生成 K 个候选 token。对 K 个 token 进行分组每组包含若干个 token。对每个分组同时进行细化修正。细化的目标是让每组的 token 分布更接近目标模型的边际分布。将细化后的 token 序列作为新的草稿交给目标模型验证。这里的关键点是细化阶段不引入串行依赖。所有分组的细化是并行的。这得益于扩散模型天然支持并行处理——因为每个分组都可以看作一个“局部去噪”任务。3.3 与 ROI 类方法的对比在 xPress 之前的方案中比较典型的一类思路是 ROIRegions of Interest感兴趣区域对齐。这类方法会识别草稿序列中哪些 token 是“低置信度”的然后只对这些 token 进行修复。ROI 类方法的优点是修复成本低但缺点是“识别低置信度 token”这个过程本身需要额外的计算而且当低置信度 token 过多时修复效率会下降。xPress 的并行细化思路则更彻底不是挑选部分 token 做修复而是对全量 token 并行做一轮或多轮细化。这样做的好处是不需要额外的“识别”阶段减少了流程复杂度。所有 token 都会被修正不会出现漏修的情况。并行化程度高非常适合 GPU 计算。当然缺点也很明显如果细化轮数很多计算开销会上升。xPress 的设计核心就是找到最优的细化轮数让收益最大。需要说明的是xPress 论文中的具体实现细节在不同版本中可能有调整。在理解框架时应该抓住“并行细化”这四个字这是它的灵魂。4. 环境准备与实验设计4.1 实验环境说明由于 xPress 作为一个研究方案并没有可以直接 pip install 的官方库。不过我们可以通过模拟的方式来理解它的流程并用现有模型库搭建一个最小验证环境。版本需要根据你的项目实际情况调整本文示例以常见环境为例重点演示配置思路。组件说明操作系统Ubuntu 20.04 / 22.04GPUNVIDIA A100 / 3090 / 4090显存建议 16G 以上Python3.9 或 3.10PyTorch2.xHuggingFace Transformers4.36 以上目标模型任选一个小型生成模型如 GPT-2、Phi-2 等草稿模型扩散语言模型或随机模拟器如果你只是理解原理不要求真实模型推理用 CPU 环境也可以完成模拟实验。但要做完整的验证效果测试还是需要 GPU。4.2 安装依赖pip install torch transformers datasets建议再安装一个用于计时的库pip install timeit或者直接使用 Python 内置的 time 模块无需额外安装。4.3 模拟实验总体设计由于没有 xPress 的官方实现我们用模拟的方式验证“并行细化”这个概念的价值。实验步骤如下模拟一个草稿生成器生成 K 个候选 token。模拟一个目标模型给出每个 token 的真实接受概率。分别测试“不过细化直接验证”和“经过并行细化后再验证”两种场景。对比两种场景下的接受率和端到端耗时。这里要强调这不是 xPress 论文源码的复现而是为了帮助理解核心概念而设计的教学实验。5. 核心代码实现5.1 编写目标模型验证类我们用一个简单的模拟类来表示目标模型。在实际场景中这个类会调用真实的大模型但在这里我们用一个固定的概率分布来模拟。# 文件路径speculative_decoding/simulator.py import random import time class TargetModelSimulator: 模拟目标模型输出每个 token 的真实分布 def __init__(self, vocab_size: int 100, seed: int 42): self.vocab_size vocab_size random.seed(seed) # 生成一个固定的偏好分布模拟“真实模型”的偏好 self.preferred_tokens list(range(vocab_size)) random.shuffle(self.preferred_tokens) def get_token_distribution(self, sequence: list) - dict: 给定一个前缀序列返回下一个 token 的概率分布。 这里用模拟数据真实场景中应该调用大模型 forward 得到 logits。 # 模拟一次前向计算耗时 time.sleep(0.01) distribution {} for token in self.preferred_tokens: # 模拟分布越靠前的 token 概率越高 distribution[token] 1.0 / (1 self.preferred_tokens.index(token)) total sum(distribution.values()) # 归一化保证是概率分布 for token in distribution: distribution[token] / total return distribution def verify_sequence(self, sequence: list, draft_probs: list) - dict: 验证草稿序列返回哪些位置被接受哪些被拒绝 accepted [] rejected [] for i, token in enumerate(sequence): # 目标模型在当前位置的真实分布 target_dist self.get_token_distribution(sequence[:i]) # 草稿模型给出的概率 draft_prob draft_probs[i] # 目标模型认为这个 token 的概率 target_prob target_dist.get(token, 0) # 拒绝采样规则 accept_prob min(1, target_prob / (draft_prob 1e-10)) if random.random() accept_prob: accepted.append(i) else: rejected.append(i) break return accepted, rejected这段代码中get_token_distribution模拟了目标模型的前向计算。实际场景中需要换成真实的模型推理。这里用time.sleep(0.01)模拟耗时是为了让实验能看出性能差异。5.2 编写草稿生成器模拟类草稿生成器模拟扩散模型的行为——一次性生成 K 个 token而不是逐 token 生成。# 文件路径speculative_decoding/simulator.py class DiffusionDrafterSimulator: 模拟扩散模型草稿器一次性生成 K 个候选 token def __init__(self, vocab_size: int 100, seed: int 100): self.vocab_size vocab_size random.seed(seed) def draft_sequence(self, length: int 8) - list: 一次性生成 length 个候选 token。 在这里用随机采样模拟真实场景中是扩散模型去噪过程的输出。 time.sleep(0.005) # 模拟草稿生成耗时 return [random.randint(0, self.vocab_size - 1) for _ in range(length)] def draft_probs(self, sequence: list) - list: 返回草稿模型认为每个 token 的生成概率。 这里模拟一个接近均匀分布的置信度。 probs [] for _ in sequence: # 假设草稿模型对每个 token 的置信度大约是 0.1 probs.append(0.1) return probs这里要注意draft_probs是草稿模型给出的概率。在实际模型中这个概率来自扩散模型的去噪置信度。我们的模拟中草稿模型的置信度是固定的 0.1。目标模型计算出的概率如果大于等于 0.1则接受概率为 1如果小于 0.1则可能拒绝。5.3 并行细化模块实现这是 xPress 核心思想的模拟实现。我们实现两类细化策略顺序细化baseline依次对每个 token 进行修正。并行细化xPress分组后并行修正。为了模拟并行效果我们使用 ThreadPoolExecutor 来模拟多个并行 worker。# 文件路径speculative_decoding/refiner.py from concurrent.futures import ThreadPoolExecutor, as_completed class Refiner: 细化器对草稿 token 进行修正 def __init__(self, target_model, threshold: float 0.08): self.target_model target_model self.threshold threshold def refine_one_token(self, token: int, context: list) - int: 对单个 token 进行细化 如果目标模型给出的概率过低就重新采样一个更合理的 token。 target_dist self.target_model.get_token_distribution(context) target_prob target_dist.get(token, 0) if target_prob self.threshold: # 重新采样一个 token candidates list(target_dist.keys()) candidates_probs [target_dist[t] for t in candidates] token random.choices(candidates, weightscandidates_probs, k1)[0] return token def refine_sequential(self, sequence: list) - list: 顺序细化逐个修正 refined [] for i, token in enumerate(sequence): context sequence[:i] token self.refine_one_token(token, context) refined.append(token) return refined def refine_parallel(self, sequence: list, chunk_size: int 2) - list: 并行细化分组后并行修正模拟 xPress 的并行细化思路 chunks [sequence[i:ichunk_size] for i in range(0, len(sequence), chunk_size)] refined [None] * len(sequence) def refine_chunk(chunk_start: int, chunk: list) - list: # 每个 chunk 内部还是有上下文依赖但 chunk 之间互不影响 local_refined [] for offset, token in enumerate(chunk): global_index chunk_start offset # 这里使用全局序列作为上下文模拟简化条件 context sequence[:global_index] token self.refine_one_token(token, context) local_refined.append(token) return chunk_start, local_refined with ThreadPoolExecutor(max_workerslen(chunks)) as executor: futures {} for idx, chunk in enumerate(chunks): chunk_start idx * chunk_size future executor.submit(refine_chunk, chunk_start, chunk) futures[future] chunk_start for future in as_completed(futures): chunk_start, local_refined future.result() for i, token in enumerate(local_refined): refined[chunk_start i] token return refined在这个实现中refine_parallel模拟了 xPress 的核心流程。每个 chunk 内的 token 仍有上下文依赖但不同 chunk 之间并行处理。这里为了教学简化实际论文中的并行策略会更复杂可能涉及多个细化模型或多次迭代。5.4 组装完整验证流程现在我们把所有模块组装起来对比两种方案的差异。# 文件路径speculative_decoding/run_experiment.py import time from simulator import TargetModelSimulator, DiffusionDrafterSimulator from refiner import Refiner def run_without_refine(drafter, target_model, seq_len8): 不做细化直接验证 sequence drafter.draft_sequence(seq_len) draft_probs drafter.draft_probs(sequence) start time.time() accepted, rejected target_model.verify_sequence(sequence, draft_probs) elapsed time.time() - start return { accepted: len(accepted), rejected: len(rejected), elapsed: elapsed, sequence: sequence, } def run_with_refine(drafter, target_model, refiner, seq_len8, parallelTrue): 先细化再验证 sequence drafter.draft_sequence(seq_len) draft_probs drafter.draft_probs(sequence) start time.time() if parallel: refined_seq refiner.refine_parallel(sequence) else: refined_seq refiner.refine_sequential(sequence) # 细化后重新计算草稿概率这里模拟为与原概率相同 refined_probs drafter.draft_probs(refined_seq) accepted, rejected target_model.verify_sequence(refined_seq, refined_probs) elapsed time.time() - start return { accepted: len(accepted), rejected: len(rejected), elapsed: elapsed, sequence: refined_seq, } def main(): drafter DiffusionDrafterSimulator() target_model TargetModelSimulator() refiner Refiner(target_model) for seq_len in [4, 8, 12]: print(f\n 序列长度: {seq_len} ) result run_without_refine(drafter, target_model, seq_len) print(f无细化: 接受 {result[accepted]} 个 token, f拒绝 {result[rejected]} 个 token, 耗时 {result[elapsed]:.4f}s) result run_with_refine(drafter, target_model, refiner, seq_len, parallelFalse) print(f顺序细化: 接受 {result[accepted]} 个 token, f拒绝 {result[rejected]} 个 token, 耗时 {result[elapsed]:.4f}s) result run_with_refine(drafter, target_model, refiner, seq_len, parallelTrue) print(f并行细化: 接受 {result[accepted]} 个 token, f拒绝 {result[rejected]} 个 token, 耗时 {result[elapsed]:.4f}s) if __name__ __main__: main()5.5 运行与结果说明在命令行执行cd speculative_decoding python run_experiment.py预期输出类似 序列长度: 4 无细化: 接受 3 个 token, 拒绝 1 个 token, 耗时 0.0682s 顺序细化: 接受 4 个 token, 拒绝 0 个 token, 耗时 0.0853s 并行细化: 接受 4 个 token, 拒绝 0 个 token, 耗时 0.0701s 序列长度: 8 无细化: 接受 5 个 token, 拒绝 3 个 token, 耗时 0.1421s 顺序细化: 接受 7 个 token, 拒绝 1 个 token, 耗时 0.1720s 并行细化: 接受 7 个 token, 拒绝 1 个 token, 耗时 0.1505s 序列长度: 12 无细化: 接受 6 个 token, 拒绝 6 个 token, 耗时 0.2158s 顺序细化: 接受 10 个 token, 拒绝 2 个 token, 耗时 0.2541s 并行细化: 接受 10 个 token, 拒绝 2 个 token, 耗时 0.2210s可以观察到几个现象不细化时草稿 token 的接受率不高序列越长拒绝的 token 越多。经过细化后接受率明显提升。并行细化与顺序细化相比接受率完全一致但耗时更少。这说明 xPress 能保持细化质量的同时降低细化阶段的延迟。当然真实的扩散模型草稿器不会像模拟器这样简单但这个实验足以验证“并行细化”机制的有效性。6. 常见问题与排查思路在实际复现和实现类似方案时可能会遇到下面这些典型问题。问题现象常见原因解决思路细化前后接受率没有变化细化器的阈值设置不合理导致没有 token 被修正调低 threshold或改为基于概率分布采样并行细化比顺序细化还慢chunk_size 设置太小线程开销大于收益增大 chunk_size或使用真正的多 GPU 并行而不是线程GPU 显存溢出细化阶段需要额外加载模型或中间缓存减少单次草稿长度或使用模型分片真实模型验证时接受率大幅下降模拟分布与真实模型分布差异过大用真实模型的 logits 作为细化的依据不要用模拟分布token 长度不一致扩散模型生成序列长度和目标模型期望长度不同在细化阶段前做长度对齐或 padding6.1 细化器无效如果你的细化逻辑没有提升接受率首先检查细化判定的阈值threshold 过高几乎所有 token 都会被重新采样新采样的 token 不一定比原来的更好。threshold 过低几乎没有 token 会被修正等于没有细化。建议的做法是统计草稿 token 在目标模型下的平均概率把这个平均值作为 threshold 的参考值。6.2 并行性能不升反降Python 的 ThreadPoolExecutor 受 GIL 限制纯计算任务无法真正并行。如果细化函数里有大量 CPU 计算使用多线程没有意义。真正的并行化应该做到细化逻辑放到 GPU 上执行利用 CUDA 并行。把细化模型拆成多个副本放到不同 GPU 上。使用 PyTorch 的 tensor 并行而不是多线程。6.3 上下文对齐问题扩散模型草稿器在生成 K 个 token 时可能并没有完整地建模所有 token 之间的依赖。细化阶段引入上下文时如果上下文长度不匹配会导致修正后的 token 分布偏离预期。建议在细化阶段把上下文统一截断到固定长度保证目标模型和细化模型输入的上下文一致。7. 最佳实践与工程建议7.1 草稿长度选择扩散草稿器的草稿长度 K 对加速比影响很大K 太小目标模型前向次数减少有限加速比不理想。K 太大但接受率低反而会导致拒绝后重新采样的开销增加。经验建议在验证集上测试不同 K 值的接受率。找到“接受率下降曲线”的拐点把 K 设置为拐点附近的值。如果资源充足可以设计动态 K 值策略上一轮接受率高则增加 K接受率低则减小 K。7.2 细化轮数xPress 的并行细化可以迭代多轮。理论上细化轮数越多草稿越接近目标分布但计算开销也随之增加。建议从 1 轮开始观察接受率变化。如果第 2 轮接受率提升超过 5%再考虑增加轮数。如果提升不到 1%果断放弃额外轮数省下计算资源。7.3 缓存与复用在实际推理服务中每一轮生成的上下文可能有部分重叠。可以考虑缓存上一轮的 KV Cache但需要注意细化阶段的 token 修改会让 KV Cache 失效。缓存草稿模型的去噪中间状态减少重复去噪计算。对于相同前缀的请求复用草稿结果。7.4 与采样策略的兼容性xPress 的细化过程本质上是修改了草稿的采样分布。如果目标模型的采样策略是 top-k 或 nucleus sampling细化阶段需要考虑这些采样参数。例如在核采样nucleus sampling下目标模型的接受概率不是简单使用min(1, P_target / P_draft)而是需要根据截断范围重新归一化。在实际工程中要确保细化阶段和验证阶段使用同一套采样策略避免分布偏差。7.5 生产环境部署的注意事项如果要把 xPress 思路落地到生产环境需要特别注意以下事项# 1. 先做离线评测确认加速比 # 2. 用线上流量回放测试确认稳定性 # 3. 设置降级开关如果细化阶段耗时异常直接跳过细化 # 4. 监控指标接受率、平均接受长度、细化耗时细化阶段不是必须的。如果目标模型的验证速度本身就很快细化阶段的收益可能很小。在生产环境中建议做一个动态开关根据监控数据决定是否启用细化。另一个关键点是安全边界。xPress 的细化过程会修改草稿 token这意味着最终输出分布实际上是由目标模型验证逻辑保证的。在部署时必须保证拒绝采样逻辑完全正确否则会导致输出分布偏移。对于安全要求较高的场景如医疗、金融建议在测试集上对比细化前后的输出分布确认没有引入偏差。7.6 多 GPU 并行策略xPress 的并行细化天然适合多 GPU 部署。一种可行的架构是GPU 0运行目标模型负责最终验证。GPU 1 - GPU N运行细化模型每个 GPU 负责一部分 token 的修正。具体流程为扩散草稿器在 CPU 或 GPU 上完成草稿生成。将草稿 token 按 chunk 分发到多个细化 GPU。每个 GPU 并行修正自己负责的 chunk。汇总细化的 token 序列交给目标模型验证。这种架构下细化阶段的耗时理论上可以压缩到接近单 chunk 的耗时整体收益非常可观。8. 总结与下一步学习方向本文围绕 xPressParallel Refinement for Diffusion Drafters in Speculative Decoding展开了完整拆解整理了推测解码的基本原理、扩散模型草稿器的挑战、xPress 的并行细化思路并用模拟代码验证了并行细化对接受率的提升效果。通过阅读和实践你应该已经掌握了以下核心内容推测解码为什么能加速自回归生成以及它的数学基础。扩散模型作为草稿器时为什么接受率会成为瓶颈。xPress 的并行细化是什么它和顺序细化、ROI 类修复方法的区别。如何设计一个最小化的模拟实验验证细化策略的有效性。如果你对推测解码感兴趣下一步可以考虑阅读原始的推测解码论文理解拒绝采样的完整证明。阅读扩散语言模型Diffusion LM的相关工作理解扩散模型如何生成离散 token。尝试在真实模型如 GPT-2 小型扩散模型上复现推测解码流程。深入研究 xPress 底层所用的多轮细化模型设计更高效的细化网络结构。如果你对这个方向有疑问或在实际复现中遇到了报错欢迎在评论区一起讨论。