NumPy+GPT

Numpy

作为一个知识浅薄的人，此处不得不先补一下numpy的相关知识

  1. NumPy 是一个免费的开源 Python 库，用于 n 维数组(也称为张量)处理和数值计算
  2. NumPy有对大型多维数组（也称为矩阵或张量）的支持，提供了一系列高级数学函数(基本的线性代数、随机模拟、傅立叶变换、三角运算和统计运算)，是 Pandas、Scikit-learn 和 SciPy 等库的基础
  3. NumPy库的中心数据结构是多维数组(值的网格)。
    1. Numpy的ndarray(同构的n维数组对象)：每个项包含大小相同的内存块，且采用统一识别方式
    2. NumPy数组是类似数据类型的编译，在内存中密集集中打包
    3. 底层是用C语言实现的目标代码
    4. 使用前需要向量化(变成数组)
  4. 简单使用
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import numpy as np

# 创建数组
np.array([1，2，3，4，5]) # 创建数组
np.zeros((3,2)) # 全零  3行2列
np.ones((2,3)) # 全1
np.arange((3,7)) # 递增或递减

a = np.zeros((3,2))
a.shape # 获取尺寸

np.linspace((0,1,5)) # 返回介于某区间等间距分布的数
np.random.rand(2,4) # 生成随机数组

# 默认数据类型：64位浮点数
a = np.zero((4,2),dtype = np.int32) # 指定数据类型
b = a.astype(int) # 转换数据类型

# 两个相同尺寸的数组可以直接进行四则运算
np.dot(a,b) # 向量点乘
a @ b # 矩阵乘法
np.sqrt(a) # 依次求平方根
# 不同尺寸数组运算：先扩展至相同尺寸，再相同位置进行计算
a.min() # 返回最小值
a.argmin() # 返回最小元素所在索引

a[4:1:2] # 2为跨度，可取负值

# 还有其他计算此处不一一赘述
# 0:2代表可以取[0,2)这个区间但是不包括2
  1. 总结：
    1. 如此，NumPy可以进行图片处理
    2. 提升了数据运算速度，核心是将问题向量化、并行化

LLM(Large Language Model)

这是一个很大的概念，此处仅补充小小的部分

  1. GPT(Generative Pre-trained Transformer)模型: 生成式 预训练的 Transformer
  2. 在Transformer架构被提出之前，语言模型的主流架构是RNN(Recurrent Neural Netwoork 循环神经网络)
    1. 按顺序逐字处理，无法并行计算
    2. 不擅长处理长序列(长文本)
  3. Transformer:有能力学习输入序列里所有词的相关性和上下文
    1. 自注意力机制(注意力权重，相关性)
    2. 位置编码(向量化)
    3. 原始文本token化，token会用一个整数(token ID)表示
    4. 传入嵌入层之后，每个token用向量表示
    5. 编码器：对向量进行加工(多头自注意力机制，前馈神经网络)，输出结果
    6. 解码器：(带掩码的多头自注意力、多头自注意力机制)、输出表示->词汇表的概率分布(线性层、softmax层)

      其实是在猜下一个最有可能的输出，所以可能出现一本正经胡说八道的情况

  4. Q,K,V向量(每个 输入token都生成三组向量)：
    1. Q(查询向量):”我要寻找什么”
    2. K(键向量):”我有什么特征”
    3. V(值向量):”我实际携带的信息”

NumPy+GPT

既然有了前置知识，那么可以开始看正文:60行NumPy手搓GPT

原理

1. 输入

  1. 很明显，输入的是文本，然后token化(分词器产生)，接着被表示为一串整数序列I(词汇表进行映射)
  2. 在实际中，1过程还有其他方法：
    1. encode(编码)：str->list[int]
    2. decode(解码)：list[int]->str

