vllm-project · gmarinho2 · May 22, 2025 · May 22, 2025 · May 22, 2025 · May 22, 2025
@@ -276,3 +276,16 @@ def _get_matching_warmup_shapes(
             if prompt_len <= shape['prompt_length']
             and max_tokens <= shape['new_tokens']
         ]
+
+    @classmethod
+    def maybe_update_max_tokens(self, prompt_len: int,
+                                default_max_tokens: int) -> int:
+        if self._warmup_shapes is None:
+            return default_max_tokens
+
+        max_new_tokens = 1
+        for shape in self._warmup_shapes:
+            if prompt_len <= shape['prompt_length']:
+                max_new_tokens = max(max_new_tokens, shape['new_tokens'])
+
+        return max_new_tokens