fix CI bugs

2025-09-21 09:29:47 +00:00 · 2024-01-09 15:18:28 +08:00
parent 2a73e828eb
commit fab294c7f4
5 changed files with 21 additions and 9 deletions
--- a/colossalai/inference/core/engine.py
+++ b/colossalai/inference/core/engine.py
@@ -191,7 +191,14 @@ class InferenceEngine:
                prompt = None
            else:
                prompt = prompts[i]
-            block_table = torch.full([self.inference_config.max_seq_len], -1, device=self.device)
+
+            max_blocks_per_sequence = (
+                self.inference_config.max_input_len
+                + self.inference_config.max_output_len
+                + self.inference_config.block_size
+                - 1
+            ) // self.inference_config.block_size
+            block_table = torch.full([max_blocks_per_sequence], -1, device=self.device)
            sequence = Sequence(
                request_id,
                prompt,