2. 输出

  1. output[i][j] 模型的预测概率：位于vocab[j]的token是下一个tokeninputs[i+1]的概率
  2. greedy sampling(贪心采样)：将具有最高概率的token作为预测

    直接选择每个输出的最大概率，直到出现终结符或最大句子长度 (贪心算法)

  3. 生成文本
    1. 自回归：迭代，实现一个解码for循环，预测未来的值，并添加回输入中
    2. 贪心采样：引入一些随机性，可基于输入产生不同的输出句子

3. 提示

  1. 条件生成：基于提示内容生成文本

实践

1
2
git clone https://github.com/jaymody/picoGPT
pip install -r requirements.txt # 安装依赖

整体代码

1. gpt2()

  1. 嵌入层：
    1. token嵌入：词向量，编码token语义信息，成为可学习的参数矩阵，语义相似的词向量距离近
    2. 位置嵌入：token在序列中的位置信息，负责处理词序
    3. 组合嵌入：token信息+位置信息
  2. 解码层

    姑且理解为嵌入层(编码)的逆过程
    输入 → 嵌入层 → [Transformer Block] × n_layer → 输出

    1. Transformer Block: 层化一归+多头注意力机制+层化一归+前馈网络(非线性)
    2. 解码器模块：
      1. 多头因果自注意力:
        1. 注意力、自身、因果、多头
        2. 优势：降低单个头维度，增加头数量，不同的头学习不同关注模式，实现并行计算
        3. 因果掩码保证自回归性
      2. 逐位置前馈神经网络(FFN):增强模型对复杂模式的建模能力,每个位置独立处理，无跨位置交互,与注意力机制互补
  3. 投影为词汇表

    Transformer输出->层归一化->与嵌入矩阵转置相乘->logits输出

    1. 层归一化: 稳定最终输出的数值范围,防止数值溢出，提高训练稳定性

2. gernerate()

自回归生成

  1. 前向传播
  2. 贪心采样(取最后一个位置)
  3. 添加到上下文

3. main()

  1. 参数：prompt(输入文本提示) n_tokens_to_generate:(要生成的token数量，默认40个) model_size:(GPT-2模型大小) models_dir:(存储模型文件的目录路径)
  2. 过程：
    1. 加载分词器(encoder)， 模型权重（params）， 超参（hparams）

      encoder转化为token IDs hparams决定模型架构尺寸 params提供训练好的权重

    2. 使用分词器将输入提示词编码为token ID
    3. 调用生成函数
    4. 将输出ID解码为字符串
  3. 理解的简化过程如下：

    加载模型组件->编码文本->长度检查->生成->解码

4. fire.Fire(main)

源文件转成一个命令行借口(Google Fire库)

其他

神经网络基础层

  1. GELU激活函数：实现非线性变换，可以拟合复杂函数
  2. Softmax():实数转为概率分布
  3. 层归一化：稳定训练过程，加速收敛(归一化)
  4. 线性变换层：最基本的全连接层/投影操作

GPT架构

  1. 文本 + 位置嵌入
  2. 解码器层
  3. 投影为词汇表的步骤

PyTorch+KVcache

PyTorch

依然犹豫见识浅薄补一下前置知识

概述

  1. 是用于张量编程的库
  2. 张量是表示深度神经网络中数据和参数的多维数组
  3. 支持在GPU上进行高性能并行计算(+cuda,加速GPU上的计算)，支持动态计算图，允许在运行时优化模型
  4. 构建一个由函数组成的有向无环图DAG来实现对张量上所有已执行操作的追踪，允许根据需要在每次迭代后更改形状、大小和操作

基本用法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import torch

# 创建一维张量(向量)
x = torch.tensor([1, 2, 3]) 
# 创建二维张量(矩阵)
matrix = torch.tensor([[1, 2, 3],
                       [4, 5, 6]])
zeros = torch.zeros(2,3) # 2行3列的全零张量
# 创建随机张量(标准正态分布)
random = torch.randn(2,3) # 2行3列

