纯Python超轻量高性能LLM推理框架 —— LightLLM

来源：AiGC面试宝典作者：宁静致远日期：2023年09月29日官方文档：https://modeltc.github.io/lightllm/ 源码：https://github.com/ModelTC/lightllm

一、引言

1.1 前言

随着ChatGPT的火爆出圈，大语言模型受到越来越多的关注。这类模型的出现极大地提高了人们的工作效率，然而，如何低成本、高吞吐地将参数量动辄千亿的模型部署到各类服务器上，成为将技术进一步大范围推广的关键。为了提高大模型服务的吞吐量，同时让更多感兴趣的人快速上手参与其中，一个轻量化的LLM推理服务框架LightLLM应运而生。

📝 通俗解释：LightLLM就像是一个"智能餐厅服务员"，它能同时接待很多客人（处理多个请求），让厨房（GPU）始终保持高效运转，不会出现空闲或拥挤的情况。

LightLLM引入了一种更细粒度的KV Cache管理算法——TokenAttention，并设计了一个与TokenAttention高效配合的Efficient Router调度实现。在TokenAttention和Efficient Router的相互作用下，LightLLM在大部分场景下都能获得比vLLM和Text Generation Inference更高的吞吐，部分场景下可以得到约4倍左右的性能提升。LightLLM灵活、易用、高效，感兴趣的同学不妨点击上方项目链接上手一试。

1.2 为什么需要 LightLLM？

大语言模型由于其出色的对话性能受到了研究人员的广泛关注，其中比较有代表性的模型结构有BLOOM、LLaMA等。这些模型不仅可以和人类进行日常对话，同时可以帮助人们完成一些生产工作，提高工作效率。然而，虽然这些模型已经表现出来了卓越的性能，但是部署大模型提升服务性能，存在以下挑战：

显存碎片化严重：几十乃至上百GB的网络权重以及推理时不断动态产生的KV Cache，极易造成显存利用率低的问题。
请求调度效率低：请求的长度随时间动态变化，易造成GPU空转或是利用率低的问题。
Kernel定制化难度高：为了高效利用显存，提高服务吞吐量，需要为网络定制化CUDA C Kernel，对于普通研究员来说，难度较高。

📝 通俗解释：
显存碎片化：就像你的电脑内存被各种程序占用得七零八落，明明总内存够用，但就是找不到一大块连续的空间来运行新程序
请求调度效率低：就像餐厅服务员不知道该什么时候叫下一位客人入座，导致厨房有时太闲有时太忙
Kernel定制化：就像要自己发明一套厨具来做菜，普通厨师很难做到

1.3 目前 LLM推理框架有哪些？

为了应对以上挑战，目前涌现了很多优秀的LLM推理框架，例如FasterTransformer、Text-Generation-Inference（简称TGI）、vLLM等。他们的核心Feature和能力矩阵如下表所示：

框架	NV Triton + FasterTransformer	TGI	vLLM	LightLLM
核心Feature	算子高效，静态速度快	Continuous Batch，流式推理	PageAttention	三进程架构，TokenAttention，Efficient Router
显存碎片化	低	高	低	低
请求调度效率	低	中	中	高
Kernel定制化难度	高	中	中	低

这些推理框架都具有自己独特的特色：

FasterTransformer：具有优秀的静态推理性能，但是没有良好的服务调度处理功能，并且主体用C++开发而成，二次开发成本较高。
TGI：具有优秀的服务接口和服务调度特性如Continuous Batch，但是推理性能、调度策略、显存管理亦有缺憾。
vLLM：显存管理优秀，但是请求调度效率不高，且整体实现细节更适合于小模型的部署。

📝 通俗解释：
FasterTransformer：像一辆方程式赛车，直线跑得很快，但不会自己找路
TGI：像一辆配置齐全的商务车，有导航（调度）功能，但发动机效率一般
vLLM：像一辆省油的车（显存管理好），但载客量有限（不适合大模型）
LightLLM：综合了各种优点，既省油又跑得快，还能智能调度

二、LightLLM 介绍

2.1 什么是 LightLLM？

为了解决这些问题，我们开发了一套基于纯Python语言的大模型推理部署框架LightLLM，方便研究员进行轻量级的本地部署和定制修改，用于快速扩展对不同模型的支持，吸纳层出不穷的优秀开源特性，探索最优服务架构。

📝 通俗解释：纯Python意味着这个框架更容易上手和修改，就像用简单的烹饪方法做菜，比从零打造一套专业厨具要容易得多。

LightLLM的核心Feature如下：

