use python int lists for free block instead of torch tensor

deepspeedai · Jan 16, 2024 · bfdb5db · bfdb5db
1 parent 22895da
commit bfdb5db
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mii/batching/ragged_batching.py b/mii/batching/ragged_batching.py
@@ -226,7 +226,7 @@ def _generate_output(self, r: Request) -> bool:
             self.result_queues[r.tid].put_nowait(output)
 
     def _schedule_token_gen(self, requests: List[Request]) -> None:
-        free_blocks = self.inference_engine.free_blocks.min().item()
+        free_blocks = min(self.inference_engine.free_blocks)
         conf_manager = self.inference_engine._config.state_manager
 
         num_schedulable = min(len(requests), conf_manager.max_ragged_sequence_count)
@@ -247,7 +247,7 @@ def _schedule_token_gen(self, requests: List[Request]) -> None:
                     self.scheduled_requests.append(r)
 
     def _schedule_prompts(self, requests: List[Request]) -> None:
-        free_blocks = self.inference_engine.free_blocks.min().item()
+        free_blocks = min(self.inference_engine.free_blocks)
         conf_manager = self.inference_engine._config.state_manager
 
         for r in requests: