如何将具有1M Tokens的8B LLM的TTFT优化到20秒

免责声明：以下内容代表我个人观点，不代表我公司或团队的观点。此外，本文不讨论长上下文LLM与RAG的比较优势。

如果您旨在将8B模型在100万个token下的TTFT优化到20秒，您可以考虑以下方法：

• A. 使用门控线性RNN或SSM（例如Mamba-2、RetNet、RWKV、DeltaNet）、稀疏注意力（例如SW、BigBird、LongNet）或线性注意力；
• B. 使用混合模型（例如Jamba、Samba）；
• C. 扩展到65+个A100或H100 GPU；
• D. 利用基于内存、基于循环或基于集群的方法来降低注意力复杂度；
• E. 要求用户等待第二轮对话；
• F. 尝试MInference。

相对于FlashAttention，需要多大的加速？

让我们做一个简单的计算。根据fp16 8B Transformer的计算，一个A100的fp16 TFLOPs = 312 TFLOPs。

一个A100的理论TTFT可计算如下：

然而，考虑到Transformer计算中的内核优化、读写、同步等待和内存移动，实际的TTFT会慢约1.6倍-2倍（取决于框架，TensorRT-LLM > vLLM > HF）。

因此，A100的估计上限 $$TTFT=14.52mins\ast 1.5x=21.78minsTTFT\; =\; 14.52\; mins\; *\; 1.5x\; =\; 21.78\; mins$$

因此，在不考虑TP和Sequence Parallel通信成本，特别是节点间通信成本的情况下，至少需要65+个A100 GPU。

如果您碰巧拥有大量财政资源，您也可以使用20+个H100 SXM GPU来达到同样的效果

（注意：量化需要除以一个系数）

假设您只有8个A100 GPU，实现1M tokens TTFT 20秒的目标需要相对于FlashAttention 8倍的加速。

为了实现这一点，由于预填充阶段是计算密集型的，优化目标相当于将注意力FLOPs减少8倍。

总而言之，使用MInference并利用长上下文注意力中的动态稀疏性，您只需要8个A100 GPU即可实现20+秒的TTFT，并且几乎保持相同的精度，特别是在高度动态的长上下文任务中。

• MInference项目页面：https://aka.ms/MInference
• 代码：https://github.com/microsoft/MInference
• 演示：https://huggingface.co/spaces/microsoft/MInference

如何协同设计算法和系统？

首先，优化的直觉是什么：注意力（Attention）的稀疏性，这源于Softmax和长上下文带来的极端稀疏性。（这在许多工作中都有分析，例如StreamingLLM、SparQ、TriForce。）

目标是利用这种先验稀疏性来设计一种GPU高效且准确召回的稀疏注意力算法。

回顾现有高效长上下文LLMs方法：

• 对于多头注意力（Multi-head Attention），它涉及两个[batch size, head number, seqlen, head dim] Q和K矩阵的逐层矩阵乘法。
• 优化可以在head级别通过聚类或共享进行，在seqlen级别通过token剪枝、稀疏注意力或线性注意力进行，在head dim级别通过低秩或topK进行。

我们专注于基于无需训练的稀疏注意力方法（因为减少比增加更容易）。

在稀疏注意力方法中，静态稀疏注意力（Static Sparse Attention）存在显著的性能损失，并且无法处理像“大海捞针”（Needle In A Haystack）和KV检索等动态任务。这是因为注意力的动态性质使其能够利用其N^2的信息收集能力。

基于检索的稀疏注意力受到注意力极端稀疏性的启发，旨在以最小的开销为每个Q获得topK的K索引，例如在head dim中进行优化（例如SparQ，尽管它仅用于解码阶段）。然而，这种方法对内核不友好，因为GPU上的topK操作可能不如CPU。此外，以最小计算开销实现高topK索引召回仍然是一个挑战，这使得它难以扩展到长上下文。

因此，设计一个对内核友好的动态稀疏注意力方法包括：