# 看形状
print(x.shape)
print(x.size())
print(x.shape[0]) # 1
print(x.shape[1]) # 2
print(x.shape[-1]) # 3,最后一维

# 重塑张量(改变形状，不改变元素总数和顺序)
x = torch.arange(12)
y = x.view(3.4) # 变成3行4列
x.view(3,-1) # -1自动推断

# 交换维度
y = x.transpose(1,2) # 交换维度1和2
y = x.transpose(-2,-1) # 交换最后两个维度
'''
[B, T, n_head, head_dim].transpose(1, 2) → [B, n_head, T, head_dim]
k.transpose(-2, -1) 用于计算 Q @ K^T
# 注意力机制中常用的张量转置操作:先考虑所有注意力头，再所有时间步，方便批量矩阵乘法
'''

# 矩阵乘法
A = torch.randn(2,3)
B = torch.randn(3,4)
C = A @ B
'''
[..., m, k] @ [..., k, n] → [..., m, n]
att = q @ k.transpose(-2, -1) # 注意力机制中计算注意力分数->注意力分数矩阵
K_T # K的转置
'''

# 拼接张量
a = torch.randn(2, 3)
b = torch.randn(2, 3) 
c = torch.cat([a, b], dim=0) # dim=0,上下拼接   dim=1,左右拼接
'''
KV Cache 的核心(缓存机制，高效处理序列生成任务)
# past_k: [B, n_head, past_len, head_dim]
# k:      [B, n_head, 1, head_dim]  (新的一个 token)
new_k = torch.cat([past_k, k], dim=2)
# new_k: [B, n_head, past_len+1, head_dim]
'''

# 分割张量
x = torch.randn(2, 12)  # [2, 12]
parts = x.split(4, dim=1) # 沿 dim=1 分成 3 份，每份大小为 4
'''
qkv = self.c_attn(x)  # [B, T, 3*n_embd]
q, k, v = qkv.split(n_embd, dim=2)  # 各 [B, T, n_embd]
'''

# 索引和切片
x = torch.randn(4, 5, 6)
# 索引
x[0]    # 第一个元素，形状 [5, 6]
x[0, 1]     # 第一个面的第二行，形状 [6]
x[0, 1, 2]  # 一个标量
# 切片
x[0:2]      # 前两个面，形状 [2, 5, 6]
x[:, 0:3]   # 所有面的前三行，形状 [4, 3, 6]
x[:, :, -1] # 所有面所有行的最后一列，形状 [4, 5]
# 常用: 取最后一个位置
last_logits = logits[0, -1, :] # 也可用于取最后一个token预测

神经网络组件

  1. nn.Linear(线性层)
    1. 自动处理批量和序列维度，只变换最后一维
    2. 对输入数据进行线性变换，主要包含矩阵乘法和加上偏置(映射到新的空间)
    3. 公式：output = input x weight_T + bias (bias为偏置向量)
    4. 应用：self.c_attn(x) 线性层调用，自注意力中的QKV投影，前馈网络(FFN)， 输出层(语言模型头)
  2. F.softmax(Softmax归一化)
    1. 将任意数值转换为概率分布(和为1)，其中dim参数指定沿哪个维度归一化
    2. 将输出层中每个值转化为一个概率分布，其中每个元素的值代表该类别的预测概率
    3. 是常用的激活函数
    4. 性质：输出为概率、最大值选择、平滑化效果
    5. 应用：att = F.softmax(att, dim=-1)
  3. masked_fill(条件填充)
    1. 将满足条件的位置填充为指定值
    2. 特点(函数)：
      1. 原地操作：直接修改原始张量
      2. 保持形状：输出张量形状与输入张量一致
    3. 语法：tensor.masked_fill(mask, value)
      1. mask: 布尔张量（True/False），形状需要可广播到tensor
      2. value: 填充的值
    4. 应用：因果掩码(实现因果注意力只能看前面，不能看后面)，带KV缓存的注意力需动态生成掩码，att = att.masked_fill(causal_mask, float(‘-inf’))