三进程架构。主要用于异步化处理Tokenize和Detokenize操作，可以避免这些耗时的CPU处理阻碍模型推理时GPU的调度执行，降低GPU的利用率，进而降低整体服务性能。

📝 通俗解释：三进程架构就像餐厅里分工明确：一个人专门负责迎宾（接收请求），一个人专门负责端菜（处理结果），一个人专门负责炒菜（模型推理）。各司其职，互不耽误。

TokenAttention。一种以Token为粒度进行KV Cache显存管理的特性，并实现了高性能的管理方法。

📝 通俗解释：TokenAttention就像一个智能仓库管理员，它不是按"房间"来分配空间，而是精确到每个"箱子"来管理显存，用完就立即回收，不浪费任何空间。

Efficient Router。配合TokenAttention用于精确管理调度请求的合并推理。

📝 通俗解释：Efficient Router就像餐厅的排班经理，它会精心安排哪些客人可以同时入座（合并Batch推理），让有限的餐桌（GPU资源）发挥最大效用。

配合基于OpenAI Triton开发的与服务调度高度配合的高效算子，LightLLM实现了优秀的吞吐性能。

LIGHT LLMA Light and Fast Inference Service for LLM

📝 通俗解释：LightLLM的口号就是"又轻又快"，目标是让大模型推理像开灯一样简单快捷。

![架构图描述：左侧request进入httpserver，httpserver输出output text（经detokenization），同时发送tokenized requests给Router。Router通过RPC与下方的Model Backend with Token Attention交互，Model Backend返回next token id给detokenization。]

目前支持模型：

BLOOM: https://huggingface.co/bigscience/bloom
LLaMA: https://github.com/facebookresearch/llama
LLaMA V2: https://huggingface.co/meta-llama

📝 通俗解释：LightLLM支持主流的开源大模型，就像一个 универсальный 转换插头，可以适配多种不同规格的电器。

2.2 TokenAttention 介绍

目前的大语言模型都是基于Transformer架构的，且对问题的回复是通过自回归解码逐Token产生的。为了快速地生成下一个Token，这些模型会把自回归解码过程中历史上下文Token在注意力模块产生的Key和Value缓存到GPU中，以便于能快速地生成下一个Token。这些缓存会占据大量的显存空间，且由于每个问题的长短不一，其大小也是高度变化和不可预测的，如果没有合理的显存管理手段，这将会造成严重的显存碎片化，产生极大的显存浪费。

📝 通俗解释：
KV Cache：就像在做阅读理解时，你会在大脑中记住之前段落的重要内容（Key和Value），这样就不用每次都重新读一遍
显存碎片化：就像房间里堆满了各种大小的盒子，明明总体积够用，但想放一个大盒子却找不到足够大的连续空间

因此我们设计了一种以Token为粒度进行KV Cache显存管理的注意力计算方法——TokenAttention，并实现了高性能的算子和高效的显存申请释放方式。TokenAttention的计算过程如下图所示：

1. Init the Memory Cache

Seq1: The capital of China is Seq2: Shanghai is

Token Table

Input	Token id
-	-
-	-
-	-
-	-
...	...
-	-
-	-

Pre-allocated Token Cache

token 1
token 2
token 3
token 4
token 5
...
token n-2
token n-1
token n

具体地，在模型初始化时，系统根据用户设置的 max_total_token_num 预先分配KV Cache，并创建Token Table来记录输入Token的实际存储位置。其中，max_total_token_num为部署环境的硬件显存一次最多能容纳的Token总量。

📝 通俗解释：就像餐厅开业前先准备好固定数量的餐桌（预分配显存），并建立一个座位表（Token Table）来记录每位客人坐在哪里。

当请求到来时，系统首先检查预分配的Token Cache中是否有可用的连续空间用于存储请求的KV缓存。系统倾向于为请求分配连续的显存，以最大限度地减少推理过程中的访存时间，仅当连续空间不足时，才会为请求分配非连续显存。分配的空间会记录到Token Table中，以便于后续的注意力计算。

📝 通俗解释：有连续空间就像让一家人坐在一起（访问快），没位置时只能分开坐（访问慢但能应急）。

对于自回归过程新生成的Token的缓存，仅需从预先分配的Token缓存中找到未使用的空间，并将相应的记录添加到Token Table中即可。

📝 通俗解释：新生成的Token就像新来的客人，只需要找一个空位坐下，并在座位表上登记即可。

为了高效地进行Cache的分配和释放，我们利用PyTorch Tensor在GPU上的并行计算特性来对预分配的Token Cache的状态进行管理，首先是状态定义：

python

