Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free Token Importance Estimation

📄 arXiv: 2502.02789v2 📥 PDF

作者: Jingyu Liu, Beidi Chen, Ce Zhang

分类: cs.CL, cs.AI

发布日期: 2025-02-05 (更新: 2025-05-19)

备注: Proceedings of the 42nd International Conference on Machine Learning (ICML 2025)

💡 一句话要点

SpecPrefill：一种轻量级、免训练的token重要性估计方法，加速LLM的TTFT。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 LLM推理 首个token生成时间 TTFT优化 token重要性估计

📋 核心要点

  1. 现有LLM推理引擎在优化TTFT方面面临挑战，因为其性能瓶颈已从自注意力机制转移到MLP部分。
  2. SpecPrefill通过轻量级模型推测prompt中局部重要的token，并将其输入主模型，从而加速TTFT。
  3. 实验表明，SpecPrefill在真实下游任务中显著提升了LLM的QPS和TTFT，例如Llama-3.1-405B-Instruct-FP8。

📝 摘要（中文）

本文提出SpecPrefill，一个免训练框架，旨在加速大型语言模型（LLM）推理中的首个token生成时间（TTFT），适用于长、中等上下文查询。SpecPrefill基于以下洞察：LLM具有足够的泛化能力，即使只使用精心挑选的prompt token子集也能保持质量。其核心是利用一个轻量级模型来推测上下文中局部重要的token。这些token连同必要的位置信息被发送到主模型进行处理。通过多样化的任务评估和全面的基准测试，包括真实端到端设置和消融研究，验证了SpecPrefill的性能提升。SpecPrefill成功地使Llama-3.1-405B-Instruct-FP8在实际下游任务中实现了高达7倍的最大端到端QPS，以及7.66倍的TTFT提升。

🔬 方法详解

问题定义：论文旨在解决大型语言模型推理过程中，首个token生成时间（TTFT）过长的问题。现有的优化方法大多集中在自注意力机制上，而忽略了MLP部分带来的瓶颈。因此，如何在计算资源有限的情况下，有效降低TTFT，成为一个重要的挑战。

核心思路：论文的核心思路是，并非prompt中的所有token都同等重要，LLM具有足够的泛化能力，即使只使用一部分精心选择的token也能保持输出质量。因此，可以通过一个轻量级的模型来预测prompt中哪些token是重要的，然后只将这些token输入到主模型中，从而减少计算量，加速TTFT。

技术框架：SpecPrefill框架包含两个主要部分：轻量级token重要性估计模型和主LLM。首先，轻量级模型接收prompt作为输入，输出每个token的重要性得分。然后，根据得分选择最重要的token子集，并保留它们的位置信息。最后，将这些选定的token和位置信息输入到主LLM中进行处理，生成后续token。

关键创新：SpecPrefill的关键创新在于提出了一种免训练的token重要性估计方法。与需要额外训练的模型不同，SpecPrefill直接利用LLM自身的知识来评估token的重要性，避免了额外的训练成本和数据依赖。此外，SpecPrefill是一种通用的框架，可以应用于不同的LLM架构。

关键设计：轻量级模型的设计是关键。论文中使用的具体实现未知，但可以推测其目标是计算效率。token重要性得分的计算方式也至关重要，需要能够准确反映token对LLM输出的影响。此外，如何选择token子集，例如选择多少个token，也需要仔细调整，以在性能和质量之间取得平衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SpecPrefill在Llama-3.1-405B-Instruct-FP8模型上取得了显著的性能提升。在真实下游任务中，SpecPrefill实现了高达7倍的最大端到端QPS提升，以及7.66倍的TTFT提升。这些结果表明，SpecPrefill是一种有效的TTFT优化方法，具有很强的实用价值。

🎯 应用场景

SpecPrefill可应用于各种需要快速响应的LLM应用场景，例如在线对话系统、实时翻译、智能客服等。通过降低TTFT，可以显著提升用户体验，提高系统的并发处理能力。此外，该方法还可以降低LLM推理的计算成本，使其更易于部署在资源受限的环境中，例如移动设备或边缘服务器。

📄 摘要（原文）

Improving time-to-first-token (TTFT) is an essentially important objective in modern large language model (LLM) inference engines. Optimizing TTFT directly results in higher maximal QPS and meets the requirements of many critical applications. However, boosting TTFT is notoriously challenging since it is compute-bounded and the performance bottleneck shifts from the self-attention that many prior works focus on to the MLP part. In this work, we present SpecPrefill, a training free framework that accelerates the inference TTFT for both long and medium context queries based on the following insight: LLMs are generalized enough to preserve the quality given only a carefully chosen subset of prompt tokens. At its core, SpecPrefill leverages a lightweight model to speculate locally important tokens based on the context. These tokens, along with the necessary positional information, are then sent to the main model for processing. We evaluate SpecPrefill with a diverse set of tasks, followed by a comprehensive benchmarking of performance improvement both in a real end-to-end setting and ablation studies. SpecPrefill manages to serve Llama-3.1-405B-Instruct-FP8 with up to 7$\times$ maximal end-to-end QPS on real downstream tasks and 7.66$\times$ TTFT improvement.