fix chunked prefill bugs and refactor

nanovllm/engine/scheduler.py

@@ -11,6 +11,7 @@ class Scheduler:
         self.max_num_seqs = config.max_num_seqs
         self.max_num_batched_tokens = config.max_num_batched_tokens
         self.eos = config.eos
+        self.block_size = config.kvcache_block_size
         self.block_manager = BlockManager(config.num_kvcache_blocks, config.kvcache_block_size)
         self.waiting: deque[Sequence] = deque()
         self.running: deque[Sequence] = deque()
@@ -29,31 +30,26 @@ class Scheduler:
         while self.waiting and len(scheduled_seqs) < self.max_num_seqs:
             seq = self.waiting[0]
             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)):
+            if remaining == 0:
                 break
-
             if not seq.block_table:
-                self.block_manager.allocate(seq)
-
-            # Re-calculate num_tokens after allocate(), as prefix caching may update 
-            # seq.num_cached_tokens during the allocation process.
-            #
-            # Using an outdated num_cached_tokens would overestimate num_scheduled_tokens, 
-            # leading to an inflated 'end' and 'end_block' in prepare_prefill (model_runner.py). 
-            # This results in an 'index out of range' at line 155 when accessing 
-            # seq.block_table[i] beyond its actual physical allocation.
-            num_tokens = max(seq.num_tokens - seq.num_cached_tokens, 1)
-
+                num_cached_blocks = self.block_manager.can_allocate(seq)
+                if num_cached_blocks == -1:
+                    break
+                num_tokens = seq.num_tokens - num_cached_blocks * self.block_size
+            else:
+                num_tokens = seq.num_tokens - seq.num_cached_tokens
             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, num_cached_blocks)
             seq.num_scheduled_tokens = min(num_tokens, remaining)
-            if seq.num_scheduled_tokens == num_tokens:
+            num_batched_tokens += seq.num_scheduled_tokens
+            if seq.num_cached_tokens + seq.num_scheduled_tokens == seq.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
@@ -69,6 +65,7 @@ class Scheduler:
                     break
             else:
                 seq.num_scheduled_tokens = 1
+                seq.is_prefill = False
                 self.block_manager.may_append(seq)
                 scheduled_seqs.append(seq)
         assert scheduled_seqs
@@ -77,19 +74,18 @@ class Scheduler:
 
     def preempt(self, seq: Sequence):
         seq.status = SequenceStatus.WAITING
+        seq.is_prefill = True
         self.block_manager.deallocate(seq)
         self.waiting.appendleft(seq)
 
     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
+            self.block_manager.hash_blocks(seq)
+            seq.num_cached_tokens += seq.num_scheduled_tokens
             seq.num_scheduled_tokens = 0
+            if is_prefill and seq.num_cached_tokens < seq.num_tokens:
+                continue
+            seq.append_token(token_id)
             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)