纯 Python 的推理框架-lightLLM

大模型技术演进的速度，用“日新月异”来形容一点不为过。在LLM推理部署这个方向上，TGI、vLLM、lightLLM这几个框架吸引了最多的关注——它们都能更高效地榨干GPU的计算能力，在访存与计算之间找到更好的平衡点。特别值得一提的是lightLLM，一个纯Python实现的推理框架。它的整个服务流程和模型推理几乎全部用Python写就，代码结构清晰，二次开发的门槛相对低了不少。这篇文章会在lightLLM官方文档的基础上，对代码做进一步的拆解，希望能帮大家理解它的设计思路。

一、概览

lightLLM的README里列出了框架的主要特性，其中与模型推理强相关的包括：Nopad、FlashAttention、Tensor Parallel、Int8KV Cache、Token Attention、Dynamic Batch。后面两个特性与调度部分关系更紧密，而前四个会在模型推理的代码中直接体现。本文的重点就是沿着模型推理的代码结构，把这些特性串起来讲清楚。

如果之前没接触过lightLLM的模型结构，官方文档（docs/AddNewModel_CN.md）是个不错的起点——它以bloom为例，展示了如何添加自定义模型backend，里面把模型相关代码拆成了三个部分：layer_weights、layer_infer、triton_kernel。简单说：layer_weights负责加载各模型参数文件并统一命名空间；layer_infer接收输入和参数完成每层运算；triton_kernel则用OpenAI的Triton库，用Pythonic的语法实现底层CUDA kernel。

在basemodel目录下，整体框架已经搭好，但基类方法的具体实现留给了子类。Triton Kernel部分提供了基础的矩阵乘法（支持int8/int4量化）。目录结构大致如下：

├── basemodel.py
├── infer_struct.py
├── __init__.py
├── layer_infer
│   ├── base_layer_infer.py
│   ├── __init__.py
│   ├── post_layer_infer.py
│   ├── pre_layer_infer.py
│   ├── template
│   │   ├── __init__.py
│   │   ├── post_layer_infer_template.py
│   │   ├── pre_layer_infer_template.py
│   │   └── transformer_layer_infer_template.py
│   └── transformer_layer_infer.py
├── layer_weights
│   ├── base_layer_weight.py
│   ├── hf_load_utils.py
│   ├── __init__.py
│   ├── pre_and_post_layer_weight.py
│   └── transformer_layer_weight.py
└── triton_kernel
    ├── apply_penalty.py
    ├── destindex_copy_kv.py
    ├── quantize_gemm_int8.py
    ├── dequantize_gemm_int8.py
    ├── dequantize_gemm_int4.py
    └── __init__.py

二、模型推理代码

lightllm/models目录下是框架支持的各模型以及推理核心代码实现。这里拿bloom模型作为例子，一步步拆开看。

模型基类解读

bloom的模型主体类继承了TpPartBaseModel基类，子类的主要工作是给基类中的一些方法赋上实际的值，比如pre_and_post_weight_class = BloomPreAndPostLayerWeight。整体的调度逻辑仍然在基类里，子类只需要负责个性化配置。

模型参数初始化

初始化过程的第一步，是从HuggingFace格式的模型文件中读参数，把命名空间统一之后再一个个拷贝给模型实例。最核心的代码在BloomPreAndPostLayerWeight和BloomTransformerLayerWeight的load_hf_weights方法里——它们才是真正创建参数实体的地方。比如下面这行代码，就把分词嵌入的权重从weights字典中取出来，切分后赋值给self.wte_weight_：

self.wte_weight_ = self._cuda(weights["word_embeddings.weight"][split_vob_size * self.tp_rank_: split_vob_size * (self.tp_rank_ + 1), :])

不少LLM的参数量巨大，需要分摊到多张GPU上，参数切分就是在这里完成的。既然有Tensor Parallel的实现，初始化时就得考虑多卡协作。这部分在test/model/model_infer.py的test_model_inference函数里能看到——通过for循环加进程的方式，完成多个子模型的并行初始化。

# test/model/model_infer.py
def test_model_inference(world_size, model_dir, model_class, batch_size, input_len, output_len):
    ans_queue = Queue()
    workers = []
    for rank_id in range(world_size):
        proc = multiprocessing.Process(target=tppart_model_infer, args=(rank_id, world_size, ans_queue, model_dir, model_class, batch_size, input_len, output_len))
        proc.start()
        workers.append(proc)
    for proc in workers:
        proc.join()
    assert not ans_queue.empty()
    while not ans_queue.empty():
        assert ans_queue.get()
    return

从模型结构上说，lightLLM把模型按transformer layer和pre_and_post layer拆分。pre_and_post layer包含word embedding、输出层等；transformer layer则包含自注意力、FFN、LayerNorm等核心组件。