1. 确保稀疏模式具有空间局部性和大的块大小，以利用Tensor核。
2. 以最小的开销在线确定和构建动态稀疏索引。
3. 确保LLMs原始注意力模式中的显著稀疏性。

基于这些原则，我们提出了MInference，一种无需训练且对内核友好的动态稀疏注意力方法，它以几乎无损的精度显著加速长上下文LLMs。

得益于我们之前在稀疏注意力优化方面的工作，特别是动态稀疏注意力（例如PIT），我们实现了对1M tokens LLMs TTFT高达10倍的端到端加速。

我想强调PIT的强大特性，它在LLM社区中经常被忽视：

1. MoE的稀疏加载和计算，与MegeBlocks同时代。
2. 有效解决/加速RLHF或SFT填充问题，提供了比TurboTransformer更好的解决方案。
3. Deja Vu或PowerInfer中动态稀疏FFN计算的解决方案。

有关MInference的更多详细信息，请参阅Yucheng的博客或我们的论文。

在这里，我想讨论一些论文中未提及的见解。

动态稀疏注意力是未来吗？

我们不确定它是否是未来，因为MInference在短上下文中的性能可能无法超越现有的vLLM或TensorRT-LLM。然而，我们确信它是一种加速50K-1M上下文LLM的有效且可行的解决方案。

首先，我们已经在同期工作中看到了类似的想法，这证实了该方向的可行性。其次，MInference表现出良好的泛化能力，仅需一个样本即可搜索最佳配置，显示出卓越的跨任务和跨长度泛化能力。此外，它在各种模型上表现良好，包括LLaMA-3-8B-1M、GLM-4-9B-1M、Yi-200K、Phi-3-mini-128K和Qwen2-7B-128K。我们还收到了LLM提供商在内部模型上测试后的积极反馈。

哪种模式最重要，为什么？

斜线模式是三种模式中最重要的一种。尽管它在BERT时代就被发现，但由于加速困难而未能得到充分利用。

这种模式不仅与RoPE有关；我们认为它在注意力中充当信息通道，专注于等距离信息。

这种方法是否与KV缓存压缩冲突，或者能否用于KV缓存复用场景？

MInference与KV缓存压缩方法是正交的。我们进行了将其与SnapKV结合的实验，结果优于单独使用SnapKV。我们还测试了多轮对话场景，MInference在大多数任务中表现良好。我们将很快更新这些结果。

如何评估长上下文LLM的能力？

评估应包括LLMs在检索、通用任务如问答（QA）、摘要（Summarization）、代码和数学推理方面的表现。检索是目前不同方法之间性能差异显著的领域。在难度方面，KV检索 > 大海捞针 > 检索数字 > 检索密码。对于后三者，LLMs可以利用SFT（监督微调）预先感知语义变化和潜在问题以获得更好的性能，而KV检索由于其动态性质仍然具有挑战性。

MInference、SSM、线性注意力（Linear Attention）和稀疏注意力（Sparse Attention）有什么区别？

一些门控线性RNN可以等效于带有KV缓存压缩的稀疏注意力，但它们压缩提示词（prompt tokens）的能力受限于先前稀疏模式中的归纳偏置。

使用SSM或SW进行预填充，并使用完全注意力进行解码，这会是一个解决方案吗？

直观地说，将SSM或SW与密集解码相结合，在正常推理中可能表现不佳。主要问题是先前稀疏模式造成的信息损失显著。此外，从头开始的稀疏注意力可能会覆盖一些动态模式，这需要进一步的实验验证。

如何优化解码中的KV缓存？

这是一个很大的话题，我们将在未来的文章中讨论（待办）。

欢迎随时开启讨论或联系我以获取更多见解。Yucheng @liyucheng 和我下周也会参加ICML，欢迎前来与我们讨论高效方法和相关话题。

• 微软展位：7月23日，上午10点-上午11点
• ES-FoMo：7月26日，下午1点-下午2:15
• LCFM：7月26日，上午9:45-上午10点