SwiftInfer

SwiftInfer 项目介绍

项目概述

SwiftInfer 是一个基于 TensorRT 实现的工具，旨在为大规模语言模型（LLM）的推理提供无限输入长度的支持。其核心技术来源于 Streaming-LLM，这是一种利用“注意力水槽”技术防止模型在注意力窗口移动时崩溃的方法。最初，Streaming-LLM 是用 PyTorch 实现的，而 SwiftInfer 则进一步提供了生产级别的 TensorRT 实现。这一实现是构建在最近发布的 TensorRT-LLM 项目基础上的。

快速开始

安装指南

SwiftInfer 依赖于 TensorRT-LLM 的 API 来构建模型和进行推理。由于 TensorRT-LLM 的 API 变化较快，SwiftInfer 将当前版本绑定在 42af740db51d6f11442fd5509ef745a4c043ce51 这次提交上（版本为 v0.6.0）。未来会根据 API 的稳定性进行升级。

如果你已经构建了 TensorRT-LLM V0.6.0，只需执行以下命令即可安装 SwiftInfer：

git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
pip install .

使用 Docker 安装 TensorRT-LLM

通过 Docker，可以按照 TensorRT-LLM 提供的安装指南来安装所需版本。具体命令为：

git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
pip install .

不使用 Docker 安装 TensorRT-LLM

在没有 Docker 的情况下，可以使用自动化脚本安装 TensorRT-LLM。需要确保系统已安装一些基本的包，例如 python、CUDA 工具包、cuDNN 等，且 TensorRT 版本需大于等于 9.1.0，CUDA 工具包版本需大于等于 12.2。

接着执行以下安装操作：

ARCH=$(uname -m)
if [ "$ARCH" = "arm64" ];then ARCH="aarch64";fi
if [ "$ARCH" = "amd64" ];then ARCH="x86_64";fi
if [ "$ARCH" = "aarch64" ];then OS="ubuntu-22.04"; else OS="linux";fi
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/9.1.0/tars/tensorrt-9.1.0.4.$OS.$ARCH-gnu.cuda-12.2.tar.gz
tar xzvf tensorrt-9.1.0.4.linux.x86_64-gnu.cuda-12.2.tar.gz
PY_VERSION=$(python -c 'import sys; print(".".join(map(str, sys.version_info[0:2])))')
PARSED_PY_VERSION=$(echo "${PY_VERSION//./}")
pip install TensorRT-9.1.0.4/python/tensorrt-*-cp${PARSED_PY_VERSION}-*.whl
export TRT_ROOT=$(realpath TensorRT-9.1.0.4)

要保证 nvcc 配置正确，可以通过 nvcc --version 来验证配置。

然后，可以通过以下命令安装 TensorRT-LLM 及 SwiftInfer：

git clone https://github.com/hpcaitech/SwiftInfer.git
cd SwiftInfer
TRT_ROOT=xxx NCCL_ROOT=xxx CUDNN_ROOT=xxx pip install .

运行 Llama 示例

要运行 Llama 模型的示例，需要先克隆 Hugging Face 上的相关仓库，比如 meta-llama/Llama-2-7b-chat-hf 或其他 Llama 变体。然后通过以下命令构建 TensorRT 引擎（需要将 <model-dir> 替换为实际模型路径）：

cd examples/llama

python build.py \
--model_dir <model-dir> \
--dtype float16 \
--enable_context_fmha \
--use_gemm_plugin float16 \
--max_input_len 2048 \
--max_output_len 1024 \
--output_dir ./output/7B-streaming-8k-1k-4-2000/trt_engines/fp16/1-gpu/ \
--max_batch_size 1

下载 MT-Bench 数据：

mkdir mt_bench_data
wget -P ./mt_bench_data https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl

最后，使用以下命令运行 Llama 示例：

python ../run_conversation.py \
--max_input_length 2048 \
--max_output_len 1024 \
--tokenizer_dir <model-dir> \
--engine_dir ./output/7B-streaming-8k-1k-4-2000/trt_engines/fp16/1-gpu/ \
--input_file ./mt_bench_data/question.jsonl \
--streaming_llm_start_size 4 \
--only_n_first 5

运行后，可以看到模型生成的输出。

性能基准

SwiftInfer 已通过与原始 PyTorch 版本的 Streaming-LLM 进行基准测试。使用的硬件包括 Nvidia H800 GPU（80GB）、Intel® Xeon® Platinum 8468 CPU 及 2TB 内存。在二十次会话中，SwiftInfer 显示出显著的性能优势，并将继续努力提高性能和适配更新版本的 TensorRT API。

项目规划

• 完成基于 TRT-LLM 的注意力实现
• 适配 KV 缓存
• 适配提前终止
• 修复连续张量问题
• 支持多轮对话的 Llama 示例

鸣谢

在开发过程中，SwiftInfer 受到了 Streaming-LLM 项目的启发，并参考了多种开源材料和学术贡献，特此鸣谢。