nona vllm源码解读 (未完成)


题主代码只会一些基本语法和刷了几道算法题的程度,直接嗯看代码实在是困难,于是采用这种笨方法边学边记录每一块的功能

就先局部局部的看着吧,知道每一块干嘛的,之后我再另写一篇看能不能从全局上有什么收获

我觉得我应该听听ai提供的学习路线建议,vllm的源码解读等回家了在wsl上学习好了

nanovllm

sampling_params.py

@dataclass 声明类的核心目的是装数据,而非藏逻辑

其内只有一个类SamplingParams
初始化了temperature、max_tokens和ignore_eox 并规定了temperature的最小值

engine

model_runner

block_manager

Block

init
block_id,ref_count,hash,token_ids

update
更新hash和token_ids

reset

BlockManager

init
初始化block_size,block池,hash和block索引表等

compute_hash
把输入编码转换为字节流
返回哈希流的整数形式

_allocate_block

sequence

Class SequenceStatus

定义Sequence的合法生命周期状态(WAITTING,RUNNING,FINISHED)

Sequence

init
给每个seq_id分配唯一递增id,规定status,初始化变量等

@property 调用函数不需要加括号,像是在访问属性一样,语义上为状态。def了很多函数,把计算的结果伪装成字段一样的属性

block(self, i)
用来取出第i个token块
assert限制边界

append_token
token_ids里添加新token_id
更新last_token和num_tokens

getstate与setstate保存状态

llm_engine

class LLMEngine

init 初始化
1.建立Config对象过滤无关的kwargs
2.初始化进程和存储子进程的列表,spawn 让子进程从干净的解释器状态启动
3.启动tensor_parallel_size - 1个子进程,并在列表中记录,对应rank = 1,2,…,tensor_parallel_size - 1
4.主进程里创建rank=0的ModelRunner实例,持有其他rank的events
5.匹配分词器(tokenizer)
6.tokenizer和config对齐结束符(eos)
7.初始化Scheduler
8.结束、释放内存

exit 退出全部进程
1.通过rank0 runner通知所有runner我们要退出了,自己把循环退掉、资源清掉哈
2.删除主进程里rank0 runner的python引用,允许其相关对象被回收
3.阻塞等待所有子进程退出,之后结束主进程

add_request 调度请求怎么跑
1.已被编码/编码prompt得到token列表
2.token列表+采样参数,创建一个Sequence(一次生成任务的任务单)
3.Sequence交给scheduler执行

step 一步一步跑seqs,对finish的seq输出其最终token id
1.调度器从等待/运行中的sequence里调一批执行,并标记为prefill阶段还是decode
2.seqs交给runner,返回结果token_ids
3.处理结果,判断seq结束后整理输出output
4.统计token量
5.返回output与token量

generate
1.初始化进度条
2.使sampling_params(后简称为sp)为列表
3.将prompt和sp分别对应然后提出推理请求
4.初始化表示输出和吞吐量的变量
5.while调度器检查是否有未完成seq,计时。计算输出和token数量,算出吞吐量
6.输出序列绑定seq id,值绑定token ids,同时进度条增加
7.排序输出,解码输出,关闭进度条,返回输出


文章作者: Austin
版权声明: 本博客所有文章除特別声明外,均采用 CC BY 4.0 许可协议。转载请注明来源 Austin !
评论
  目录