LightLLM 轻量级推理框架

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

源码：https://github.com/ModelTC/lightllm

1. 引言

1.1 前言

随着 ChatGPT 的火爆出圈，大语言模型受到越来越多的关注。然而，如何低成本、高吞吐地将参数量动辄千亿的模型部署到各类服务器上，成为技术大范围推广的关键。

LightLLM 应运而生，它引入了更细粒度的 KV Cache 管理算法 TokenAttention，并设计了与 TokenAttention 高效配合的 Efficient Router 调度实现。

在 TokenAttention 和 Efficient Router 的相互作用下，LightLLM 在大部分场景下都能获得比 vLLM 和 Text Generation Inference 更高的吞吐，部分场景下可得到 4 倍左右的性能提升。

1.2 为什么需要 LightLLM？

部署大模型提升服务性能存在以下挑战：

挑战	说明
显存碎片化严重	几十乃至上百 GB 的网络权重以及推理时不断动态产生的 KV Cache，极易造成显存利用率低
请求调度效率低	请求的长度随时间动态变化，易造成 GPU 空转或利用率低
kernel 定制化难度高	需要为网络定制化 CUDA C kernel，对普通研究员难度较高

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

框架	优势	不足
FasterTransformer	优秀的静态推理性能	无良好的服务调度功能，C++ 开发，二次开发成本高
TGI	优秀的服务接口和 Continuous Batch	推理性能、调度策略、显存管理有缺陷
vLLM	显存管理优秀	请求调度效率不高，更适合小模型部署

2. LightLLM 介绍

2.1 什么是 LightLLM？

LightLLM 是一套基于纯 Python 语言的大模型推理部署框架，方便研究员进行轻量级的本地部署和定制修改。

核心特性

特性	说明
三进程架构	异步化处理 tokenize 和 detokenize 操作，避免 CPU 处理阻碍 GPU 调度
Token Attention	以 Token 为粒度进行 KV Cache 显存管理，实现高性能管理方法
Efficient Router	配合 Token Attention 精确管理调度请求的合并推理
Triton 算子	基于 OpenAI Triton 开发的高效算子，实现优秀吞吐性能

支持模型

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

2.2 Token Attention 介绍

目前的大语言模型都是基于 Transformer 架构的，对问题的回复通过自回归解码逐 token 产生。模型会将历史上下文 token 在注意力模块产生的 key 和 value 缓存到 GPU 中，以便快速生成下一个 token。

这些缓存占据大量显存空间，且大小高度变化和不可预测，如果没有合理的显存管理手段，将造成显存碎片化严重。

TokenAttention 工作流程

1. 模型初始化：系统根据用户设置的 max_total_token_num 预先分配 KV Cache，并创建 Token Table 记录输入 token 的实际存储位置
2. 请求到来：系统检查预分配的 Token Cache 中是否有可用连续空间存储请求的 KV 缓存
3. 空间分配：系统倾向于分配连续显存以减少访存时间，仅当连续空间不足时才分配非连续显存
4. 自回归生成：新生成 token 的缓存从预分配缓存中找到未使用空间，记录到 Token Table

Cache 状态管理

# 状态定义
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

# 分配过程：利用 GPU 并行特性高效筛选未使用空间
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)

# 释放过程：仅需删除 Token Table 中的记录
self.can_use_mem_size += free_index.shape[0]
self.mem_state[free_index] = 1

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

2.3 Efficient Router 介绍

Router 的主要功能是管理到达的请求，并动态判断该请求能否和已在运行的 Batch 融合到一起进行推理。

判断核心

估计合并后的整个推理过程中，Token 的最大占用量是否小于可以容纳的容量 max_total_token_num。

Token 最大使用量计算

假设将新请求融合到 Batch 中进行推理，token 的最大使用量必然在以下三个时刻之一到来：

时刻	Token 总占用量
Time 1	黄色格子数量 + 绿色格子数量（所有请求的已运行 token + 新请求的全部待生成 token）
Time 2	黄色格子数量 + 绿色格子数量
Time 3	黄色格子数量（仅已运行 token）

伪代码实现

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 使用量 <= max_total_token_num，说明新的请求可以进行合并 Batch 推理。

3. LightLLM 性能表现

在数据集 ShareGPT_Vicuna_unfiltered 上与主流推理框架进行性能对比：

性能对比结论

框架	表现
TGI	显存碎片化严重，很难达到较高的吞吐量
vLLM	引入 PageAttention，但大模型并发性能不理想（默认配置）
LightLLM	各种大小的模型下都保持稳健性能，大模型上相对 TGI 和 vLLM 实现 3 倍左右提升

TGI 兼容与消融分析

分析维度	结果
TGI 兼容	引入 TokenAttention 和 Router 后，给原始 TGI 带来 4 倍以上性能提升
长短不齐请求	Router 在问题长度差异很大时，带来近 50% 的性能提升

ShareGPT_Vicuna_unfiltered 数据集问题长短差异不显著，Router 优势不明显；在问题差异更大的请求集合中，Router 性能优势显著。

4. LightLLM 依赖包

依赖	要求
PyTorch	>= 1.3
CUDA	11.8
Python	3.9

5. LightLLM 安装

5.1 下载 LightLLM

$ git clone https://github.com/ModelTC/lightllm
$ cd lightllm
$ pip install -r requirements.txt

5.2 安装 LightLLM

$ 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 内核

6. LightLLM 使用

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

6.1 启动 LightLLM 服务

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

参数说明：

参数	说明
--model_dir	模型权重路径
--tp	张量并行数
--max_total_token_num	受 GPU 显存影响，值越大允许处理更多并发请求，提高系统并发性

6.2 调用 LightLLM 服务

使用 curl 调用：

$ 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 调用：

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)

7. LightLLM 支持模型

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

LightLLM 作为纯 Python 实现的轻量级推理框架，通过 TokenAttention 和 Efficient Router 的协同作用，在大模型推理场景下实现了优异的吞吐性能，特别适合需要快速定制和扩展的研究场景。