support chunked prefill and fix minor bug

nanovllm/engine/scheduler.py

@@ -22,26 +22,32 @@ class Scheduler:
         self.waiting.append(seq)
 
     def schedule(self) -> tuple[list[Sequence], bool]:
-        # prefill
         scheduled_seqs = []
-        num_seqs = 0
         num_batched_tokens = 0
-        while self.waiting and num_seqs < self.max_num_seqs:
+
+        # prefill
+        while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
             seq = self.waiting[0]
-            if num_batched_tokens + len(seq) > self.max_num_batched_tokens or not self.block_manager.can_allocate(seq):
+            num_tokens = max(seq.num_tokens - seq.num_cached_tokens, 1)
+            remaining = self.max_num_batched_tokens - num_batched_tokens
+            if remaining == 0 or (not seq.block_table and not self.block_manager.can_allocate(seq)):    # no budget
                 break
-            num_seqs += 1
-            self.block_manager.allocate(seq)
-            num_batched_tokens += len(seq) - seq.num_cached_tokens
-            seq.status = SequenceStatus.RUNNING
-            self.waiting.popleft()
-            self.running.append(seq)
+            if remaining < num_tokens and scheduled_seqs:    # only allow chunked prefill for the first seq
+                break
+            if not seq.block_table:
+                self.block_manager.allocate(seq)
+            seq.num_scheduled_tokens = min(num_tokens, remaining)
+            if seq.num_scheduled_tokens == num_tokens:
+                seq.status = SequenceStatus.RUNNING
+                self.waiting.popleft()
+                self.running.append(seq)
             scheduled_seqs.append(seq)
+            num_batched_tokens += seq.num_scheduled_tokens
         if scheduled_seqs:
             return scheduled_seqs, True
 
         # decode
-        while self.running and num_seqs < self.max_num_seqs:
+        while self.running and len(scheduled_seqs) < self.max_num_seqs:
             seq = self.running.popleft()
             while not self.block_manager.can_append(seq):
                 if self.running:
@@ -50,7 +56,7 @@ class Scheduler:
                     self.preempt(seq)
                     break
             else:
-                num_seqs += 1
+                seq.num_scheduled_tokens = 1
                 self.block_manager.may_append(seq)
                 scheduled_seqs.append(seq)
         assert scheduled_seqs
@@ -62,9 +68,16 @@ class Scheduler:
         self.block_manager.deallocate(seq)
         self.waiting.appendleft(seq)
 
-    def postprocess(self, seqs: list[Sequence], token_ids: list[int]):
+    def postprocess(self, seqs: list[Sequence], token_ids: list[int], is_prefill: bool):
         for seq, token_id in zip(seqs, token_ids):
+            if is_prefill:
+                seq.num_cached_tokens = min(seq.num_cached_tokens + seq.num_scheduled_tokens, seq.num_tokens)
+                if seq.num_cached_tokens < seq.num_tokens or seq.num_completion_tokens > 0:    # chunked prefill or re prefill after preemption
+                    seq.num_scheduled_tokens = 0
+                    continue
             seq.append_token(token_id)
+            seq.num_cached_tokens += 1
+            seq.num_scheduled_tokens = 0
             if (not seq.ignore_eos and token_id == self.eos) or seq.num_completion_tokens == seq.max_tokens:
                 seq.status = SequenceStatus.FINISHED
                 self.block_manager.deallocate(seq)