今天来读vllm的scheduler.py。两个关键词:调度器,PD分离
调度器
说到调度器,我便会不自觉地往操作系统的调度器去靠,但是vllm的调度器有诸多不同,从队列分类上看,不同于操作系统常见的MLFQ,vllm的调度器只有两个队列,而且不以优先级区分,而是以状态区分:running和waiting。
running first
vllm的scheduler逻辑其实很直观,他会在一开始给出一个整体的算token的预算,先把running里面的请求跑完,还有预算的话再去跑waiting里面的任务。为什么要running first?如果waiting的请求一直进来,将没有能够完成的任务!而且running一直占着KV,只有先把它们跑完才能释放显存。
任务
我感觉,waiting队列里面的任务其实可以看作是都是做prefill的任务,而running是prefill和decode都有,但是实际上scheduler是不会感知到具体是哪一类任务的,而每个任务就是 num_computed_tokens 去追 num_tokens(_with_spec),scheduler 每步只填缺口 num_new = num_tokens − num_computed。prefill阶段num_tokens_with_spec就是prefill词的数量,最后一个prefill词进来后worker会把生成的最后一个词算进num_tokens_with_spec,进入事实上的decode阶段,这样就会使num_tokens_with_spec和num_computed的数量差1,调度器让每次worker生成一个词。
抢占
当running队列任务太多,allocate_slots返回none(显存不足),就会发生抢占,把running队列的任务放在waiting队列的前面。相对地,Waiting里面发生allocate_slots不足就会直接break掉,不再对队列进行处理,这也可以看出调度器对于两个队列的任务侧重点不同,running first不只是一种处理顺序,更是一种策略。
抢占有两种调度策略:FCFS和优先级。
被踢出的任务,V1会进行重算,V0会swap。但是重算重新做prefill往往不是从零开始,因为还可能会命中prefix cache。
被抢占任务的恢复
没有特别的恢复逻辑,因为任务的num_tokens_with_spec一直会更新,而抢占时会把任务的num_computed_tokens设置为0,所以其实就是把prefill重做一遍,只不过带上上次decode生成的token,也相当于是做了一次prefill。
远程KV——从特殊的waiting任务讲起
有这么一种特殊的waiting任务,它的num_computed_tokens大于0!可以看出,这类waiting任务肯定不是被抢占出来的,也不是刚到的任务,原来,他是一种表示等待远端KV-Cache传输过来的特殊waiting任务,被放在skipped_waiting上(因为外部资源依赖而被阻塞的任务,即使预算足够,由于前提条件不足,也没办法resume)。这就像是OS中进程通过系统调用去读块设备那样,要等待块设备把数据读完了才能resume,不然强行拉回来running也没用。
这种特殊的KV有对应的处理逻辑,首先它不能通过一般的resume手段去resume,需要通过特殊的判断把它转为一般的waiting任务,也就是有一个关于是否为WAITING_FOR_REMOTE_KVS的判断,还有一个_try_promote_blocked_waiting_request的方法,来判断远端kv-cache是否已经准备好了,如果是那就把它转成普通的waiting。
此外,在allocate_slots时也需要有特殊的排布,因为它比起一般的waiting会有的prefix命中部分和未命中部分,中间还多了一个KV在远端准备好但是需要KV-Connector传过来,于是需要预先申请好空间的部分。
为什么?——PD分离
可是为什么会有把远程的KV搬过来本地这种情况呢?之前也有过介绍,prefill是compute-bound,decode是memory-bound,各有侧重,与其把这两组工作放在同一张卡上相互拖累,还不如分成两个池子,各干各的事,但是拆开之后就得有从P池到D池的流向(KV),Connector就是干这个的。

参与讨论
(Participate in the discussion)
参与讨论
没有发现评论
暂无评论