self.mem_state = torch.ones((size,), dtype=torch.bool, device="cuda")
self._mem_cum_sum = torch.empty((size,), dtype=torch.int32, device="cuda")
self.indexes = torch.arange(0, size, dtype=torch.long, device="cuda")
self.can_use_mem_size = size

其中：

mem_state 是使用状态数组，1为未使用，0为已使用
_mem_cum_sum 数组用于mem_state的累加和，以此用于高效地筛选出未使用的空间，用于Cache的分配

分配过程如下：

python

torch.cumsum(self.mem_state, dim=0, dtype=torch.int32, out=self._mem_cum_sum)
select_index = torch.logical_and(self._mem_cum_sum <= need_size, self.mem_state == 1)
select_index = self.indexes[select_index]
self.mem_state[select_index] = 0
self.can_use_mem_size -= len(select_index)

📝 通俗解释：这段代码利用GPU的并行计算能力，快速找出所有可用的空位（未使用的显存）。就像用扫描仪快速检查哪些餐桌是空的。

可以看到，我们的Cache状态管理只在GPU上完成，充分利用了PyTorch在GPU上的并行特性，从而使得系统能够非常高效地进行Cache空间的申请和释放。

对于已经完成的请求，仅需删除Token Table中的记录，就可以完成空间的高效释放。

释放代码如下：

python

self.can_use_mem_size += free_index.shape[0]
self.mem_state[free_index] = 1

由于TokenAttention是以Token为粒度进行显存的管理，所以其可以做到显存空间的零浪费，且能够准确地计算出系统还能容纳多少新Token的计算。因此，配合一个高性能的Router对请求进行管理，可以在推理过程中不断地加入新请求，充分压榨GPU的每一块显存，最大化GPU利用率。

📝 通俗解释：TokenAttention的精准管理让显存利用率接近100%，就像一个优秀的空间利用大师，能把仓库的每一寸空间都用到极致。

2.3 Efficient Router 介绍

Router的主要功能是管理到达的请求，并动态地判断该请求能否和已经在运行的Batch融合到一起进行推理。这个判断的核心就是估计合并后的整个推理过程中，Token的最大占用量是否小于可以容纳的容量。这里我们将这个最大容量设为 $max_total_token_num$，由于有TokenAttention特性的支持，我们可以精确地管理Token的使用量，合理的配置可以使其永远没有发生OOM的风险。

📝 通俗解释：Router就像餐厅的排队长官，它会估算"如果让这批新客人进来，会不会超过餐厅的最大接待能力"。有了精确的座位统计（TokenAttention），这个估算可以非常准确，永远不会超载。

![图表描述：图片展示了一个时间轴（Time 0, Time 1, Time 2, Time 3）上的请求处理甘特图。左侧标注"new req"指向中间的一行绿色格子。图中有黄色格子（代表已有请求的历史或当前处理）、绿色格子（代表新请求）、灰色/浅蓝色格子（代表待生成的token）。]

如上图所示，每行代表一个请求当前的运行状态，深色代表已经运行完的历史KV Cache Token，每个格子代表一个Token，灰色代表待生成的Token，待生成的Token数由每个请求设置的最大输出长度和已经生成的Token数决定。上图中的第二行绿色格子所在的行代表一个新到达的请求，图中将所有请求按照待输出的长度进行从大到小的顺序排列。

📝 通俗解释：这张图就像餐厅的时间排班表，横轴是时间，纵轴是不同客人的需求。黄色是已经在吃的客人，绿色是新来的客人，灰色是还没端上来的菜。

如果我们假设将新请求融合到Batch中进行推理，那Token的最大使用量必然在时刻Time 1、Time 2、Time 3中某个时刻到来。我们只需要计算这三个时刻对应的Token使用量都不超过 $max_total_token_num$，则代表新的请求可以加入到Batch中进行融合推理。

Time 1 Token的总占用量：黄色格子数量 + 绿色格子数量
Time 2 Token的总占用量：黄色格子数量 + 绿色格子数量
Time 3 Token的总占用量：黄色格子数量

📝 通俗解释：排队长官会检查三个关键时刻：现在开始融合、新请求结束时、最后结束时。只要这三个时刻都不超过最大容量，就可以让新请求加入。

实际的Token最大使用量，必然是Time 1、Time 2、Time 3其中之一。只要保证动态推理过程中的最大Token使用量 <= max_total_token_num，说明新的请求可以进行合并Batch推理。

为了快速地计算一个Batch的所有请求需要的最大Token使用量，我们利用NumPy实现了一个高效的示例实现，下面是Python伪代码：

python

import numpy as np

def demo():
    max_total_token_num = 100
    req_list = [(5, 4), (4, 3), (5, 3), (3, 2), (4, 2)]  # (run_len, left_output_len)
    req_list.sort(key=lambda x: -x[1])

    left_out_len_array = np.array([e[1] for e in req_list])
    has_run_len_array = np.array([e[0] for e in req_list])
    cum_run_len_array = np.cumsum(has_run_len_array)
    size_array = np.arange(1, len(req_list) + 1, 1)
    need_max_token_num = (left_out_len_array * size_array + cum_run_len_array).max()

    if need_max_token_num <= max_total_token_num:
        print("ok")
    else:
        print("oom")

📝 通俗解释：这段代码计算"如果让这批请求一起跑，需要多少座位"。它按待输出长度排序，然后计算每个时刻的最大Token使用量，找出峰值来判断是否会爆满。

三、LightLLM 性能表现介绍

我们在数据集ShareGPT_Vicuna_unfiltered上和目前主流的推理框架TGI、NV Triton + FasterTransformer以及vLLM进行了性能对比，结果如下图所示。

可以看到，LightLLM在不同大小的模型下都获得了更高的吞吐量。

TGI：由于显存碎片化严重，所以很难达到较高的吞吐量
vLLM：因引入了PageAttention，但是由于整体实现细节更有利于小模型推理，所以在大模型上的并发性能并不是十分理想（使用的默认配置）
相比之下，LightLLM则可以在各种大小的模型下都保持稳健的性能，在大模型上（LLaMA-65B）相对TGI和vLLM实现了约3倍的提升

📝 通俗解释：LightLLM就像一个全能型选手，在不同体量的比赛（不同大小的模型）中都能保持好成绩，而其他选手要么只擅长小比赛，要么容易发挥失常。

图表内容描述：

图表标题为 Throughput (requests/s)，展示了不同模型在不同推理框架下的吞吐量对比。

X轴分别为三个模型配置：

BLOOM-7B (1 x A800-80G)
LLaMA-7B (1 x A800-80G)
LLaMA-65B (8 x A800-80G)

图例包括：NV Triton + FasterTransformer（蓝色），TGI（橙色），vLLM-0.2.1（灰色），LightLLM（黄色）。

具体数据如下：

BLOOM-7B:
- NV Triton + FasterTransformer: 2.43
- TGI: 1.33
- vLLM-0.2.1: 7.05
- LightLLM: 9.95
LLaMA-7B:
- TGI: 2.7
- vLLM-0.2.1: 5.38
- LightLLM: 10.33
LLaMA-65B:
- TGI: 2.86
- vLLM-0.2.1: 2.48
- LightLLM: 7.31

3.1 TGI兼容&消融分析

为了进一步验证TokenAttention和Router的有效性，我们同样将这些特性接入到了TGI中，来改善其显存碎片化的问题，结果如下图（左）所示。可以看到，在引入TokenAttention以及Router以后，可以给原始TGI带来4倍以上的性能提升。

📝 通俗解释：消融实验就像"对比实验"，证明LightLLM的每个关键技术（TokenAttention和Router）都是有效的"功臣"。

图表内容描述（左图 Throughput requests/s）:

LLaMA-7B (1 x A800-80G):
- TGI: 2.7
- TGI + TokenAttention: 10.33
- TGI + TokenAttention + Efficient Router: 10.85
LLaMA-65B (8 x A800-80G):
- TGI: 2.86
- TGI + TokenAttention: 8.22
- TGI + TokenAttention + Efficient Router: 8.53

3.2 长短不齐请求情况下的提升

从下图（左）中可以发现，Router的引入并未带来较为明显的性能提升，这是由于ShareGPT_Vicuna_unfiltered的数据集中问题长短差异并不显著。为此我们构建了一个问题差异更大的请求集合，对我们的Efficient Router的性能进行了验证，结果如下图（右）所示。

可以看到，我们的高性能Router可以更好地利用GPU资源，在问题长度差异很大的请求下，可以带来近50%的性能提升。

📝 通俗解释：当请求长短差异大时（有的要生成很短的回答，有的要生成很长的回答），Router的智能调度能力更能体现价值。就像排班时让"吃快餐"和"吃大餐"的客人错开安排，能显著提高翻台率。

图表内容描述（右图 Throughput requests/s）:

LLaMA-65B (8 x A800-80G):
- TGI + TokenAttention: 4.02
- TGI + TokenAttention + Efficient Router: 6.09 (提升约51.5%)

四、LightLLM 依赖包有哪些？

PyTorch >= 1.13（更新版本推荐1.13+）
CUDA 11.8
Python 3.9+

