# 在谷歌TPU上实现3倍加速：UCSD利用扩散式推测解码优化LLM推理 - 来源：Google Developers Blog（RSS） - 发布时间：2026-05-04 00:00 - AIHOT 分数：66 - AIHOT 标记：精选 - AIHOT 链接：https://aihot.virxact.com/items/cmorfwui900bbslahbt29u50b - 原文链接：https://developers.googleblog.com/supercharging-llm-inference-on-google-tpus-achieving-3x-speedups-with-diffusion-style-speculative-decoding ## 精选理由 把扩散式的 speculative decoding 在 TPU 上跑出了 3 倍推理加速，峰值快到 EAGLE-3 的两倍，还直接集成了 vLLM，做推理优化的赶紧试一下。 ## AI 摘要 加州大学圣地亚哥分校的研究团队在谷歌TPU上成功部署了DFlash，一种基于块扩散的推测解码方法。该方法突破传统自回归草稿生成的序列性瓶颈，通过单次前向传播并行“绘制”整个候选令牌块，而非逐个预测。系统平均实现了3.13倍的推理加速，峰值性能接近EAGLE-3等现有方法的两倍。这一开源方案已集成至vLLM生态系统，通过利用“免费”的并行验证能力和针对复杂推理任务的高质量草稿预测，显著优化了TPU硬件的利用效率。 ## 正文 2026年5月4日 于伟仁 穆亚荣 高级首席软件工程师 Google Cloud 冉立豪 Google Cloud 冯兆祥 研究助理 加州大学圣迭戈分校 赵一鸣 研究助理 加州大学圣迭戈分校 张昊 助理教授 加州大学圣迭戈分校 当前大语言模型加速领域由自回归推测解码主导，其流程是轻量级草稿模型在目标验证之前顺序预测模型 token。然而，这种串行草稿生成方式带来了一个根本性的执行瓶颈：生成 K 个候选模型 token 需要 K 次顺序前向传播。这种逐步依赖关系迫使系统必须等待每个模型 token 预测完成后才能开始下一个，从根本上限制了草稿阶段的加速潜力。为了突破这一效率天花板，研究人员开始从逐 token 草稿生成转向块扩散，这是一种范式转变，能够以 O(1) 单次前向传播生成整个候选模型 token 块。 我们很自豪能够支持那些推动 AI 硬件边界的外部研究人员。今天，我们非常高兴地介绍来自加州大学圣迭戈分校研究人员的一个重要开源里程碑，该项目由张昊领导，他是分页注意力机制以及预填充/解码分离服务的共同发明人。他们成功地在 Google TPU 上实现了块扩散推测解码（即 DFlash，一种由加州大学圣迭戈分校 Z Lab 的刘志坚、陈健等人开发的更优的扩散式推测解码）。 通过将这种新颖架构直接集成到开源 vLLM TPU 推理生态系统中，加州大学圣迭戈分校团队在 TPU v5p 上实现了平均每秒模型 token 数 3.13 倍的提升，在复杂数学任务上峰值加速比接近 6 倍。在 TPU v5p 上进行的 DFlash 与 EAGLE-3 直接服务对比中，DFlash 实现了 2.29 倍的端到端服务加速，几乎是 EAGLE-3 的 1.30 倍性能提升的两倍。 以下是加州大学圣迭戈分校研究人员提供的技术深度解析，详细介绍了他们如何构建该系统、性能基准测试结果，以及这对 Google TPU 生态系统未来的意义。 克服自回归瓶颈 标准大语言模型推理采用自回归方式生成文本。这意味着模型每生成一个 token 都需要执行一次完整的前向传播，这严重低估了 TPU 等 AI 加速器的大规模并行计算能力，尤其是在批量较小时。 推测解码通过使用一个更小、更高效的"草稿"模型（或机制）来同时预测多个未来 token，从而缓解这一问题。随后，更大的"目标"模型在单次并行前向传播中验证这些草稿 token。如果草稿 token 准确，系统就能以单步成本接受多个 token，大幅降低延迟。 然而，推测解码的潜力往往受限于草稿模型本身。现有大多数方法依赖自回归草稿机制来顺序生成候选 token。这意味着，虽然目标模型的验证是并行的，但草稿阶段仍受限于 O(K) 步串行操作。结果，"猜测" token 所花费的时间开始侵蚀验证所节省的时间，从而限制了实际加速潜力。 谷歌 TPU 上的扩散式草稿生成 扩散大语言模型（dLLM）通过用块扩散机制取代这一顺序过程，从根本上改变了游戏规则。dLLM 不是猜测下一个词，而是"绘制"整个块。一种基于 dLLM 的显著草稿生成方法是 DFlash。通过利用从目标模型中提取的隐藏特征，DFlash 能在单次前向传播中生成整个草稿 token 块。这种从 O(K) 到 O(1) 的复杂度转变，将草稿延迟降低到几乎可以忽略的水平，使其成为 TPU 高带宽矩阵乘法单元（MXU）的理想架构匹配。 加州大学圣地亚哥分校的研究团队将 DFlash 集成到了 vLLM TPU 推理框架中。DFlash 是一种新颖的推测解码方法，它利用块扩散机制来提出具有极高接受长度（T）的草稿 token。 在 Google TPU 上实现这一方案需要深度优化。在 Google Cloud 工程师的架构指导下，UCSD 团队将开销降至最低，确保内存带宽和矩阵乘法单元得到充分利用。通过将 DFlash 提议器与验证流水线高效映射到 TPU 架构上，他们既最小化了草稿阶段的开销，又最大化了目标模型的并行验证吞吐量。 将 DFlash 引入 TPU/JAX 将 DFlash 从原始的 GPU/PyTorch 实现移植到 Google TPU/JAX AI 栈生态系统，并非简单的代码翻译；它需要重新设计系统，以适配 TPU 独特的架构优势。以下是 UCSD 团队如何攻克三大主要技术难题。 用于注意力机制的“双缓存”方案 在 PyTorch 环境中，DFlash 依赖简单、动态的 KV 管理。然而，通过 tpu-inference 实现的高性能 TPU 服务采用基于 Pallas 内核的分页注意力机制——这种机制将内存划分为固定大小的页面，以最大化效率。 问题在于？DFlash 的非因果块扩散——正是它能够“绘制”一个 token 块的关键特性——从根本上与标准分页注意力机制不兼容。为了解决这个问题，研究人员设计了一种双缓存架构。目标模型继续使用分页 KV 缓存，确保其受益于大规模服务所需的高性能 Pallas 内核。草稿模型则使用一条专用路径，采用静态设备端 JAX 数组，成功复现了原始 DFlash 的设计，同时保持了 TPU 原生性能。 智能上下文管理 DFlash 的独特之处在于，草稿模型是“目标条件化”的——它通过观察目标模型的中间推理步骤来保持智能。这些“隐藏状态”存储在一个随时间增长的上下文缓冲区中。 为了尽可能加快主机 CPU 与 TPU 加速器之间的通信速度，该团队实现了一种 2 的幂次填充策略。这确保了当新投影的特征被追加到缓冲区时，它们能以优化后的数据块进行传输。通过精确追踪草稿模型已经“消耗”了多少上下文，他们避免了任何重复处理或数据丢失，从而保证了并行草稿生成的高度准确性。 弥合 TPU 推理中的元数据鸿沟 与标准的草稿生成方法不同，DFlash 是独特的“有状态”方法，它依赖于跨迭代的持久化状态（包括上下文缓冲区、KV cache 位置和 RoPE 偏移量）来维持其并行块预测。在 TPU 优化的 vLLM 流水线中，转发给提议模型的元数据包含了当前正在验证的草稿 token。虽然这对大多数模型来说是标准做法，但对于基于扩散模型的架构，这导致了“序列长度膨胀”——一种内部草稿状态偏离目标模型真实状态的错位问题。 通过重新设计提议模型，使其严格与真实接受的 token 数量同步，研究团队恢复了两个模型之间的完美对齐。这一调整使得块扩散逻辑能够在 TPU 硬件上以完全的数学精度运行，从而实现了最终结果中看到的显著加速。 为 TPU 服务的未来进行基准测试 正面交锋：DFlash 与 EAGLE-3 在 TPU v5p 上的对决 为确保严格且公平的比较，加州大学圣地亚哥分校的研究人员将 DFlash 与当前 TPU 上的主流推测解码方法 EAGLE-3 进行了基准测试。在这项对比研究中，研究人员对两者使用了完全相同的硬件（TPU v5p）和相同的目标模型（Llama-3.1-8B）。 DFlash 与 Eagle3 在 vLLM TPU 流水线及 v5p TPU 上的对比 模型：Llama-3.1-8B-Instruct（目标模型）+ z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat（DFlash 草稿模型，K=10）/ yuhuili/EAGLE3-LLaMA3.1-Instruct-8B（EAGLE-3 草稿模型，K=2） 这种配置代表了两种方法最实用的部署场景，因为 K 值的选择基于它们各自官方开源检查点，开箱即用，无需额外的微调或重新配置。像 EAGLE-3 这样的自回归草稿模型会产生顺序延迟惩罚，该延迟随 K 值线性增长，这通常将其限制在较小的推测预算内，以保持较低的每 token 延迟。相比之下，DFlash 使用并行块扩散在单次前向传播中预测所有 token，使得草稿生成成本在很大程度上对 K 值不敏感。结果很明确：DFlash 实现了 2.29 倍的加速，而 EAGLE-3 仅提供了 1.30 倍的提升。在 mbpp 等编码任务上，DFlash 将每个 token 的生成时间从 9.81 毫秒压缩到 3.48 毫秒，提升了 2.83 倍。 为什么差距如此之大？EAGLE-3 每一步自回归地预测 2 个 token，需要顺序的前向传播，并且每一步之间都有 Python 编排开销。而 DFlash 在单次前向传播中生成一个包含 10 个高质量候选 token 的块，完全消除了这种串行瓶颈。在 TPU 上，这种“高质量、高数量”的草稿输出直接转化为更高的平均接受长度，将 TPU 巨大的计算潜力转化为实际的服务吞吐量。 TPU v5p 上的基准测试结果 为了评估 DFlash 在 Google TPU 上的影响，加州大学圣地亚哥分校团队在 TPU v5p 上跨多个领域对其实现进行了基准测试，重点关注复杂推理、数学和编码——这些领域的长上下文生成通常受高延迟困扰。 加州大学圣地亚哥分校团队构建了一个独立的 JAX 基准测试来评估 DFlash 的结果。通过剥离服务层的开销，他们能够隔离 DFlash-on-TPU 算法的原始性能。他们观察到所有数据集上的平均加速比为 3.13 倍，在数学推理方面达到了显著的峰值。 v5p TPU 上的基准测试结果 模型：Qwen/Qwen3-4B（目标模型）+ z-lab/Qwen3-4B-DFlash-b16（草稿模型，K=16）；贪婪解码 对于像 math500 这样严谨的数学任务，DFlash 将生成时间从每个 token 8.02 毫秒降低到了每个 token 1.40 毫秒。在 humaneval 等代码评估中，生成速度提升了超过 3.5 倍。 投机效率的深度洞察 “K-Flat”突破：为何更宽是免费的 在优化过程中，研究团队发现了一个硬件特性，它改变了工程师对投机限制的思考方式：K-Flat 验证。 在 TPU v5p 这类数据中心级加速器上，他们的系统性实验揭示了一个令人惊讶的现实：验证 1024 个 token 的成本与仅验证 16 个 token 的成本几乎相同。这一现象之所以发生，是因为在高端硬件上，时间消耗主要来自加载模型权重，而非针对这些序列长度的注意力机制的原始数学计算。换句话说，硬件的计算天花板如此之高，以至于检查一个长得多的“猜测”所带来的额外工作基本上是免费的。 这一发现改变了整个研究前沿。它证明了投机解码的瓶颈并非“验证成本”，而是“草稿质量”。知道更宽的块在计算上是免费的，这使得开发者能够大胆地扩展草稿块的大小，利用更丰富的双向上下文来提高准确性，而无需担心拖慢硬件速度。 扩展理论：质量胜于数量 虽然数据中心级 AI 加速器使得增加块大小（K）几乎“免费”，但他们的扩展理论揭示，仅仅增加更多 token 会带来递减的回报。在他们当前的操作点上，块大小 K=16 已经捕获了超过 90% 的理论最大加速比。事实上，将 K 从 16 一直扩展到 128，每一步可能净增不到一个被接受的 token。 性能的真正杠杆在于质量而非数量。他们的分析表明，提升每个位置的接受概率（a）比增大块大小 K 的价值高出 2-3 倍。这改变了研究重点：在验证成本恒定的环境下，主要瓶颈不再是系统能检查多少个模型 token，而是系统能多准确地预测它们。大语言模型服务的下一个前沿在于更智能的草稿训练，而不仅仅是更宽的推测窗口。 可预测性因素：任务驱动的加速 接受概率与任务的可预测性密切相关。团队观察到一种自然的“位置衰减”现象，即块末尾的模型 token 比开头的更难猜测。在数学和编程等逻辑驱动领域，这种衰减非常缓慢，即使在块深处也能保持高接受率。然而，对话式聊天则更具随机性，在前几个模型 token 之后准确率就会急剧下降。 这种可预测性直接驱动了加速。由于结构化推理能产生更可预测的序列，数学和代码任务允许更长的已接受块，从而更有效地利用 TPU 的并行验证能力。因此，DFlash 在数学推理中获得了最高的增益，其次是编程，而对话式任务的提升则相对温和。 与 vLLM 的开源集成 此次合作的核心原则之一是丰富开源生态系统。完整的实现并未作为内部研究原型保留，而是已提交至 vLLM tpu-inference 仓库，包括： PR #1868：DFlash 模型与提议者架构。 PR #1869：用于推测性解码的端到端流水线集成。 PR #1870：全面的 CI 和端到端测试框架。 UCSD 团队正在积极添加一个 torchax 提议者，以便 DFlash 也能在 PyTorch 服务路径上运行。 拓展推测系统的前沿 这一里程碑为谷歌 TPU 下一波创新奠定了基础。通过利用 DFlash 独特的并行采样技术，他们正在为推测性解码（SSD）铺平道路，使用推测缓存来大幅降低高吞吐环境中的延迟。为了捕捉更丰富的上下文并提高复杂推理的接受率，他们计划使用 TPU RL 栈 Tunix 和 MaxText 扩展到更宽的草稿块。此外，新开发的高性能 JAX 内核为支持基于扩散的目标模型提供了基础，使 vLLM-TPU 生态系统始终处于高效非自回归生成的最前沿。 您可以通过 Colab 笔记本查看相关的技术报告和实现细节，或直接深入 vLLM GitHub 仓库中的代码。 研究提案征集 这项工作得益于 TPU 构建者计划，该计划体现了我们的使命，即通过提供高性能硬件和谷歌云积分来赋能学术界和开源社区。如果您有兴趣将 TPU 用于研究、教学或开源开发，我们期待您的来信！请发送电子邮件至 tpu-builders-support@google.com 与我们联系。 致谢：衷心感谢加州大学圣地亚哥分校的研究团队，包括 Zhongyan Luo、Son Nguyen、Andy Huang。特别感谢 Kyuyeun Kim、Brittany Rockwell、Chris Chan、Mitali Singh、Yixin Shi 和 Gang Ji 的团队与研究团队合作完成 PR，以及 Josh Gordon、Edgar Chen、Aditi Joshi、Shubha Rao、Mani Varadarajan、Joe Pamer、Fenghui Zhang、Hassan Sipra 和 Bill Jia 对 TPU 构建者计划研究合作伙伴关系的坚定支持和投入。 AI 云 教程 案例研究 解决方案 学习 探索 影响力