模型前向运算

前向的实现集中在各模型下的layer_infer目录。以bloom为例，三个核心脚本分别对应三种层：

├── layer_infer
│   ├── __init__.py
│   ├── post_layer_infer.py
│   ├── pre_layer_infer.py
│   └── transformer_layer_infer.py

每种层的前向有两个核心函数：context_forward和token_forward。在模型类的forward方法中，通过is_prefill标志位决定用哪个。从test/model_model_infer.py的调用方式看，context_forward会一次性处理整个上下文（比如用户的prompt+问题，可以看到多个token），而token_forward则逐token追加，每次循环调用时标记序列起始位置的index逐步加1。对应到实际场景：第一次模型调用得到用户输入，用context_forward；之后每次预测一个新token，用token_forward循环完成。

在lightllm/models/bloom/layer_infer/pre_layer_infer.py中，前向计算主要包括两部分：word embedding和LayerNorm。embedding直接用PyTorch内置函数实现：input_embdings = torch.embedding(layer_weight.wte_weight_, tmp_input_ids, padding_idx=-1)，多卡拆分则通过dist.all_reduce完成。LayerNorm的计算则用Triton Kernel实现，效率更高。

Transformer layer是更复杂也更核心的部分，按算法组件拆成attention和ffn两个子模块。这一层的实现在基类与子类都有重要代码：基类TransformerLayerInferTpl负责组合attention和ffn模块按序计算，并管理KV缓存。KV缓存管理尤其关键——LLM在自回归推理时，为避免历史信息重复计算，会将注意力机制中的k/v历史值缓存起来。基类通过_pre_cache_kv与_post_cache_kv在attention计算前后完成kv值的获取和新kv值的缓存。具体的高效管理方式，后面会单独讲。

postlayer的实现位于lightllm/models/bloom/layer_infer/post_layer_infer.py。有意思的是，它只有token_forward实现，没有context_forward——因为postlayer是做输出预测的层，用户提交的文本部分不需要预测，只有答案部分才需要。

TritonKernel

前面介绍的代码主要是在搭建框架结构，具体的内存管理、基础矩阵乘（比如torch.addmm）都有了，但像attention和LayerNorm这类复杂算子的具体实现并不在里面。这些算子可以通过计算顺序重规划、算子融合来提升访存和运算效率，lightLLM用OpenAI的Triton库来实现这一层。

bloom目录下实现了attention和LayerNorm相关的Triton Kernel：

├── triton_kernel
│   ├── __init__.py
│   ├── layernorm.py
│   ├── token_flashattention_nopad.py
│   ├── context_flashattention_nopad.py
│   └── ...

LayerNorm的实现与Triton官方示例一致。计算时要先沿C维度算mean和var，再减均值除方差。代码分三段，看懂一段其他两段也就自然明白了。将输入视作二维矩阵，比如(BT, C)维度的输入会被划分到形状的硬件资源上并行计算。每次计算一行，通过X += row * stride定位第row行的起始位置。每行有N个元素，每个元素由BLOCK_SIZE个线程并行处理，所以每行需要range(0, N, BLOCK_SIZE)次循环来完成：

row = tl.program_id(0)
Y += row * stride
X += row * stride
# Compute mean
mean = 0
_mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
for off in range(0, N, BLOCK_SIZE): # N > BLOCK_SIZE
    cols = off + tl.arange(0, BLOCK_SIZE)
    a = tl.load(X + cols, mask=cols < N, other=0.).to(tl.float32)
    _mean += a
mean = tl.sum(_mean, axis=0) / N

ContextFlashAttention的实现要复杂得多，核心代码有几十行。思路是：q@kT以及softmax(q@kT)的中间结果不写回HBM，减少内存访问次数。这部分代码解读篇幅不小，后续可以单独开一章。nopad特性也体现在这里——在load数据和计算时，通过mask特性避开无效的计算。

三、总结

随着LLM的发展，新的工程技术也在快速迭代。有意思的是，Nvidia的TensorRT在Transformer模型大放异彩的这段时间里，迭代速度反而没跟上；而AI编译器技术衍生出了新的可行路线。从lightLLM的代码可以很清楚地看到这一点——模型推理部分完全舍弃了onnx、TensorRT等库，转而用PyTorch做内存和算子的组合管理，用Triton实现底层kernel，用HuggingFace Transformers的模型仓库加载权重。

LLM对算力的巨大需求无疑让Nvidia凭借软硬件优势继续高速增长，但在软件层面，以Python动态思想为核心的方式，正在简洁度、代码易写性和灵活性上获得越来越大的优势。lightLLM就是这条路线上一个很有代表性的成果。