KV Cache

  1. 由前文NumPy+GPT板块中对Transformer架构的介绍中，已知每个token会生成Q,K,V三个向量LLM与Transformer架构
  2. 在做注意力权重计算时，每个token只和自己前面的token做注意力计算，结果不受新的token的影响，则不需要做再次计算，而gernerate方法默认会保存前面收成，生成KV向量，从而加快计算，而新加入向量的token计算只用到前面token的KV,所以直保留这两个即可，而新加入的token也会保留自己的KV，这是KV cache
  3. 所以对话时大模型速度没有减慢，是因为只需计算当前token的KV,之前token的已经缓存了
  4. KV cache利用率只有20%-40%：在大模型推理中，按可生成最长序列长度分配显存，此时会出现三种浪费：
    1. 预分配，但不会用到
    2. 预分配，但尚未用到
    3. 显存之间的间隔碎片，不足以预分配给下一个文本生成

PyTorch+KVcache

由NumPy到PyTorch

根据.py文档的要求完成，测试结果输出为：

1
2
3
4
5
6
Generated text:
Hello world, I'm not sure what to say.

"I'm sorry, but I'm not

Time: 0.50s (39.9 tokens/s)

后来对比了一下readme文档中提到的答案，修改了TODO 7中的激活函数，输出结果仍为上文

x = F.gelu(x) 改为 x = F.gelu(x, approximate=’tanh’)
略微对比了一下，NumPy和PyTorch实现的区别：

  1. NumPy
    1. 纯数学的实现。操作时函数调用
    2. 要手动处理所有运算
    3. 参数通过字典传递
  2. PyTorch
    1. 面向对象的神经网络实现，使用类封装模型结构
    2. 使用了神经网络的相关组件(nn.Module子类)
    3. 可以使用cuda加速

参考文献阅读

KV cache详解

LLM Inference Series: 3. KV caching explained

原文阅读需要一些魔术技巧

Transformer注意力层快速回顾
  1. 单个注意力头的工作机制：
    1. 每个输入向量生成三个低维稠密向量：Q(查询向量),K(键向量),V(值向量)
    2. 对每个Q,生成的O(输出向量) = 所有V的线性组合，该组合的系数->注意力分数
    3. 在计算Q的输出时，不能使用序列中在该标记之后出现的标记对应的V (掩码，把禁止访问的V的注意力分数设为0)
    4. 最后，每个注意力头的output会被拼接起来，通过一次最终的线性变换，以生成最终的output
  2. 注意力计算的二次方增长问题
    1. 已知序列：批次大小b,总长度t
    2. 整体注意力分数的计算量随t呈二次方增长
    3. 涉及模型权重的矩阵乘法的计算需求与总序列长度t呈线性关系
掩码机制导致生成阶段计算冗余

之前的token在不同的迭代过程中相同，特定token的输出在后续迭代中也相同->计算冗余

KV cache

除了最新的token，所有token的K,V都已经在之前的迭代中计算过了，则可以缓存并重用之前迭代中的K,V

  1. 此时为最新token计算输出表示：
    1. 为最新token计算Q,K,V
    2. 从缓存中获取之前token的K,V,并将他们与最新计算的K,V拼接
    3. 使用最新token的Q和所有K计算注意力分数
    4. 使用注意力分数和所有V计算最新token的O
  2. 使用KV cache时，模型实际input是最后生成的token以及KV cache
  3. 预填充阶段不受KV cache策略影响，预填充阶段实际上会计算所有输入标记的K和V，并用它们(预先)填充KV cache
  4. 解码阶段不再使用整个序列做input，而仅用最后生成的token以及KV cache
  5. 在每个生成步骤中，除了权重矩阵之外，还必须从内存中获取一个不断增长的缓存，而仅仅是为了处理每个序列的单个标记
