Skip to content

SwiftInfer:基于TensorRT的流式推理加速框架

内容整理自学习笔记,仅供面试备考参考;不构成录用、培训或考试承诺。

1. 什么是 SwiftInfer?

SwiftInfer 是 Colossal-AI 团队发布的基于 TensorRT 的 StreamingLLM 实现方案,旨在进一步提升大模型多轮对话推理的性能,为大模型推理提供低成本、低延迟、高吞吐的最佳实践。

  • 基础方法:StreamingLLM(Efficient Streaming Language Models with Attention Sinks
  • 论文地址https://arxiv.org/pdf/2309.17453.pdf
  • 核心提升:在原始 StreamingLLM 基础上最多带来额外 46% 的推理吞吐速度提升

2. 为什么需要 SwiftInfer?

StreamingLLM 虽然解决了无限长度流式输入的问题,但其使用原生 PyTorch 实现,在多轮对话推理场景的落地应用中仍有优化空间:

需求原生 StreamingLLMSwiftInfer
低成本一般优化(TensorRT 加速)
低延迟一般优化(KV Cache 重实现)
高吞吐一般优化(推理引擎集成)

3. SwiftInfer 的核心思路

3.1 KV Cache 机制

假设窗口大小为 10 个 token,随着生成的 token 增加:

KV Cache 窗口示意:

初始状态:  [S1][S2][T1][T2][T3][T4][T5][T6][T7][T8]
           ↑ 蓝色: 始终保留的初始token
                    ↑ 黄色: 中间动态token

生成新token后:
          [S1][S2][T2][T3][T4][T5][T6][T7][T8][T9]
                              T1被踢出,T9新增

继续生成:
          [S1][S2][T3][T4][T5][T6][T7][T8][T9][T10]
                              T2被踢出,T10新增

3.2 位置信息重注入

由于黄色方块(中间 token)的位置会发生变化,在计算注意力时需要重新注入位置信息。

注意:StreamingLLM 不会直接提高模型能访问的上下文窗口,而是能够在支持流式超多轮对话的同时保证模型的生成效果。

4. SwiftInfer 的优点

4.1 性能提升

指标数值
流式输入 token 上限超过 400 万
速度提升(vs 滑动窗口重计算)22.2 倍
额外推理吞吐提升最多 46%

4.2 技术优势

  1. StreamingLLM + TensorRT 结合:不仅拥有原始 StreamingLLM 的所有优点,还具有更高的运行效率
  2. KV Cache 重实现:重新实现了 KV Cache 机制以及带有位置偏移的注意力模块
  3. 生态兼容:TensorRT-LLM 团队也在同期对 StreamingLLM 进行了类似支持

5. StreamingLLM 篇:背景知识

5.1 为什么需要 StreamingLLM?

大语言模型在多轮对话场景中面临以下核心问题:

  1. 上下文长度限制:大模型能够记住的上下文长度有限,影响与用户互动的质量
  2. 预训练窗口限制:LLM 在预训练期间只能在有限的注意力窗口下训练
  3. KV Cache 内存膨胀:多轮对话场景下,key 和 value 缓存消耗大量内存,无法在有限显存下无限扩展上下文
  4. 泛化能力不足:二次微调后的模型无法很好地泛化到比训练序列长度更长的文本

5.2 StreamingLLM 核心思路

通过观察注意力模块中 Softmax 的输出,发现了 attention sink 现象:

  • 注意力机制会为每一个 token 分配注意力值
  • 文本最初的几个 token 总是被分配到很多无用的注意力
  • 当这几个 token 被踢出滑动窗口时,模型生成效果迅速崩溃
  • 只要保留这几个 token 在窗口内,模型就能稳定生成高质量文本

5.3 注意力机制对比

机制计算复杂度生成效果无限长度
Dense Attention不支持
Window Attention差(溢出崩溃)不支持
Sliding Window w/ Re-computing支持(慢)
StreamingLLM支持

提示:StreamingLLM 不需要重新训练模型,可直接兼容主流大语言模型并改善推理性能。

6. 总结

SwiftInfer 作为 StreamingLLM 的工程化实现,通过结合 TensorRT 推理优化技术,为大模型多轮对话推理提供了低成本、低延迟、高吞吐的最佳实践方案。其核心价值在于将学术研究成果转化为可实际部署的高性能推理引擎。