nano-scheduler

nano-vllm/nanovllm/engine/scheduler.py · bb823b3 · L25–92 · 固定 commit 上游

这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。

25    def schedule(self) -> tuple[list[Sequence], bool]:
26        scheduled_seqs = []
27        num_batched_tokens = 0
28
29        # prefill
30        while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
31            seq = self.waiting[0]
32            remaining = self.max_num_batched_tokens - num_batched_tokens
33            if remaining == 0:
34                break
35            if not seq.block_table:
36                num_cached_blocks = self.block_manager.can_allocate(seq)
37                if num_cached_blocks == -1:
38                    break
39                num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
40            else:
41                num_tokens = seq.num_tokens - seq.num_cached_tokens
42            if remaining < num_tokens and scheduled_seqs:  # only allow chunked prefill for the first seq
43                break
44            if not seq.block_table:
45                self.block_manager.allocate(seq, num_cached_blocks)
46            seq.num_scheduled_tokens = min(num_tokens, remaining)
47            num_batched_tokens += seq.num_scheduled_tokens
48            if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.num_tokens:
49                seq.status = SequenceStatus.RUNNING
50                self.waiting.popleft()
51                self.running.append(seq)
52            scheduled_seqs.append(seq)
53
54        if scheduled_seqs:
55            return scheduled_seqs, True
56
57        # decode
58        while self.running and len(scheduled_seqs) < self.max_num_seqs:
59            seq = self.running.popleft()
60            while not self.block_manager.can_append(seq):
61                if self.running:
62                    self.preempt(self.running.pop())
63                else:
64                    self.preempt(seq)
65                    break
66            else:
67                seq.num_scheduled_tokens = 1
68                seq.is_prefill = False
69                self.block_manager.may_append(seq)
70                scheduled_seqs.append(seq)
71        assert scheduled_seqs
72        self.running.extendleft(reversed(scheduled_seqs))
73        return scheduled_seqs, False
74
75    def preempt(self, seq: Sequence):
76        seq.status = SequenceStatus.WAITING
77        seq.is_prefill = True
78        self.block_manager.deallocate(seq)
79        self.waiting.appendleft(seq)
80
81    def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):
82        for seq, token_id in zip(seqs, token_ids):
83            self.block_manager.hash_blocks(seq)
84            seq.num_cached_tokens += seq.num_scheduled_tokens
85            seq.num_scheduled_tokens = 0
86            if is_prefill and seq.num_cached_tokens < seq.num_tokens:
87                continue
88            seq.append_token(token_id)
89            if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
90                seq.status = SequenceStatus.FINISHED
91                self.block_manager.deallocate(seq)
92                self.running.remove(seq)