KV缓存实现注意力计算线性增长
  1. 但是客户端闲置过长，GPU内存有限，会通过缓存驱逐策略丢弃过时的对话缓存，当客户端恢复时，要为整个对话的历史重新计算缓存(与对话总长度呈二次方增长的计算成本)
  2. KV cache是一种权衡，用更高的内存使用量和数据传输来换取计算量的减少

KV cache深入探讨

LLM Inference Series: 4. KV caching, a deeper look

这篇文章在讨论KV cache可以增长到多大容量以及其带来的挑战

缓存大小随总序列长度增长
  1. KV cache随批次大小，尤其是总序列长度的增加而线性增长：
    1. GPU内存容量有限，但是KV cache的大小实际上无上限
    2. 总序列长度无法预先确定。KV cache的内存需求也不可预知
    3. 对长序列场景，KV cache的内存消耗可能超过加载模型权重所需的内存容量
    4. KV cache限制了对超长序列(即长上下文窗口)的处理或生成能力，同时也制约了批量处理的规模，从而影响了硬件效率的最大化
  2. 为KV cache争取尽可能多的空间：
    1. 降低模型权重的内存占用(权重量化)
    2. 减少KV cache的内存占用(下文)
    3. 通过跨多GPU分片模型(模型并行)实现多设备内存聚合，但需承担网络通信开销；或利用CPU内存/磁盘等其他存储介质(卸载技术)
  3. 解码步骤涉及大量数据传输,改善延迟的两种途径：
    1. 提升内存带宽(即升级硬件)
    2. 减少数据传输量
  4. 降低内存占用的策略具有三重效益：既能提升硬件利用率和成本效率，又能降低延迟并增加吞吐量
关于输入token也需计费

当输入提示处理完成时(即预填充阶段结束时),消耗了GPU内存(用于存储每个输入令牌的键和值张量)和计算资源(将提示令牌通过模型前向传播)

减少批次大小

有助于降低KV缓存的内存占用从而减少延迟，但会降低硬件利用率，进而影响成本效益

事实上，我们希望尽可能增大批次规模

降低对总序列长度的依赖
  1. 明确决定在每次迭代中重新计算缺失的部分

    用浮点运算来换取GPU内存的节省，目前实践中并未广泛应用此方法

  2. 对于模型几乎不关注或极少关注的令牌，可以不必存储其K V
    1. 滑动窗口注意力(SWA)或局部注意力：确保每个序列在KV cache中存储的张量对永远不会超过窗口大小
    2. 利用注意力层在序列令牌间分配注意力的规律模式：丢弃对输出的贡献始终微乎其微的token，将其对应的注意力分数设为0，从而用更稀疏的注意力矩阵来近似原始矩阵
    3. 相关方法：H2O方法,Scissorhands方法,FastGen方法
    4. 打破对不可预测的总序列长度的依赖，为每个序列分配确定的内存预算，从而极大简化内存管理
减少层数

适用于较小的模型

减少注意力头的数量

模型规模主要由层数和注意力头数决定，减少头数可以选择更小的模型

  1. 只需减少K V头的数量，Q头的数量不影响KV cache大小 (多头查询注意力MQA和分组查询注意力GQA架构的核心思想)
  2. 完全削减说有注意力头对大型模型的影响更显著
注意力头隐藏维度

根据模型系列的不同，注意力头的隐藏维度可能在所有规模变体中保持恒定，因此选择同一系列中的较小变体并无助益

降低每个参数的字节数
  1. 只有同时对权重和“激活值”(即所有权重以外的数据，包括KV cache)进行量化的算法，如LLM.int8()或SmoothQuant，才能生成量化的KV cache
  2. 推理的自回归阶段实际上受内存带宽限制，则单纯提升计算速度的价值有限
  3. 因此，只需要在将张量移入GPU内存前进行量化，并在从GPU内存读取后进行反量化(尽管会引入额外开销)，这就足够了