📝 通俗解释：这些是运行LightLLM的基本环境要求，就像开车需要合适的汽油和驾驶证一样。

五、LightLLM 如何安装？

5.1 下载 LightLLM

bash

$ git clone https://github.com/ModelTC/lightllm
$ cd lightllm

5.2 安装 LightLLM 依赖

bash

$ pip install -r requirements.txt

5.3 安装 LightLLM

bash

$ python setup.py install

⚠️ 注意：
该代码已在一系列GPU上进行了测试，包括A100、A800、4090和H800
如果您在A100、A800等上运行代码，建议使用triton 2.0.0.dev20221202
如果您在4090、H800等平台上运行代码，则需要从GitHub仓库编译并安装triton 2.1.0的源代码
如果代码在其他GPU上不起作用，请尝试修改模型推理中使用的Triton内核

📝 通俗解释：不同GPU需要不同版本的Triton（一种加速库），就像不同型号的手机需要不同版本的操作系统。

六、LightLLM 如何使用？

通过高效的Router和TokenAttention，LightLLM可以作为一项服务进行部署，并实现最先进的吞吐量性能。

6.1 启动 LightLLM 服务

bash

$ python -m lightllm.server.api_server --model_dir /path/llama-7B --tp 1 --max_total_token_num 120000

参数说明：

model_dir：模型权重文件所在目录
tp（tensor_parallel）：Tensor并行度，用于多卡部署
max_total_token_num：受部署环境的GPU显存影响。此参数的值越大，允许处理更多并发请求，从而提高系统并发性

📝 通俗解释：max_total_token_num就像餐厅的最大客容量，设得太高会爆满（OOM），设得太低会浪费资源。

6.2 调用 LightLLM 服务

使用curl调用：

bash

$ curl 127.0.0.1:8000/generate \
-X POST \
-d '{"inputs":"What is AI?","parameters":{"max_new_tokens":17, "frequency_penalty":1}}' \
-H 'Content-Type: application/json'

使用Python调用：

python

import time
import requests
import json

url = 'http://localhost:8000/generate'
headers = {'Content-Type': 'application/json'}
data = {
    'inputs': 'What is AI?',
    "parameters": {
        'do_sample': False,
        'ignore_eos': False,
        'max_new_tokens': 1024,
    }
}
response = requests.post(url, headers=headers, data=json.dumps(data))
if response.status_code == 200:
    print(response.json())
else:
    print('Error:', response.status_code, response.text)

附录：LightLLM 支持模型列表

目前 LightLLM 支持模型：

BLOOM: https://huggingface.co/bigscience/bloom
LLaMA: https://github.com/facebookresearch/llama
LLaMA V2: https://huggingface.co/meta-llama

📝 通俗解释：随着项目发展，LightLLM支持模型还在持续增加中，具体请参考官方GitHub仓库的最新列表。

整理日期：2024年8月来源：知识星球 - AiGC面试宝典

纯Python超轻量高性能LLM推理框架 —— LightLLM ​

一、引言 ​

1.1 前言 ​

1.2 为什么需要 LightLLM？ ​

1.3 目前 LLM推理框架 有哪些？ ​

二、LightLLM 介绍 ​

2.1 什么是 LightLLM？ ​

2.2 TokenAttention 介绍 ​

2.3 Efficient Router 介绍 ​

三、LightLLM 性能表现 介绍 ​

3.1 TGI兼容&消融分析 ​

3.2 长短不齐请求情况下的提升 ​

四、LightLLM 依赖包 有哪些？ ​

五、LightLLM 如何安装？ ​

5.1 下载 LightLLM ​

5.2 安装 LightLLM 依赖 ​

5.3 安装 LightLLM ​

六、LightLLM 如何使用？ ​

6.1 启动 LightLLM 服务 ​

6.2 调用 LightLLM 服务 ​

附录：LightLLM 支持模型列表 ​

纯Python超轻量高性能LLM推理框架 —— LightLLM

一、引言

1.1 前言

1.2 为什么需要 LightLLM？

1.3 目前 LLM推理框架有哪些？

二、LightLLM 介绍

2.1 什么是 LightLLM？

2.2 TokenAttention 介绍

2.3 Efficient Router 介绍

三、LightLLM 性能表现介绍

3.1 TGI兼容&消融分析

3.2 长短不齐请求情况下的提升

四、LightLLM 依赖包有哪些？

五、LightLLM 如何安装？

5.1 下载 LightLLM

5.2 安装 LightLLM 依赖

5.3 安装 LightLLM

六、LightLLM 如何使用？

6.1 启动 LightLLM 服务

6.2 调用 LightLLM 服务

附录：LightLLM 支持模型列表