fix.

huggingface · Oct 23, 2024 · 82a6cb8 · 82a6cb8
1 parent 5c3efbc
commit 82a6cb8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/server/text_generation_server/models/flash_causal_lm.py b/server/text_generation_server/models/flash_causal_lm.py
@@ -1418,9 +1418,9 @@ def warmup(
                 )
                 max_total_tokens = available_blocks
             else:
-                max_total_tokens = len(batch.input_ids)
+                max_total_tokens = sum(len(input_ids) for input_ids in batch.input_ids)
                 max_input_tokens = (
-                    batch.num_blocks - 1
+                    max_total_tokens - 1
                     if max_input_tokens is None
                     else max_input_tokens
                 )