高效内存管理的重要性
  1. 实际上，简单粗糙的内存管理策略可能导致显著的内存浪费(可能高达80%)
    1. 内存预留浪费
    2. 未充分利用的内存
    3. 重复存储浪费
  2. PagedAttention算法：
    1. 分配固定大小且相对较小的内存块(称为“块”)。每个块可容纳固定数量的token，并可在不同请求间共享
    2. 按需分配和小块大小减轻了内部内存碎片，而相同大小的块则消除了外部内存碎片
    3. 实现了接近零浪费的KV缓存内存利用率(低于4%[21])
  3. 另一个优化：跨请求重用键值缓存 (PagedAttention未涵盖):
    1. 适用于提示共享共同前缀的情况，常见于多轮对话或使用提示模板的场景
    2. 能够跨请求重用KV cache可显著降低延迟(尤其是首个令牌延迟)并提升吞吐量(通过大幅减少具有共享前缀的并发请求的内存占用)
  4. RadixAttention技术:
    1. 实现了此类KV cache重用
    2. 在完成生成请求后并不丢弃KV缓存，而是将其保留在GPU内存中，并向一个专用数据结构(基数树)添加新条目，该结构将令牌序列映射到其KV缓存张量。当新请求到达时，调度器使用基数树进行前缀匹配。如果缓存命中，调度器便复用缓存的KV张量来处理请求
    3. 由于GPU内存有限，缓存的KV张量无法永久保留，采用淘汰策略
卸载到CPU内存甚至disk

可行，但存储介质速度较慢，会有显著的延迟代价，通常面向吞吐量优先的用例(如：离线批处理)

假设在单台主机内运行，受限于所能获取的最大多GPU实例的存储容量

使用多GPU
  1. 通过跨设备分片模型，可以同时利用聚合内存容量和内存带宽来缓解内存压力
    1. 流水线并行：模型和KV缓存会沿层维度分片
    2. 张量并行：KV缓存会沿注意力头维度分片，但MQA此时效率低下：无法将单个头跨设备分片，KV cache必须在所有设备上复制，从而丧失了MQA的优势

      假设在单台主机内运行，受限于所能获取的最大多GPU实例的存储容量

运行结果

根据readme文档的要求及py文件中的指引，完成了KV cache的实现
以下是运行结果

1
2
3
4
5
6
, I'm not sure what to say.

"I'm sorry, but I'm not
, I'm not sure what to say.

"I'm sorry, but I'm not

和没有实现KV cache时的运行结果相同，但是generate的速度明显变快，可见KV cache在即时文本生成中能很好的提高运算速度

nano-vLLM

(能在在2080ti上跑的仓库地址)[https://github.com/duanzhaol/nano-vllm]

依然先过一下前置知识

vLLM

  1. 用于解决运行大型AI模型所带来的速度和内存挑战
  2. 传统LLM模型：
    1. 有时会低效分配GPU内存
    2. 模型或响应的延迟，获得响应的时间也存在问题：与LLM交互的用户越多->响应速度越慢(批处理瓶颈)
    3. scalling:为了获得一个模型并将其提供给大型组织，会超过单个GPU的内存和浮点运算能力，需要复杂的设置和分布式环境
  3. vLLM旨在解决内存碎片到批量执行和分布式推理的问题，在降低延迟的同时提高了性能和GPU资源利用率
  4. vLLM优化：
    1. 分页注意力算法(page attention)：
      1. 帮助更好地管理用于生成下一个令牌的K V:
        1. 使用KV block来管理KV cache
        2. 每个请求需要的KV cache被划分到显存不同的KV block中
        3. 克服了预分配的问题：按需分配，不提前占用显存
        4. 按块分配，减少了内存碎片
      2. 不是将所有内容一次性加载并保持连续的内存空间，而是将内存分成可管理的块
      3. 仅在必要时访问所需的内容
      4. 将请求与所谓的连续批处理捆绑在一起，在序列完成后立即填充GPU插槽
      5. 用更大的batch size来处理请求，提高了系统的吞吐量
    2. 虚拟内存(参考操作系统)
    3. 连续批处理(continous batching)：
      1. 动态添加新请求到正在运行的批次中
      2. iteration level:
        1. 每生成一个新token，称为一轮迭代，每个请求经过多轮迭代才生成最终结果，每轮迭代结束后，调度器和执行器就交互一次，如果检测到有已完成的请求，将其拿出batch后加入新的请求
        2. 优势：
          1. 新到来的请求能被及时处理->解码就绪的请求更多->组合更大的batch->提高解码阶段计算资源利用率
          2. 首字延迟降低->及时发现已经生成结束的请求，及时返回请求处理结果，将新进的请求插入到下一个batch中推理，缩短首字时间TTFT
        3. 局限：解码阶段被预填充阶段干扰而中断->词间延迟(TBT)增加(生成停滞)
      3. request level:
        1. 请求层级，即每个请求的结果生成完毕后，统一的返回结果，调度器和执行器在请求和结束时才进行交互
        2. 优势：调度逻辑简单，无需动态调整批次内容(适合请求到达速率稳定、对解码阶段优先级不敏感的场景)
        3. 局限：请求的token生成阶段可能很长->新到达的请求需要等待很久->高TTFT、高E2E延迟

nano-vLLM

  1. 其他优化：
    1. Prefix Caching：避免重复计算
    2. CUDA Graph：减少启动开销(预先录制计算图，避免每次decode时的kernel启动延迟)
  2. 技术栈：
    1. 基础框架：PyTorch + CUDA
    2. 注意力计算：FlashAttention
    3. 并行策略：Tensor Parallelisam(张量并行)
    4. 内存管理：PagedAttention + Prefix Caching
    5. 计算优化：CUDA Graph
  3. 数据流过程

    用户输入 → Tokenizer → Sequence → Scheduler → ModelRunner → Model → Sampler → 输出

  4. 工作流程：
    1. 初始化：
      1. 用户创建实例 LLM(model_path,…)
      2. LLMEngine被初始化：
        1. 读取模型配置(config.py)
        2. 创建多个ModelRunner进程，每个进程负责一个GPU(根据tensor_parallel_size)
        3. ModelRunner 在自己的进程中加载模型 (qwen3.py) 的一部分权重，并分配KV cache。为了加速，它还会选择性地使用CUDA Graph捕获模型计算图
        4. LLMEngine创建一个Scheduler实例
    2. 接受请求：
      1. 发起推理请求：调用llm.generate(prompts, sampling_params)
      2. 每个prompt和sampling_params被封装成一个Sequence对象
      3. Sequence对象被添加到Scheduler的等待队列(waiting queue)中
    3. 调度与执行(循环进行直到所有请求完成)：
      1. 调度：
        1. 调度器从wating队列中取出Sequence
        2. 处理预填充(Prefill)阶段的序列，调度器检查BlockManager是否有足够空间为序列分配KV cache block:
          1. 空间充足：序列状态waiting->runnint，队列移动：waiting->running
          2. KV cache不足：可能发生抢占(一个正在运行的序列被暂停，其占用的KV cache被释放，然后它被移回到 waiting 队列的头部，等待后续调度)
        3. 预填充完成，调度器处理decode阶段的序列，从running队列中取出序列，为其准备下一个token的生成
      2. 模型执行：
        1. 调度器将准备好的序列(seqs)和阶段信息(is_prefill)传递给 ModelRunner
        2. ModelRunner 根据阶段准备输入数据
        3. 调用模型(Qwen3ForCausalLM)进行前向传播，计算出logits
        4. ModelRunner(rank == 0)调用Sampler,根据logits和sampling_params生成下一个 token_id
      3. 后处理：
        1. Scheduler获取新生成的token_id，将其追加到对应的 Sequence中
        2. 检查序列是否达到结束条件(生成了EOS token或达到了max_tokens)
        3. 若序列完成，状态变为FINISHED,从running队列中移除，占用的KV cache block会被BlockManager释放
    4. 返回结果：当所有请求都完成后(scheduler.is_finished()为True)，llm.generate方法会收集所有完成的序列，使用tokenizer将token_ids解码成文本，并返回给用户

代码

1
2
3
conda deactivate
source ~/uv/nanovllm/bin/activate
# 激活uv环境
  1. 先简单跑一下bench.py测试代码，结果如下(总token数146，总时间3.73s，吞吐量39.17tok/s，解码速度14tok/s),解码速度偏低
1
2
3
4
(nanovllm) dongwanqing@tank-108:~/nano-vllm$ /home/dongwanqing/uv/nanovllm/bin/python /home/dongwanqing/nano-vllm/bench.py
`torch_dtype` is deprecated! Use `dtype` instead!
Generating: 100%|█| 1/1 [00:06<00:00,  6.30s/it, Prefill=2tok/s, Decode=14t
Total: 146tok, Time: 3.73s, Throughput: 39.17tok/s
  1. SJF实现思路(schedule)：
    1. 首先检查是否有等待的序列
    2. 创建一个列表来存储候选序列(长度，索引，序列),遍历列表，检查序列是否能被调度(资源是否够用)
    3. 按prompt长度排序(短优先)
    4. 按排序顺序调度序列，并再次检查资源，然后调度这个序列，最后标记其已被选择
    5. 从waiting队列中移除已调度的序列
  2. 实现SJF后：速度反而慢了，反思后：
    1. SJF调度使KV cache更加碎片化
    2. 与PagedAttention不兼容(破坏了空间局部性，block频繁分配/释放，增加了块表的维护开销)
1
2
3
4
(nanovllm) dongwanqing@tank-108:~/nano-vllm$ /home/dongwanqing/uv/nanovllm/bin/python /home/dongwanqing/nano-vllm/bench.py
`torch_dtype` is deprecated! Use `dtype` instead!
Generating: 100%|█████████████| 1/1 [00:07<00:00,  7.55s/it, Prefill=2tok/s, Decode=11tok/s]
Total: 146tok, Time: 4.82s, Throughput: 30.30tok/s
  1. 在Prefill阶段保持FIFO(原策略)，只在Decode中使用SJF：
    1. 定义新的Decode调度(SJF):
      1. 定期对running队列排序(SJF：按剩余token数):
        1. 不是每次调度都排序，避免开销
        2. 按剩余token数排序(最短剩余时间优先)
      2. 接上原Decode逻辑
    2. 在Sequence类中添加剩余token属性：
      1. remaining_tokens(self)：计算剩余需要生成的token数(用于SJF调度)
      2. is_short_request(self): 判断是否为短请求
    3. 初始化阶段增加SJF相关参数
    4. 为什么在decode实现SJF:
      1. 所有序列都已分配KV cache,避免碎片化，没有额外的内存管理开销
      2. running队列中的序列数通常较少，排序开销小
      3. decode阶段的SJF可能只是调整了生成顺序，保持了连续批处理的优势
  2. 重新运行，发现总时间基本不变，吞吐量有轻微提升
1
2
3
4
(nanovllm) dongwanqing@tank-108:~/nano-vllm$ /home/dongwanqing/uv/nanovllm/bin/python /home/dongwanqing/nano-vllm/bench.py
`torch_dtype` is deprecated! Use `dtype` instead!
Generating: 100%|████████| 1/1 [00:06<00:00,  6.27s/it, Prefill=2tok/s, Decode=14tok/s]
Total: 146tok, Time: 3.71s, Throughput: 39.40tok/s